Lesson 11. The AdamW Optimizer Step

Mission link: This is stage 4's capstone: lesson 10's backward pass produces a gradient for every weight, and the optimizer is what actually turns those gradients into updated weights, closing the loop that lets the model train from scratch.
Primary source: Paper: "Decoupled Weight Decay Regularization", Loshchilov and Hutter, 2019
Prerequisites: Lesson 10, Scaled dot-product attention

Warm-up

1. ▢ What does calling loss.backward() populate, and why must gradients typically be zeroed before the next call?

2. ▢ Why does a residual connection guarantee some gradient reaches an earlier layer regardless of how small the sublayer's own local derivative is?

Know this

The simplest update: plain gradient descent

The most basic optimizer moves each weight a small step opposite its gradient: w = w - lr × grad, where lr is the learning rate. Since the gradient points in the direction of steepest increase in the loss, subtracting a scaled version of it moves the weight toward lower loss. Too large a learning rate overshoots and can diverge; too small makes training crawl. Plain gradient descent applies this identical rule to every parameter, with no memory of past updates and no distinction between parameters that need large or small steps.

Momentum: smoothing noisy per-batch gradients

Adam improves on this in two ways. First, it tracks a first moment, m, an exponentially decaying running average of past gradients: m = beta1 × m + (1 - beta1) × grad. Rather than reacting to each batch's gradient in isolation, which is noisy since different batches produce somewhat different gradients, the update follows this smoothed running average, letting the optimizer build consistent "velocity" in a direction the gradient has repeatedly pointed toward, rather than jittering with every batch's individual noise.

Adaptive scaling: a per-parameter step size

Second, Adam tracks a second moment, v, an exponentially decaying running average of the squared gradient: v = beta2 × v + (1 - beta2) × grad^2. This estimates each parameter's typical gradient magnitude. Dividing the update by sqrt(v) gives a parameter whose gradients are usually small a relatively larger effective step, and a parameter whose gradients are usually large a relatively smaller one, adapting the step size per parameter rather than using one fixed learning rate for every weight in the model. The full update, with bias correction (m_hat, v_hat, correcting for m and v both starting at zero and being biased toward zero early in training) and a small constant eps to avoid dividing by zero, is:

w = w - lr × m_hat / (sqrt(v_hat) + eps)

AdamW's fix: decouple weight decay from the adaptive update

Weight decay is a regularization technique that shrinks every weight slightly toward zero each step, discouraging overly large weights and generally helping generalization. The original way to add it to Adam folded it directly into the gradient before computing m and v: effectively adding weight_decay × w to grad. That interacts badly with Adam's adaptive scaling: the decay term itself then gets divided by sqrt(v_hat) along with the rest of the gradient, so a parameter with large typical gradients gets its decay shrunk, and one with small typical gradients gets its decay amplified, neither of which is the fixed, uniform shrinkage weight decay is supposed to apply. AdamW fixes this by applying weight decay directly to the weights, separately from the gradient-based Adam update:

w = w - lr × (m_hat / (sqrt(v_hat) + eps) + weight_decay × w)

The decay term is added to the update directly, never passed through m, v, or the adaptive denominator, which is exactly why AdamW ("W" for weight decay) rather than plain Adam is the standard optimizer for training transformers.

Practice

1. ▢ Using plain gradient descent, a weight is w = 5.0, its gradient is 2.0, and the learning rate is 0.1. Compute the updated weight.

2. ▢ Why does tracking a running average of past gradients (Adam's first moment, m) help compared to using each batch's raw gradient directly?

3. ▢ Why does dividing the update by sqrt(v), Adam's second moment, help parameters that have different typical gradient magnitudes?

4. ▢ What goes wrong with folding weight decay directly into the gradient in original Adam, and what does AdamW do differently?

5. ▢ Which claim is true of AdamW compared to plain gradient descent and original Adam?

   • a) AdamW uses the same fixed step size for every parameter, the same as plain gradient descent
   • b) AdamW adds momentum and per-parameter adaptive scaling like Adam, but applies weight decay directly to the weights rather than folding it into the gradient
   • c) AdamW removes adaptive per-parameter scaling entirely, relying only on momentum
   • d) Weight decay in AdamW is scaled by the same adaptive denominator as the gradient-based update

Real-world reps

• [ ] Implement plain gradient descent, then Adam, then AdamW from raw operations for a small set of parameters, and confirm each successive version changes the update in the way this lesson describes.
• [ ] Run a few training steps with weight decay folded into the gradient (original Adam style) versus applied directly (AdamW style) on the same small model, and compare how the weights' magnitudes evolve.
• [ ] Tomorrow: check what optimizer and hyperparameters (learning rate, weight decay, beta1, beta2) a real training script you have access to (such as nanoGPT's) actually uses.

Going further

• Paper: "Decoupled Weight Decay Regularization", Loshchilov and Hutter, 2019
• Paper: "Adam: A Method for Stochastic Optimization", Kingma and Ba, 2015
• Resources

Not landing? Reread the primary source at the top, since this lesson compresses it and compression is where understanding leaks. Check the glossary for any term that felt slippery.

If the lesson itself is unclear rather than the material, that is a defect: open an issue.