# 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`](./3.1._The_Sequential_Model.md) model. You also pick a [loss function](./3.3._Loss_Functions.md) at the same time. After that, `fit` drives the optimizer for you.
RustyML has 5 optimizers: [`SGD`](https://docs.rs/rustyml), [`Adam`](https://docs.rs/rustyml), [`AdamW`](https://docs.rs/rustyml), [`RMSprop`](https://docs.rs/rustyml), and [`AdaGrad`](https://docs.rs/rustyml). 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):
```rust,ignore
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](../Chapter-07/7.3._Performance_Tuning_and_Parallelism.md)).
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
```rust,ignore
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:
```text
param -= lr * grad
```
Set `momentum > 0.0` to accumulate a per-parameter velocity buffer. You can also add the Nesterov look-ahead step:
```text
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.
```rust
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
```rust,ignore
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:
```text
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`](./3.1._The_Sequential_Model.md), `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.
```rust
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
```rust,ignore
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.
```rust
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
```rust,ignore
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:
```text
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
```rust,ignore
AdaGrad::new(learning_rate, epsilon, weight_decay) // all f32, no rho or beta
```
AdaGrad accumulates squared gradients and never forgets them:
```text
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:
| Parameter | Accepted values | Applies to |
| --- | --- | --- |
| `learning_rate` | positive and finite (`> 0`) | all five |
| `momentum` | non-negative and finite (`>= 0`) | SGD |
| `nesterov` | any `bool` (never validated) | SGD |
| `beta1`, `beta2` | in `[0, 1)` and finite | Adam, AdamW |
| `rho` | in `[0, 1)` and finite | RMSprop |
| `epsilon` | positive and finite (`> 0`) | Adam, AdamW, RMSprop, AdaGrad |
| `weight_decay` | non-negative and finite (`>= 0`) | all five |
| `global_clipnorm` | positive and finite (`> 0`) | all five (via `with_global_clipnorm`) |
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.
```rust
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.
```rust
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:
| Optimizer | Per-parameter state | Extra memory vs. parameters |
| --- | --- | --- |
| SGD, `momentum = 0.0` | none | 0x |
| SGD, `momentum > 0.0` | velocity | 1x |
| RMSprop | squared-gradient cache | 1x |
| AdaGrad | squared-gradient accumulator | 1x |
| Adam, AdamW | first moment `m` + second moment `v` | 2x |
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](./3.9._Saving_and_Loading_Weights.md) and [7.2. Model Persistence in Depth](../Chapter-07/7.2._Model_Persistence_in_Depth.md) 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](./3.3._Loss_Functions.md) for more on this. If you change the loss, you will likely need to retune the learning rate, whichever optimizer you chose.