Optimizers

The optimizers in MLX can be used both with mlx.nn but also with pure mlx.core functions. A typical example involves calling Optimizer.update() to update a model’s parameters based on the loss gradients and subsequently calling mlx.core.eval() to evaluate both the model’s parameters and the optimizer state.

# Create a model
model = MLP(num_layers, train_images.shape[-1], hidden_dim, num_classes)
mx.eval(model.parameters())

# Create the gradient function and the optimizer
loss_and_grad_fn = nn.value_and_grad(model, loss_fn)
optimizer = optim.SGD(learning_rate=learning_rate)

for e in range(num_epochs):
    for X, y in batch_iterate(batch_size, train_images, train_labels):
        loss, grads = loss_and_grad_fn(model, X, y)

        # Update the model with the gradients. So far no computation has happened.
        optimizer.update(model, grads)

        # Compute the new parameters but also the optimizer state.
        mx.eval(model.parameters(), optimizer.state)

Saving and Loading

To serialize an optimizer, save its state. To load an optimizer, load and set the saved state. Here’s a simple example:

import mlx.core as mx
from mlx.utils import tree_flatten, tree_unflatten
import mlx.optimizers as optim

optimizer = optim.Adam(learning_rate=1e-2)

# Perform some updates with the optimizer
model = {"w" : mx.zeros((5, 5))}
grads = {"w" : mx.ones((5, 5))}
optimizer.update(model, grads)

# Save the state
state = tree_flatten(optimizer.state)
mx.save_safetensors("optimizer.safetensors", dict(state))

# Later on, for example when loading from a checkpoint,
# recreate the optimizer and load the state
optimizer = optim.Adam(learning_rate=1e-2)

state = tree_unflatten(list(mx.load("optimizer.safetensors").items()))
optimizer.state = state

Note, not every optimizer configuation parameter is saved in the state. For example, for Adam the learning rate is saved but the betas and eps parameters are not. A good rule of thumb is if the parameter can be scheduled then it will be included in the optimizer state.

• Optimizer
  • Optimizer
  • mlx.optimizers.Optimizer.state
    • Optimizer.state
  • mlx.optimizers.Optimizer.apply_gradients
    • Optimizer.apply_gradients()
  • mlx.optimizers.Optimizer.init
    • Optimizer.init()
  • mlx.optimizers.Optimizer.update
    • Optimizer.update()
• Common Optimizers
  • mlx.optimizers.SGD
    • SGD
  • mlx.optimizers.RMSprop
    • RMSprop
  • mlx.optimizers.Adagrad
    • Adagrad
  • mlx.optimizers.Adafactor
    • Adafactor
  • mlx.optimizers.AdaDelta
    • AdaDelta
  • mlx.optimizers.Adam
    • Adam
  • mlx.optimizers.AdamW
    • AdamW
  • mlx.optimizers.Adamax
    • Adamax
  • mlx.optimizers.Lion
    • Lion
• Schedulers
  • mlx.optimizers.cosine_decay
    • cosine_decay()
  • mlx.optimizers.exponential_decay
    • exponential_decay()
  • mlx.optimizers.join_schedules
    • join_schedules()
  • mlx.optimizers.linear_schedule
    • linear_schedule()
  • mlx.optimizers.step_decay
    • step_decay()

clip_grad_norm(grads, max_norm)

Clips the global norm of the gradients.