3.4. Optimizers

An optimizer in RustyML turns gradients into parameter updates. The backward pass of a layer adds a gradient for every trainable tensor to the gradient store of the context. The optimizer then walks those tensors and moves each one downhill. compile on a Sequential model picks the optimizer, and it picks a loss function at the same time. After that, fit drives the optimizer.

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.

2 properties set these optimizers apart from the textbook forms. The constructors take positional arguments, and no argument has a default value. Gradient clipping and weight decay are also part of the optimizer itself, and not part of 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 a model uses all of them. The training loop calls 3 of them. A learning-rate schedule reads and writes the other 2, learning_rate and set_learning_rate (see 3.4.8):

pub trait Optimizer: Send + Sync {
    fn step(&mut self);                               // once per batch
    fn global_clipnorm(&self) -> Option<f32>;         // clip threshold, or None
    fn update(                                        // once per layer
        &mut self,
        scope: usize,                                 // position of the layer, from the input
        layer: &mut dyn LayerBase,                    // the layer whose parameters move
        grads: &Grads,                                // every gradient the backward pass produced
        grad_scale: f32,                              // 1.0, or the clip factor
    );
    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, and gives it the index of that layer as scope. A stateful optimizer advances its notion of time inside step. Adam and AdamW increment a bias-correction timestep there. SGD, RMSprop, and AdaGrad hold no such counter, so step does nothing for them. step rewinds nothing, because no optimizer walks a cursor.

This is why Adam advances its timestep once per batch, and not once per layer. A hand-written training loop must call step once per batch, and not once per layer or once per parameter. A more frequent call breaks the bias-correction math of Adam without any error message.

update receives 3 values from the training loop. The first is scope, the index of the layer counted from the input, so the layer nearest the input has scope 0. The second is grads, the store that holds every gradient of the pass. The third is a grad_scale factor, which the loop computes 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_mut(). That call returns 1 ParamRef per trainable tensor, and a ParamRef holds a name, a mutable value: &mut [f32], and a decays: bool flag. It holds no gradient. The optimizer builds a ParamId::new(scope, name) and looks the gradient up in grads. The optimizer skips a parameter that holds no gradient, without holding back any other tensor of the same layer. A layer therefore yields every trainable tensor it owns on every call, whether the backward pass gave that tensor a gradient or not.

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, whatever the weight_decay value is. Layers set this flag, and optimizers respect it. No caller manages it.

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

ParamId pairs the scope of the layer with the name the layer gives the tensor. The 2 halves together give each parameter a stable identity, so the optimizer state follows the tensor and not its position in a list. The names are kernel, recurrent_kernel, depthwise_kernel, pointwise_kernel, bias, embeddings, alpha, gamma, and beta. The order in which a layer yields its tensors is therefore free. Because the name is the address, no layer may give 2 of its arrays 1 name, and a model build refuses a layer that does. A layer that holds other layers must put its own prefix in front of the name of each array it passes on.

Sequential passes the layer index as scope. A hand-written loop must pass the same scope for the same layer on every step. A Graph model passes the arena position of the layer instead. 1 layer that several nodes share therefore takes 1 update, from the sum of its gradients. A given optimizer instance belongs to 1 model, because scope 0 of one model is not scope 0 of another. 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. The Nesterov look-ahead step is an option on top of that:

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 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 = SequentialBuilder::new()
        .add(Dense::new(1, Activation::Linear).unwrap())
        .build(&Shape::known(x.shape()))
        .unwrap();
    // learning_rate, momentum, nesterov, weight_decay
    model.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 for 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.

A division 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 counter 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, and it is deliberate. Each optimizer keeps the form that its own literature uses, rather than 1 form forced on all 3.

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

The weight_decay of Adam 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 the wanted behavior. 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. The 2 names then stand for 2 different algorithms, and the choice between them is a real choice. A framework that gives both names the decoupled rule makes Adam(weight_decay = x) and AdamW(weight_decay = x) the same optimizer under 2 spellings.

use rustyml::prelude::*;
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 = SequentialBuilder::new()
        .add(Dense::new(8, Activation::ReLU).unwrap())
        .add(Dense::new(1, Activation::Linear).unwrap())
        .build(&Shape::known(x.shape()))
        .unwrap();
    // learning_rate, beta1, beta2, epsilon, weight_decay
    model.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 the right optimizer is not obvious. 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 decay is 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 the coupled scheme of Adam, 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 the gradient history of each weight. 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 for a model that needs regularization.

The practical rule is this: use AdamW for any non-zero weight decay with an adaptive optimizer. Do not use the weight_decay of Adam for that purpose. Use plain Adam, with weight_decay = 0.0, where no regularization applies 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 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 = SequentialBuilder::new()
        .add(Dense::new(16, Activation::ReLU).unwrap())
        .add(Dense::new(1, Activation::Linear).unwrap())
        .build(&Shape::known(x.shape()))
        .unwrap();
    // decoupled weight_decay = 0.01
    model.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 1 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 the answer of RMSprop to the main flaw of AdaGrad. cache is an exponential moving average, and not a running sum. cache therefore forgets old gradients and never grows without bound. The effective step size then stabilizes, instead of decaying to zero.

RMSprop is 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 it 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, as sqrt(cache + epsilon). This placement is not cosmetic. It decides the scale on which epsilon is measured, because 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, lives on the grad scale. The 2 scales correspond roughly as eps_inside = eps_outside^2.

An epsilon of 1e-7 inside the root gives about 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 the defining property of AdaGrad, 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. The decay factor of RMSprop and the second-moment average of Adam 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 the one of RMSprop, 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. That is how a caller turns off those features. 0.0 is invalid for learning_rate and epsilon.

The constructor validates epsilon the same way for all 4 adaptive optimizers, but epsilon 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 1 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, and 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 the threshold, the loop scales every gradient by a single factor, max_norm / global_norm, before the update. This 1 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.

RustyML has the global form alone. A per-variable clip norm, which renormalizes the gradient of each tensor on its own, is a different rule. The 2 rules give different directions as soon as more than 1 tensor is over the limit. A threshold tuned for the per-variable rule is therefore not a threshold for this one.

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 before 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 the moments of Adam, the velocities of SGD, and the cache of RMSprop. The call retunes the same optimizer instance. It does not reset it.

RustyML has no built-in scheduler object. Write the schedule as a plain loop instead. The getter exists so that loop needs no copy of the learning rate of its own. 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, and not rejected.

use rustyml::prelude::*;
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 = SequentialBuilder::new()
        .add(Dense::new(8, Activation::ReLU).unwrap())
        .add(Dense::new(1, Activation::Linear).unwrap())
        .build(&Shape::known(x.shape()))
        .unwrap();
    model.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-parameter state, memory, and persistence

Each optimizer allocates its state lazily. It creates 1 buffer per parameter tensor, sized to that tensor, the first time update reaches it. The optimizer keys the buffers by ParamId, so the state of a parameter follows the pair of the layer position and the tensor name. A tensor that changes length under its own name gets a fresh buffer of the new size, and no other buffer moves.

That key is what makes a model with several layers correct. A cursor over a flat list would tie the state of a parameter to its position in that list.

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 the layer arrays and their validation tags only. It excludes the optimizer and the loss. After load_from_path, the model needs a compile with a fresh optimizer. Its moments, velocities, and timestep all start from zero.

A reload of weights and a resumed training run therefore restart the bias-correction warmup of Adam. This is usually harmless. It matters for a checkpoint taken 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. That is the right property while the architecture is still under exploration.

Use SGD with momentum, momentum = 0.9 and nesterov = true, for fine control and a tuned 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 the coupled weight_decay of Adam 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.

1 subtlety cuts across all of these choices. The effective learning rate is tied to the averaging convention of the loss function. A switch 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. A change of loss usually needs a retuned learning rate, whichever optimizer the model holds.