Adam, Momentum, Schedulers

A triage model can learn the wrong lesson even while its loss goes down. Suppose it uses one weight for the common signal "escalated incident" and another for the rarer contrast "security issue rather than routine failure." A training step can correct the common weight too aggressively, jump past the best setting, and barely improve the contrast weight.

A matrix can have strong and weak directions. Training has a related geometry: a loss can bend sharply in one parameter direction and gently in another. A gradient says which way is uphill; an optimizer decides how much to move downhill, what history to remember, and how its pace changes over time.^{[1]Reference 1Deep Learning.https://www.deeplearningbook.org/}

One loss surface has two speeds

Measure everything on the smallest surface that exposes the problem. Let a be error in the common-signal weight and b be error in the contrast weight. The best possible setting is [a, b] = [0, 0]. Near that setting, use this loss:

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

The multiplier 100 says that the loss curves much more sharply along a than along b. It doesn't say that common tickets matter 100 times more to customers. It describes local training geometry.

Differentiate the loss to get its gradient:

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

At the initial parameter error [1, 1], the loss is 50.5 and the gradient is [100, 1]. Both parameters are equally far from their target, but the first coordinate asks for a step 100 times larger.

Run this program first to expose the imbalance before changing optimizers.

The loss falls after one step, but look at the coordinates: a crosses from 1.0 to -0.8, while b only reaches 0.982. That crossing is the zigzag. The optimizer is reducing the steep error by bouncing across zero while the shallow error moves slowly.

SGD gives every coordinate one learning rate

Stochastic gradient descent (SGD) updates parameters from a gradient computed on a minibatch of examples:

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

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

Our tiny loss is deterministic rather than minibatched, so it removes batch noise and leaves the geometry exposed. Try three learning rates for eight steps:

The small rate is calm but slow on b. Over only eight steps it also reports the lowest loss of these three runs, which is a reminder to measure rather than equate a larger rate with faster learning. The middle rate drives down a while it still crawls on b. The rate just above 0.02 makes the steep coordinate grow in magnitude, because its multiplicative update factor is 1 - 100 * 0.021 = -1.1.

Why can't increasing SGD's learning rate solve the slow progress on b here?

Answer

The same rate also multiplies the much steeper gradient on a. Once the rate exceeds 0.02, the a update flips sign and grows instead of shrinking. A rate safe for a leaves b comparatively slow.

Momentum remembers directions that agree

SGD reacts to the current gradient only. Momentum stores a moving average of gradients. Use the exponential-moving-average convention because it matches Adam's first buffer:

$$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 previous direction contributes 0.9. Some libraries use an unnormalized velocity convention instead; that changes the numerical learning-rate scale, not the idea.

Across the steep wall, gradients alternate sign. Along the shallow valley floor, they tend to agree. Watch the moving average cancel one alternating component while accumulating the consistent component:

The first memory coordinate keeps changing direction and stays much smaller than the raw wall gradients. The second remains positive because each step agrees that b should shrink.

Momentum doesn't grant one perfect learning rate for every surface. It changes the path: useful repeated direction builds up, while alternating motion is damped. The next experiment uses the same loss and compares measured results rather than promising a winner in every task.

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.

Nesterov momentum changes when the gradient is measured: it asks for the gradient near the point the current velocity is about to reach, rather than only at the current point. It can anticipate a turn. Focus on understanding stored direction; later training labs can compare library momentum variants on real minibatches.

Adam stores direction and squared-gradient scale

Momentum still uses one global learning rate after smoothing. Adam (Adaptive Moment Estimation) adds a second moving average: 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}$$

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}

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.

One more subtle point: on Adam's very first step, a nonzero gradient of 100 and a nonzero gradient of 1 both produce an update close to one learning rate after bias correction. One observation isn't enough to infer which coordinate repeatedly needs caution.

Now provide a short gradient history. The first coordinate alternates with large values, while the second keeps agreeing on a small positive direction. Adam's step reflects both pieces of state.

The alternating large coordinate receives a small signed move once its directional evidence conflicts. The consistently positive coordinate keeps a positive update. Adam isn't finding curvature or proving which feature matters; it transforms recent gradient history into coordinate-wise steps.

Bias correction repairs zero-initialized state

Both moving averages begin at zero. Early in training, they 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.

This program compares the first update with and without correction for one scalar parameter:

Bias correction isn't a cosmetic detail. It makes the startup behavior correspond to estimates that aren't artificially pulled toward the all-zero initial buffers.

Why is v_t not the statistical variance of a coordinate's gradient?

Answer

Adam updates v_t from squared gradients, g_t ** 2, without subtracting a mean. It is an uncentered second-moment estimate: a running measure of squared magnitude.

AdamW separates data gradients from weight shrinkage

Regularization often nudges parameter magnitudes downward. With ordinary SGD, adding an L2 penalty to the loss 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 when the data gradient is zero. A parameter of 10.0 should decay a small amount. If its L2 penalty is instead fed through Adam on the first step, normalization can produce a very different move:

This isn't a recommendation to choose a decay coefficient from a single number. It proves a narrower point: AdamW decay and Adam with an L2 term implement different update rules.

Schedules change pace across the run

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

The original Transformer training recipe used Adam with a schedule that increased its 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} That's evidence for one successful recipe, not a universal setting.

Warmup followed by cosine decay is another useful shape to understand. It reaches a peak slowly, then follows a smooth curve toward a configured final rate. This self-contained example makes the schedule inspectable:

A schedule choice needs an experiment: log the learning rate beside training loss, validation loss, and gradient norm. A bad curve doesn't identify its own cause.

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, measures its L2 norm, and rescales all gradients together when the norm exceeds a limit:

$$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$ is the clipping threshold. Because every coordinate gets the same scaling factor, clipping preserves the gradient's direction before it enters the optimizer.

Clipping limits the gradient passed to the optimizer. It doesn't guarantee a bound on an AdamW parameter update, because Adam state, learning rate, and weight decay also affect that update. If clipping fires on nearly every batch, investigate data, loss scaling, model stability, or the threshold rather than hiding the signal.

Build it: optimize the valley with AdamW

Now implement AdamW from the equations and test it against SGD on the same loss surface. The test logs final loss and both coordinate errors. This test doesn't declare an optimizer universally superior; it proves that your implementation handles this measured failure case.

Your acceptance test is explicit: both coordinates should shrink and the reported loss should be finite. A comparison on this constructed surface supports reasoning about this surface only. 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. This short example optimizes the same two-coordinate loss, clips the gradient before optimizer.step(), and applies a decaying learning rate. PyTorch's clip_grad_norm_() returns the total norm before clipping. Capture it when you need evidence about how often the guardrail fires, and pass error_if_nonfinite=True so a NaN or Inf norm raises before it enters optimizer state.

The sequence matters:

1. Clear old gradients.
2. Compute a loss and run backpropagation.
3. Inspect or clip gradients before the optimizer consumes them; reject non-finite norms.
4. Set the scheduled learning rate for this update.
5. Apply the optimizer step.

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. PyTorch documents that ordering explicitly.

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.

In the PyTorch loop, why must gradient clipping happen after loss.backward() but before optimizer.step()?

Answer

Backpropagation must create gradients before they can be measured or clipped, and clipping must modify them before the optimizer consumes them. Clipping after the step cannot protect that update.

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 are stored as 32-bit floats, they require:

$$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:

That 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

No training curve proves a cause by itself. Collect a small evidence set before changing an optimizer:

This is how optimizer work becomes research practice: state a failure hypothesis, log the quantity that would expose it, change one mechanism, and re-run.

Practice with answers

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.

Optimizer release checklist

Before accepting a training change, log loss, learning rate, gradient norm, update norm, validation metric, and clipping frequency on the same data slice. Reject non-finite values, verify schedule boundaries, and include optimizer-state bytes in the device-memory budget. Change one mechanism per comparison so traces can support a cause.