3.4. Optimizers

An optimizer in RustyML turns gradients into parameter updates. A layer’s backward pass stashes a gradient for every trainable tensor. The optimizer then walks those tensors and moves each one downhill. You pick an optimizer when you call compile on a Sequential model. You also pick a loss function at the same time. After that, fit drives the optimizer for you.

RustyML has 5 optimizers: SGD, Adam, AdamW, RMSprop, and AdaGrad. All 5 live in rustyml::neural_network::optimizers. The prelude re-exports them. Every constructor returns Result<Self, Error> because it checks its hyperparameters first. Each example below ends its constructor call with .unwrap(), or with real error handling.

The algorithms match the ones in Keras and PyTorch. 2 things differ. The constructors take positional arguments. No argument has a default value. Gradient clipping and weight decay are also part of the optimizer itself, not the training loop.

3.4.1. The optimizer interface and how updates flow

Every optimizer implements the Optimizer trait. The trait has 5 methods, and you use all of them. The training loop calls 3 of the methods for you. A learning-rate schedule reads and writes the other 2, learning_rate and set_learning_rate (see 3.4.8):

pub trait Optimizer {
    fn step(&mut self);                                         // once per batch
    fn update(&mut self, layer: &mut dyn Layer, grad_scale: f32); // once per layer
    fn global_clipnorm(&self) -> Option<f32>;                        // clip threshold, or None
    fn learning_rate(&self) -> f32;                            // current step size
    fn set_learning_rate(&mut self, learning_rate: f32);       // scheduling hook
}

The training loop calls step exactly once per batch, before it touches any layer. It then calls update once per layer. A stateful optimizer advances its notion of time inside step. Adam and AdamW increment a bias-correction timestep there. SGD, RMSprop, and AdaGrad only use step to rewind an internal cursor.

This is why Adam advances its timestep once per batch, not once per layer. A hand-rolled training loop must call step once per batch, not once per layer or once per parameter. Calling it more often breaks Adam’s bias-correction math without any error message.

update receives a grad_scale factor from the training loop. The loop computes grad_scale from global_clipnorm (see below). When clipping is off, grad_scale is 1.0. Inside update, the optimizer asks the layer for its parameters through layer.parameters(). This call returns a flat ParamGrad for each trainable tensor. A ParamGrad holds a mutable value slice, a matching grad slice, and a decays: bool flag.

The decays flag keeps weight decay honest. Weight matrices and convolution or recurrent kernels carry decays = true. Biases and normalization gamma and beta carry decays = false. RustyML never decays a parameter with decays = false, no matter what weight_decay you set. Layers set this flag, and optimizers respect it. You do not manage it yourself.

Layers expose parameters as flat &mut [f32] slices. Because of this, a single per-element kernel handles any tensor shape. Each kernel also switches to a Rayon parallel path once a tensor crosses an element-count threshold (see 7.3).

An optimizer keys its per-parameter state by the position at which a layer yields each tensor. The parameter order must therefore stay stable across steps, and it does. A given optimizer instance belongs to one model. Do not share an optimizer instance across models.

3.4.2. SGD

SGD::new(learning_rate, momentum, nesterov, weight_decay) // all f32 except nesterov: bool

SGD is plain stochastic gradient descent. It supports optional momentum, Nesterov acceleration, and decoupled weight decay. With momentum = 0.0, the update is the textbook form:

param -= lr * grad

Set momentum > 0.0 to accumulate a per-parameter velocity buffer. You can also add the Nesterov look-ahead step:

v    = momentum * v + grad
step = grad + momentum * v   (nesterov = true)   or   v   (nesterov = false)
param -= lr * step

0.9 is the conventional momentum value. nesterov only matters when momentum is non-zero. The weight_decay here is decoupled, the SGDW formulation. Before the gradient step, SGD shrinks weight tensors by param *= (1 - lr * weight_decay), independent of the gradient. AdamW uses this same decoupling, applied here to SGD.

use rustyml::prelude::*;
use rustyml::neural_network::sequential::Sequential;
use ndarray::Array;

fn main() {
    // Tiny inline regression: y = 2*x
    let x = Array::from_shape_vec((4, 1), vec![0.5_f32, 1.0, 1.5, 2.0])
        .unwrap()
        .into_dyn();
    let y = Array::from_shape_vec((4, 1), vec![1.0_f32, 2.0, 3.0, 4.0])
        .unwrap()
        .into_dyn();

    let mut model = Sequential::new();
    model
        .add(Dense::new(1, 1, Activation::Linear).unwrap())
        // learning_rate, momentum, nesterov, weight_decay
        .compile(SGD::new(0.05, 0.9, true, 0.0).unwrap(), MeanSquaredError::new());

    let before = model.predict(&x).unwrap();
    model.fit(&x, &y, 10).unwrap();
    let after = model.predict(&x).unwrap();
    println!("shapes: {:?} -> {:?}", before.shape(), after.shape());
}

Use SGD with momentum when you want fine control and predictable behavior. Its dynamics are well understood. Its single moving average is cheap to keep. SGD with momentum still sets the generalization standard that other methods get measured against.

The cost is sensitivity to the learning rate. A rate too large makes training diverge. A rate too small makes training slow. Adam removes much of this sensitivity.

3.4.3. Adam

Adam::new(learning_rate, beta1, beta2, epsilon, weight_decay) // all f32

Adam keeps 2 per-parameter moving averages. The first moment, m, tracks the mean gradient. The second moment, v, tracks the mean squared gradient. Adam uses both to give every parameter its own adaptive step size. The full update, run per element, is:

t += 1                                  (advanced once per batch, in step())
m = beta1*m + (1 - beta1)*grad
v = beta2*v + (1 - beta2)*grad^2
m_hat = m / (1 - beta1^t)               <- bias correction
v_hat = v / (1 - beta2^t)               <- bias correction
param -= lr * m_hat / (sqrt(v_hat) + epsilon)

Bias correction matters because m and v start at zero. On the first few steps, this pulls both moments toward zero. The pull is strong for v, because beta2 = 0.999 keeps v close to zero at t = 1.

Dividing by (1 - beta^t) rescales the moments. At t = 1, the denominator 1 - beta1 cancels the (1 - beta1) factor already in m. This recovers the raw gradient. As t grows, beta^t approaches 0, and the correction fades to a no-op.

This is why the timestep advances inside step, once per batch, and not inside update. With full-batch fit, t equals the epoch count. With fit_with_batches, t advances once per mini-batch. The timestep saturates instead of overflowing. Extremely long training runs therefore stay well defined.

Adam adds epsilon outside the square root. RMSprop and AdaGrad (3.4.5 and 3.4.6) add their epsilon inside it. This looks like an inconsistency, but it is deliberate. Keras splits epsilon placement the same way. RustyML matches Keras optimizer by optimizer, instead of forcing one form on all 3.

As a result, epsilon is not on the same scale in every optimizer. Read the scale note under RMSprop before you carry an epsilon value from one optimizer to another.

Adam’s weight_decay implements classic coupled L2 regularization. Adam folds weight_decay * param into the gradient before the moment update. The penalty therefore flows through m and v, and the adaptive 1 / (sqrt(v_hat) + epsilon) denominator rescales it. This coupling is usually not what you want. See AdamW, next, for the alternative.

With weight_decay = 0.0, the coupling has no effect, and Adam and AdamW become byte-for-byte identical. RustyML keeps the coupling in Adam on purpose. This is a deliberate divergence from Keras 3.

Keras 3’s base optimizer applies decoupled decay to every optimizer. This makes Adam(weight_decay=x) and AdamW(weight_decay=x) numerically the same in Keras. RustyML instead follows PyTorch. In PyTorch, torch.optim.Adam(weight_decay=) is coupled, and AdamW’s is not. This keeps the 2 names genuinely distinct in RustyML.

use rustyml::prelude::*;
use rustyml::neural_network::sequential::Sequential;
use ndarray::Array;

fn main() {
    // Tiny inline regression: y = x0 + 2*x1
    let x = Array::from_shape_vec((4, 2), vec![0.0_f32, 1.0, 1.0, 0.0, 1.0, 1.0, 2.0, 1.0])
        .unwrap()
        .into_dyn();
    let y = Array::from_shape_vec((4, 1), vec![2.0_f32, 1.0, 3.0, 4.0])
        .unwrap()
        .into_dyn();

    let mut model = Sequential::new();
    model
        .add(Dense::new(2, 8, Activation::ReLU).unwrap())
        .add(Dense::new(8, 1, Activation::Linear).unwrap())
        // learning_rate, beta1, beta2, epsilon, weight_decay
        .compile(Adam::new(0.01, 0.9, 0.999, 1e-8, 0.0).unwrap(), MeanSquaredError::new());

    model.fit(&x, &y, 50).unwrap();

    let preds = model.predict(&x).unwrap();
    println!("prediction shape: {:?}", preds.shape());
}

Adam::new(0.001, 0.9, 0.999, 1e-8, 0.0) is the default choice. Use this configuration when you are not sure which optimizer to pick. It tolerates a wide range of learning rates. It converges quickly on messy loss surfaces. It needs almost no tuning to start training.

3.4.4. AdamW

AdamW::new(learning_rate, beta1, beta2, epsilon, weight_decay) // identical signature to Adam

AdamW runs the same moment math and bias correction as Adam. The only difference is where weight decay enters the update. This difference matters as soon as you turn decay on. AdamW is decoupled, the Loshchilov and Hutter formulation.

AdamW shrinks the weight directly, param *= (1 - lr * weight_decay), before an ordinary Adam step. The decay never touches m or v. The adaptive denominator never divides it.

This distinction is not cosmetic. In Adam’s coupled scheme, the weight_decay * param term rides through the second moment v. A parameter that has seen large gradients gets a large v and a large denominator. It therefore gets less effective decay than a quiet parameter. Regularization strength ends up tangled with each weight’s gradient history. This is the opposite of the uniform shrink that weight decay is supposed to give.

AdamW severs this link. Every weight decays by the same (1 - lr * weight_decay) factor, regardless of its gradients. For this reason, AdamW generalizes better. It is the standard choice when you care about regularizing a model.

The practical rule is this: use AdamW for any non-zero weight decay with an adaptive optimizer, not Adam’s weight_decay. Use plain Adam, with weight_decay = 0.0, when you apply no regularization at all. At weight_decay = 0.0, the 2 optimizers run the same algorithm. There is no reason to prefer Adam over AdamW, except habit.

use rustyml::prelude::*;
use rustyml::neural_network::sequential::Sequential;
use ndarray::Array;

fn main() {
    let x = Array::from_shape_vec(
        (4, 3),
        vec![0.0_f32, 1.0, 2.0, 1.0, 0.0, 1.0, 2.0, 1.0, 0.0, 1.0, 1.0, 1.0],
    )
    .unwrap()
    .into_dyn();
    let y = Array::from_shape_vec((4, 1), vec![1.0_f32, 0.5, 0.5, 0.75])
        .unwrap()
        .into_dyn();

    let mut model = Sequential::new();
    model
        .add(Dense::new(3, 16, Activation::ReLU).unwrap())
        .add(Dense::new(16, 1, Activation::Linear).unwrap())
        // decoupled weight_decay = 0.01
        .compile(
            AdamW::new(0.01, 0.9, 0.999, 1e-8, 0.01).unwrap(),
            MeanSquaredError::new(),
        );

    model.fit(&x, &y, 30).unwrap();
    println!("trained: {:?}", model.predict(&x).unwrap().shape());
}

3.4.5. RMSprop

RMSprop::new(learning_rate, rho, epsilon, weight_decay) // all f32

The type name is RMSprop, with a lowercase p. RMSprop keeps one moving average of squared gradients per parameter. It normalizes each step by the square root of that average:

cache = rho*cache + (1 - rho)*grad^2
param -= lr * grad / sqrt(cache + epsilon)

rho is the squared-gradient decay rate, conventionally 0.9. It is RMSprop’s answer to AdaGrad’s main flaw. cache is an exponential moving average, not a running sum. Because of this, cache forgets old gradients and never grows without bound. The effective step size therefore stabilizes, instead of decaying to zero.

You can think of RMSprop as Adam without the first moment. It gives adaptive per-parameter scaling, but no momentum and no bias correction. Its weight_decay is decoupled, in the AdamW style, and applies to weight tensors before the adaptive step. RMSprop suits recurrent networks and non-stationary objectives. On most other problems, Adam covers the same ground.

epsilon goes inside the root here, matching Keras’ sqrt(velocity + epsilon). This placement is not cosmetic. It decides the scale on which epsilon is measured. cache accumulates squared gradients.

An epsilon added before the root therefore lives on the grad^2 scale. An epsilon added after the root, as in Adam above and in PyTorch’s RMSprop, lives on the grad scale. The 2 scales correspond roughly as eps_inside = eps_outside^2.

Keras’ default of 1e-7 inside the root gives the same guard as 3.2e-4 outside it. Do not reuse the same epsilon value in both forms. It is off by orders of magnitude in the other form. AdaGrad, below, also uses the inside form.

3.4.6. AdaGrad

AdaGrad::new(learning_rate, epsilon, weight_decay) // all f32, no rho or beta

AdaGrad accumulates squared gradients and never forgets them:

accumulator += grad^2
param -= lr * grad / sqrt(accumulator + epsilon)

The accumulator only grows. sqrt(accumulator) therefore only ever increases. The effective learning rate, lr / sqrt(accumulator + epsilon), decays monotonically toward zero. This is AdaGrad’s defining property, and it cuts both ways.

For convex problems and sparse features, this decay is a benefit. A rarely-seen parameter keeps a large step. A frequently-updated parameter anneals automatically. The decay to zero also gives clean convergence guarantees.

For a deep network trained over many steps, this decay is a drawback. The step size shrinks until learning effectively stalls. RMSprop’s decay factor and Adam’s second-moment average exist to fix exactly this problem.

learning_rate here is genuinely an initial rate, typically 0.01. It sets a ceiling, and the accumulator only pulls the effective rate down from there. Its epsilon, like RMSprop’s, sits inside the root, so it is on the squared-gradient scale. Weight decay, as with the other optimizers, is decoupled and applies to weights only.

3.4.7. Parameter validation

Every constructor validates its hyperparameters before it returns. A bad hyperparameter therefore fails at construction time, with Error::InvalidParameter, instead of producing NaN values well into training. The table below lists the rules, taken directly from the validators:

The decay-rate ranges are half-open on purpose. beta1 = 0.0 and rho = 0.0 are valid, an inclusive lower bound. 1.0 is invalid, an exclusive upper bound. A decay of 1.0 would freeze the moving average and never take in a new gradient.

SGD has no epsilon, because it never divides by an adaptive denominator. 0.0 is valid for momentum and weight_decay. This is how you turn off those features. 0.0 is invalid for learning_rate and epsilon.

epsilon is validated the same way for all 4 adaptive optimizers, but it does not mean the same thing in each. RMSprop and AdaGrad add epsilon inside the square root, on the squared-gradient scale. Adam and AdamW add it outside, on the gradient scale. A value that is sane for one pair is roughly the square, or the square root, of a sane value for the other pair.

use rustyml::neural_network::optimizers::{AdaGrad, Adam, AdamW, RMSprop, SGD};
use rustyml::error::Error;

fn main() {
    // learning_rate must be positive and finite
    assert!(matches!(
        SGD::new(0.0, 0.0, false, 0.0),
        Err(Error::InvalidParameter { .. })
    ));
    // beta1 must be in [0, 1): 1.0 is out of range
    assert!(matches!(
        Adam::new(0.001, 1.0, 0.999, 1e-8, 0.0),
        Err(Error::InvalidParameter { .. })
    ));
    // epsilon must be positive and finite
    assert!(matches!(
        RMSprop::new(0.01, 0.9, 0.0, 0.0),
        Err(Error::InvalidParameter { .. })
    ));
    // weight_decay must be non-negative and finite
    assert!(matches!(
        AdaGrad::new(0.01, 1e-8, -0.1),
        Err(Error::InvalidParameter { .. })
    ));
    // global_clipnorm must be positive and finite
    assert!(matches!(
        AdamW::new(0.001, 0.9, 0.999, 1e-8, 0.0)
            .unwrap()
            .with_global_clipnorm(0.0),
        Err(Error::InvalidParameter { .. })
    ));

    // A fully valid configuration, clipping enabled
    let opt = Adam::new(0.001, 0.9, 0.999, 1e-8, 0.0).unwrap();
    let _clipped = opt.with_global_clipnorm(1.0).unwrap();
    println!("validation checks passed");
}

3.4.8. Gradient clipping and learning-rate scheduling

Gradient clipping is off by default. Enable it through a consuming builder, with_global_clipnorm, present on all 5 optimizers. This is clip-by-global-norm, not per-element clamping. The training loop sums the squared gradients across every tensor in the model and takes the global L2 norm.

If the global norm exceeds your threshold, the loop scales every gradient by a single factor, max_norm / global_norm, before the update. This one uniform factor preserves the descent direction exactly. Per-element clamping would bend that direction instead. A non-finite global norm is deliberately left unscaled, so a genuine divergence still shows up as NaN, instead of being masked. This clipping is the recommended way to tame large but finite gradients, for example in RNNs or deep stacks. The backward pass itself applies no clamping.

The name matches Keras’ global_clipnorm exactly. If you port a Keras model that sets clipnorm instead, note the difference. Keras’ clipnorm renormalizes each variable’s gradient independently. Once more than one tensor is over the limit, this points in a different direction entirely. RustyML has only the global form. A clipnorm threshold from Keras therefore cannot carry over unchanged.

Learning-rate scheduling uses a read and write pair, learning_rate and set_learning_rate. The model exposes them as Sequential::learning_rate, which returns Option<f32> and is None until you compile, and as Sequential::set_learning_rate. Call the setter between epochs or batches to retune the step size. This preserves all accumulated state, such as Adam’s moments, SGD’s velocities, and RMSprop’s cache. You retune the same optimizer instance. You do not reset it.

RustyML has no built-in scheduler object. You write the schedule as a plain loop instead. The getter exists so that loop does not need its own copy of the learning rate. A separate copy would go stale the moment anything else retunes the optimizer.

The getter returns whatever was last written. Unlike the constructors, set_learning_rate validates nothing. A zero or negative rate is therefore stored and read back unchanged, not rejected.

use rustyml::prelude::*;
use rustyml::neural_network::sequential::Sequential;
use ndarray::Array;

fn main() {
    let x = Array::from_shape_vec((4, 2), vec![0.0_f32, 1.0, 1.0, 0.0, 1.0, 1.0, 2.0, 1.0])
        .unwrap()
        .into_dyn();
    let y = Array::from_shape_vec((4, 1), vec![2.0_f32, 1.0, 3.0, 4.0])
        .unwrap()
        .into_dyn();

    let mut model = Sequential::new();
    model
        .add(Dense::new(2, 8, Activation::ReLU).unwrap())
        .add(Dense::new(8, 1, Activation::Linear).unwrap())
        .compile(
            Adam::new(0.01, 0.9, 0.999, 1e-8, 0.0)
                .unwrap()
                .with_global_clipnorm(1.0) // clip the global gradient norm to 1.0
                .unwrap(),
            MeanSquaredError::new(),
        );

    // Manual step decay: halve the LR each block. The optimizer keeps all state
    for _ in 0..3 {
        model.fit(&x, &y, 10).unwrap();
        // The optimizer holds the rate, so read it back instead of shadowing it
        let lr = model.learning_rate().unwrap();
        model.set_learning_rate(lr * 0.5);
    }

    assert_eq!(model.learning_rate(), Some(0.01_f32 / 8.0));
}

3.4.9. Per-layer state, memory, and persistence

Each optimizer allocates its state lazily. It creates one buffer per parameter tensor, sized to that tensor, the first time update reaches it. The buffers are indexed by the order in which layers yield parameters. If a tensor’s length changes at a given position, RustyML resets that buffer to match.

This is why the 2-layer convergence tests check correct buffer allocation across layers. The optimizer’s state must line up with the parameters it shadows.

The memory cost is the number and size of those buffers:

For a model with P weights in f32, Adam and AdamW carry roughly 2 * 4 * P bytes of optimizer state. This is on top of the parameters and their gradients themselves. This is the price of adaptivity. It is also why a model that trains fine under SGD can run out of memory under Adam. Plain SGD carries no extra state. On a tight memory budget, this can be the deciding factor.

This state lives entirely in the optimizer instance. RustyML does not persist it. Sequential::save_to_path serializes layer architecture and weights only. It explicitly excludes the optimizer and the loss. After load_from_path, you must compile a fresh optimizer. Its moments, velocities, and timestep all start from zero.

Reloading weights and resuming training therefore restarts Adam’s bias-correction warmup. This is usually harmless. It matters if you checkpoint mid run. See 3.9. Saving and Loading Weights and 7.2. Model Persistence in Depth for the full persistence story.

3.4.10. Choosing an optimizer

Start with Adam, using 0.001, 0.9, 0.999, 1e-8, 0.0. Adam is the least fussy optimizer about learning rate. It gets almost any model training. This is what you want while you are still exploring the architecture.

Use SGD with momentum, momentum = 0.9 and nesterov = true, when you want fine control and are willing to tune the learning rate. Its generalization is the reference standard. Its state is cheap. Its behavior is easy to predict.

Use AdamW with a non-zero weight_decay whenever regularization is the goal. Its decoupled decay is the correct way to regularize an adaptive optimizer. Prefer it over Adam’s coupled weight_decay in every case where decay is on.

RMSprop is a good fallback for recurrent networks and non-stationary problems. AdaGrad performs well on convex, sparse-feature problems, where its decaying step size is an asset rather than a liability. Avoid AdaGrad for long deep-network runs, where that same decay stalls learning.

One subtlety cuts across all of these choices. The effective learning rate is tied to your loss function’s averaging convention. Switching between a per-element loss and a per-prediction-site loss changes the gradient magnitude. Per-element losses include MSE, MAE, and binary cross-entropy. A per-prediction-site loss is categorical cross-entropy. This change in magnitude also changes the step size.

See the averaging note in 3.3. Loss Functions for more on this. If you change the loss, you will likely need to retune the learning rate, whichever optimizer you chose.