Adam, Momentum, Schedulers

Two model parameters are each one unit from their best value. The loss gradient asks you to change one 100 times more than the other. How big should a training step be?

The linear algebra lesson showed how a matrix can stretch different directions by different amounts. The same imbalance appears in optimization: a learning rate that barely moves one parameter can send another past its target.

A gradient gives the loss's local slope in each parameter direction. An optimizer uses those slopes to choose an update. Moving against the gradient lowers a smooth loss for a sufficiently small step, but a large step or an update based on past gradients needn't lower it.^{[1]Reference 1Deep Learning.https://www.deeplearningbook.org/}

We'll keep both coordinates visible as we add gradient history, adaptive scaling, weight decay, and a schedule. The examples use NumPy on a CPU; the final training loop uses PyTorch. No model download or GPU is needed.

A 100-to-1 valley

Let a and b be the two parameter errors. Their target is [0, 0].

Use a quadratic loss so every gradient and update is easy to check:

$$L(a, b) = \frac{1}{2}(100a^2 + b^2)$$

The coefficient 100 means the slope changes 100 times faster along a than along b. This is curvature, not a claim that one feature is inherently more important. In matrix form:

$$L(a, b) = \frac{1}{2} \begin{bmatrix} a & b \end{bmatrix} \begin{bmatrix} 100 & 0 \\ 0 & 1 \end{bmatrix} \begin{bmatrix} a \\ b \end{bmatrix}$$

$\mathrm{diag}(100, 1)$ has condition number 100. A single step size is therefore fighting the same uneven stretch that made inverse-like solves fragile in the last chapter.

Differentiate to get

$$\nabla L(a, b) = [100a, b]$$

At [1, 1], the loss is 50.5 and the gradient is [100, 1]. Equal coordinate error hasn't produced equal correction requests: a asks for a step 100 times larger. What should happen if one learning rate serves both?

Run this small calculation and inspect both coordinates, not only the loss.

The loss falls after one step, but the coordinates tell a less comfortable story. a crosses from 1.0 to -0.8, while b only reaches 0.982.

That crossing is the zigzag: the steep error bounces across zero while the shallow error barely walks. A lower rate would calm the wall, but it would also slow the floor.

SGD gives every coordinate one learning rate

Before naming the update stochastic gradient descent (SGD), separate two choices: where the gradient comes from and how the optimizer applies it.

This walkthrough uses a deterministic two-coordinate gradient so the loss geometry stays visible. It stands in for a minibatch gradient without adding sampling noise. Real training adds that noise, so log batch size and gradient statistics when comparing optimizers.

Once the gradient source is fixed, SGD applies a single rule:

$$w_t = w_{t-1} - \eta g_t$$

Here, $w_t$ is the parameter vector after step $t$ and $g_t$ is the minibatch gradient. $\eta$ (eta) is the learning rate, the step-size multiplier.

Try three rates for eight steps on the same valley. Predict which rate keeps a stable and which moves b farther.

The small rate is calm, but b is still 0.961 after eight steps. It also reports the lowest loss of these three runs, so a bigger rate isn't automatically faster learning.

The middle rate drives down a while b still crawls. To see the stability boundary, substitute the gradient into the update: a_next = (1 - 100 * lr) * a. Its magnitude shrinks only when abs(1 - 100 * lr) < 1, or 0 < lr < 0.02. At exactly 0.02, a flips sign without shrinking. At 0.021, the factor is -1.1, so the oscillation grows.

Momentum remembers directions that agree

SGD reacts only to the current gradient. When the steep coordinate flips sign, it forgets that the previous step was already heading back across the wall.

Momentum keeps a signed moving average of recent gradients. Use the exponential-moving-average convention here because it matches Adam's first buffer. This is a normalized teaching form of Polyak's heavy-ball idea, described in his 1964 paper on accelerating iterative methods:

$$m_t = \beta m_{t-1} + (1 - \beta)g_t$$ $$w_t = w_{t-1} - \eta m_t$$

The coefficient $\beta$ controls memory. With $\beta = 0.9$, the new gradient contributes 0.1 and the previous direction contributes 0.9.

PyTorch's torch.optim.SGD uses an unnormalized buffer: with default dampening, it starts at the first gradient and then computes buffer = beta * buffer + gradient. Our m is (1 - beta) times that buffer. To reproduce our updates in that convention, multiply our learning rate by (1 - beta). Check initialization and dampening too before translating another implementation. PyTorch documents these choices.

Feed momentum a prescribed four-gradient history: large alternating values for a, small positive values for b. These are inputs chosen to expose the buffers, not gradients produced by a momentum run. We'll reuse them for RMSProp and Adam. Which signs cancel, and which reinforce one another?

The first buffer stays positive throughout this trace, even when the current gradient is negative. Opposing gradients partially cancel: on step two, 0.9 * 10 + 0.1 * (-80) = 1. The second buffer grows from 0.1 to about 0.333 because every gradient reinforces the same sign. These signed averages are what momentum subtracts from the parameters after multiplying by the learning rate.

Momentum still needs a suitable learning rate for each surface. It changes the path: useful repeated direction builds up, while alternating motion is damped. Now compare measured paths on the same loss instead of assuming that adding a buffer guarantees a win.

At this unchanged numeric learning rate, SGD ends lower than our EMA-style momentum run. That's a valid result, not a failed lesson.

Momentum changes the update scale and usually needs its own learning-rate and memory-coefficient sweep. The benefit to look for is a better-tuned path, not a guaranteed win from adding a buffer.

RMSProp rescales by recent gradient magnitude

Momentum smooths signed direction, but it still sends every coordinate through one global learning rate.

RMSProp keeps a second, unsigned history: an exponential moving average of squared gradients. The method was documented by Tieleman and Hinton in their 2012 RMSProp lecture note.

$$s_t = \rho s_{t-1} + (1 - \rho)g_t^2$$ $$w_t = w_{t-1} - \eta \frac{g_t}{\sqrt{s_t} + \epsilon}$$

The s buffer forgets old squared magnitudes. A large recent gradient makes that coordinate's denominator large, so its update is scaled down.

Unlike momentum's signed buffer, squaring keeps s from canceling when signs alternate. Each RMSProp coordinate keeps the current gradient's sign, but rescaling coordinates differently changes the direction of the whole update vector. This basic form has no Adam-style bias correction.

PyTorch calls rho alpha; its documented RMSprop defaults are lr=0.01, alpha=0.99, and eps=1e-8 (official docs).

PyTorch adds $\epsilon$ after the square root; some other libraries add it inside. Copying $\epsilon$ between them can silently change the update.

Prediction check: This fixture uses rho=0.9 so four steps make the state change easy to inspect. Before running it, predict whether the first update magnitudes will still differ by 100:1.

Use the same deliberately uneven gradient history to test that prediction:

The first row gives both coordinates nearly the same update magnitude, despite the 100:1 raw gradients. With rho=0.9, the first denominator is about sqrt(0.1) * abs(g), so the update magnitude is lr / sqrt(0.1), about 0.03162. Later, the first coordinate's denominator stays large, but its update follows the latest gradient sign. Here and below, update means the vector we subtract from the parameters.

That is the handoff to Adam: add momentum's signed m buffer to this RMSProp-like v buffer, then correct both buffers for their zero start.

Adam stores direction and squared-gradient scale

The Adam optimizer (Adaptive Moment Estimation) keeps the signed average for direction and adds a second moving average for scale: the elementwise square of gradients. Its original formulation is:^{[2]Reference 2Adam: A Method for Stochastic Optimization.https://arxiv.org/abs/1412.6980}

$$m_t = \beta_1 m_{t-1} + (1 - \beta_1)g_t$$ $$v_t = \beta_2 v_{t-1} + (1 - \beta_2)g_t^2$$ $$\hat{m}_t = \frac{m_t}{1 - \beta_1^t}, \qquad \hat{v}_t = \frac{v_t}{1 - \beta_2^t}$$ $$w_t = w_{t-1} - \eta \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon}$$

Read the four lines as two decisions. m chooses a direction from recent signs; v measures how large recent gradients have been. Bias correction turns those startup-shrunken buffers into the hatted estimates used by the update.

The symbols have precise jobs:

Kingma and Ba's usual defaults are $\beta_1 = 0.9$, $\beta_2 = 0.999$, and $\epsilon = 10^{-8}$, with $\epsilon$ added outside the square root in the denominator.^{[2]Reference 2Adam: A Method for Stochastic Optimization.https://arxiv.org/abs/1412.6980}

PyTorch's documented AdamW constructor starts with lr=1e-3, the same betas and eps, and weight_decay=1e-2. They're starting points, not guarantees for a new model or dataset. PyTorch AdamW docs

v is often casually called variance. That wording is misleading here: Adam averages g ** 2; it doesn't subtract the mean gradient to calculate statistical variance.

The first step is a useful sanity check. With bias correction, a nonzero gradient of 100 and one of 1 both produce an update close to one learning rate.

Prediction check: Before running the code, predict what will be equal and what will still differ in the two moment buffers.

Now replay the same prescribed history used for momentum and RMSProp. Look especially at steps two and four, when a's raw gradient is negative:

Unlike RMSProp, Adam still subtracts a positive value from a on steps two and four. Its signed average hasn't turned negative yet. Momentum can keep a parameter moving against the current gradient, so smoothing doesn't guarantee descent on every step. The consistently positive b history keeps its update near 0.01.

Adam isn't finding curvature or proving which feature matters. It transforms recent gradient history into coordinate-wise steps. That distinction matters when you diagnose a run: adaptive scaling isn't a geometry oracle.

Bias correction repairs zero-initialized state

Both moving averages begin at zero, so early values are pulled toward that initialization.

With Adam's usual coefficients, the first raw first-moment buffer is 0.1 * g, while the first raw second-moment buffer is 0.001 * g**2. Dividing those raw values would make the startup update too large. Bias correction removes that startup pull.

Compare the first update with and without correction for one scalar parameter:

Why divide by 1 - beta**t? If every gradient were the same value g, the first two averages would be (1 - beta) * g and (1 - beta**2) * g. After t steps, the weights on observed gradients sum to 1 - beta**t; division restores that sum to one. It removes the zero-start effect, not the uncertainty or lag in a changing gradient history.

AdamW separates data gradients from weight shrinkage

Regularization often nudges parameter magnitudes downward. Define the L2 penalty here as $\frac{\lambda}{2}\lVert w \rVert_2^2$, whose gradient is $\lambda w$.

With ordinary SGD, that penalty can produce the same shrinkage behavior as weight decay after adjusting coefficients. With adaptive optimizers, those operations aren't equivalent: putting the penalty inside Adam's gradient also sends it through coordinate-wise scaling.^{[3]Reference 3Decoupled Weight Decay Regularization.https://arxiv.org/abs/1711.05101}

AdamW applies a decoupled decay directly to the prior parameter value:

$$w_t = w_{t-1} - \eta \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon} - \eta \lambda w_{t-1}$$

$\lambda$ is the weight-decay coefficient. The data-gradient update and shrinkage term are visible separately.

The difference appears even on the first step with a zero data gradient and zero moment buffers. Compare a parameter of 10.0 under the two rules:

The penalty has entered both Adam buffers in the coupled version. In AdamW, it enters neither: the parameter shrinks by the separate factor 1 - lr * weight_decay. Decoupled doesn't mean independent of the schedule, because lr still multiplies the decay. In PyTorch, a zero gradient tensor also differs from grad=None: the optimizer skips parameters with no gradient, including their decay. AdamW API

In PyTorch 2.14, torch.optim.Adam(..., decoupled_weight_decay=True) also selects AdamW behavior. The default is False, so the weight_decay argument alone still selects coupled L2. Inspect the actual constructor rather than inferring the rule from the optimizer's name. Adam API

Schedules change pace across the run

An optimizer decides how to interpret gradients at one step. A learning-rate schedule changes the global multiplier $\eta$ across many steps.

The original Transformer training recipe used Adam with $\beta_1 = 0.9$, $\beta_2 = 0.98$, and $\epsilon = 10^{-9}$, plus a schedule that increased the learning rate linearly for 4000 warmup steps and then decreased it in proportion to the inverse square root of the step number.^{[4]Reference 4Attention Is All You Need.https://arxiv.org/abs/1706.03762} Those betas and $\epsilon$ aren't Kingma and Ba's defaults.

Treat the whole recipe as evidence that one successful run existed, not as a universal setting. Cosine annealing and warm restarts were studied separately by Loshchilov and Hutter; a schedule still needs measurement on your task.

Warmup followed by cosine decay is a shape you can inspect by hand. In this zero-based convention, step counts optimizer updates, not batches or epochs. Update zero uses rate zero, update warmup_steps reaches the peak, and update total_steps - 1 reaches the final rate. Other recipes start at a small nonzero rate, so check the indexing when porting a schedule.

Here the peak is at index 2 and the final rate at index 9. With gradient accumulation, several batches can contribute to one optimizer update; don't advance this schedule on each of those batches. Log the actual rate used alongside the loss so an off-by-one error is visible.

Gradient clipping bounds the input to the optimizer

A malformed or unusually difficult batch can create an unusually large gradient. Global-norm gradient clipping treats every parameter gradient as one long vector and measures its L2 norm.

When that norm exceeds a limit, clipping rescales all gradients together:

$$g_{\text{used}} = \begin{cases} g, & \lVert g \rVert_2 \leq c \\ g \frac{c}{\lVert g \rVert_2}, & \lVert g \rVert_2 > c \end{cases}$$

Here, $c > 0$ is the clipping threshold. Every coordinate gets the same scaling factor, so clipping preserves the gradient's direction before it enters the optimizer. Predict which of the two vectors below will change.

Clipping bounds the optimizer's gradient input, not the full AdamW update by that same threshold. Adam state, learning rate, and weight decay also affect how far the parameters move. For example, at Adam's first step, both [300, 400] and its clipped version [3, 4] yield updates close to [lr, lr] when epsilon is negligible. Clipping still changes the history stored for later steps.

If clipping fires on nearly every batch, investigate data, loss scaling, model stability, or the threshold rather than hiding the signal.

Clip before optimizer.step(). step() uses its input gradient to update AdamW's m and v buffers as well as the parameters, so clipping first keeps one outlier from entering all three paths.

Build it: optimize the valley with AdamW

Use the equations to implement AdamW and run it beside SGD on the same loss surface. Log final loss and both coordinate errors.

The peak rates differ because adaptive normalization changes update scale, so this isn't an optimizer leaderboard. Check that each run stays finite and both coordinate errors end smaller than they started. That's a basic sanity check, not a convergence guarantee.

Why do AdamW's two coordinates end almost equal despite the 100:1 curvature? Starting from equal values, their gradient histories are nearly constant multiples of one another. Adam's numerator and denominator scale together, almost canceling that factor. SGD has no such normalization, so its b error remains much larger. This symmetry belongs to the constructed quadratic; real model selection still needs validation metrics.

Use PyTorch after understanding the state

In a neural-network training loop, PyTorch computes gradients by backpropagation and provides tested optimizer implementations.

Clear old gradients with zero_grad(), run backward(), clip before optimizer.step(), set this step's learning rate, then call step().

Capture the pre-clipping norm so you can see how often the guardrail fires. PyTorch's clip_grad_norm_() returns that total norm. Pass error_if_nonfinite=True so a NaN or Inf raises before it enters optimizer state.

This example uses the same two-coordinate loss, but adds clipping and runs for 60 updates rather than 80. Its final numbers therefore aren't a direct check against the unclipped NumPy run.

This example assigns the current update's learning rate directly before optimizer.step(). If you replace that assignment with a built-in PyTorch scheduler, call scheduler.step() after optimizer.step().

Calling it first skips the schedule's first value. Choose the scheduler's cadence deliberately: some advance per optimizer update, others per epoch, and ReduceLROnPlateau consumes a measured validation metric. PyTorch scheduler guidance

If you later use mixed precision with GradScaler, unscale the gradients before measuring or clipping them with scaler.unscale_(optimizer). Otherwise the clipping threshold applies to artificially scaled values. PyTorch AMP clipping example

Later training chapters will add minibatches, validation data, mixed precision, checkpoints, and distributed state. The optimizer logic you traced here remains inside that larger loop.

AdamW isn't the endpoint of optimizer research. For example, PyTorch 2.14 provides Muon, which treats two-dimensional hidden-layer weight matrices as matrices rather than independent coordinates. Its documentation recommends a method such as AdamW for other parameters, including biases and embeddings. A model can therefore use different optimizers for different parameter groups; compare the whole recipe under the same training budget and validation metric.

Optimizer state is also a memory budget

For each trained parameter, plain SGD has no moving-average buffer, momentum stores one buffer, and AdamW stores two (m and v). If both AdamW buffers use 32-bit floats, their cost is:

$$2 \times 4 \text{ bytes} = 8 \text{ bytes per parameter}$$

A model with 7 billion parameters therefore needs about 56 GB for those two buffers alone. Estimate that number before running the code:

The 56 GB count omits parameters, gradients, activations, and any master-weight copies used by a training setup. At large scale, sharding optimizer state across devices is one reason systems such as ZeRO exist.^{[5]Reference 5ZeRO: Memory Optimizations Toward Training Trillion Parameter Models.https://arxiv.org/abs/1910.02054}

Read optimizer failures from traces

A training curve gives you a symptom, not a cause. Collect a small evidence set before changing an optimizer:

Start with one failure hypothesis and log the quantity that could expose it. Change one mechanism, then re-run on the same data slice.

Worked checks

Use the same valley for each check. Predict first, then compare your work with the solution.

1. Starting from [a, b] = [1, 1], compute one SGD update with lr = 0.01. Does either coordinate cross zero?
2. At Adam step one, take gradient [20, 0.2], lr = 0.001, and ignore epsilon. What is the bias-corrected update?
3. A global gradient has norm 50 and threshold 5. What scaling factor does clipping apply?
4. Your run clips 95 percent of minibatches. Is "clipping works" an adequate conclusion?

Solution checks:

1. The gradient is [100, 1]; the new vector is [0, 0.99]. Neither coordinate overshoots, but the second is still slow.
2. The update is [0.001, 0.001]: first-step normalization removes magnitude differences for nonzero coordinates.
3. The factor is 5 / 50 = 0.1; the direction stays the same.
4. No. Persistent clipping is evidence that your threshold, learning rate, data, or numerical path needs inspection.