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)

OptimizerState	The optimizer state implements a recursively defined collections.defaultdict, namely a missing key in an optimizer state is an OptimizerState.
Optimizer()	The base class for all optimizers.
SGD(learning_rate[, momentum, weight_decay, ...])	The stochastic gradient descent optimizer.
RMSprop(learning_rate[, alpha, eps])	The RMSprop optimizer [1].
Adagrad(learning_rate[, eps])	The Adagrad optimizer [1].
Adafactor([learning_rate, eps, ...])	The Adafactor optimizer.
AdaDelta(learning_rate[, rho, eps])	The AdaDelta optimizer with a learning rate [1].
Adam(learning_rate[, betas, eps])	The Adam optimizer [1].
AdamW(learning_rate[, betas, eps, weight_decay])	The AdamW optimizer [1].
Adamax(learning_rate[, betas, eps])	The Adamax optimizer, a variant of Adam based on the infinity norm [1].
Lion(learning_rate[, betas, weight_decay])	The Lion optimizer [1].