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)

• 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.