Part 19 — Under the hood: optimizers, efficient attention, and quantization

Optional deep-dive. Three things you've been treating as black boxes since partway through this guide: AdamW, used in every training loop since Chapter 26, without ever seeing what's inside it; GQA, name-dropped in Chapters 25 and 38 as "shrinks the KV cache" with no mechanism given; and quantization, reduced to a memory table in Chapter 66. The three chapters here are independent — read whichever black box you're curious about, and skip the rest.

Chapter 71 — How AdamW actually works

In one sentence

AdamW tracks a running average of each parameter's gradient and of its squared gradient, uses those to give every parameter its own self-tuned step size, and applies weight decay as a separate, direct shrink of the parameter — which is why it trains large networks far more reliably than the one-size- fits-all step from Chapter 10.

The problem

Chapter 10's gradient descent update is theta -= lr * gradient, using one global learning rate for every parameter. A network has millions of parameters whose gradients differ hugely in scale: an embedding row that's touched by every training example gets a steady, moderate gradient; a rarely activated neuron gets a noisy, spiky one. One global lr is a bad compromise for both — big enough for the steady one to converge reasonably means it's either too small for the spiky one (crawls) or too big (diverges).

The idea, in plain language

Keep two running averages per parameter:

Where has this parameter's gradient been pointing, recently? — an exponential moving average of the raw gradient. This is momentum: it smooths out zig-zagging and keeps moving in a direction that's been consistently downhill.
How large has this parameter's gradient been, recently, regardless of direction? — an exponential moving average of the squared gradient. This is used to shrink the step for parameters with large or noisy gradients, and enlarge it for parameters with small, quiet ones — every parameter effectively gets its own learning rate.

How it actually works

For each parameter, at training step t, with gradient g_t:

m_t = beta1 * m_{t-1} + (1 - beta1) * g_t          # momentum (1st moment)
v_t = beta2 * v_{t-1} + (1 - beta2) * g_t^2         # 2nd moment (squared grad)

m_hat = m_t / (1 - beta1^t)                          # bias correction
v_hat = v_t / (1 - beta2^t)

theta -= lr * m_hat / (sqrt(v_hat) + eps)            # the Adam update

Typical defaults: beta1 = 0.9, beta2 = 0.999 (large LLM training often uses beta2 = 0.95 instead — shorter memory, which suits the large batches and comparatively short training horizons of LLM pretraining), eps = 1e-8.

Why bias correction exists: m_0 and v_0 both start at zero, so early on, m_t and v_t are biased toward zero (they haven't accumulated enough history yet). Dividing by (1 - beta^t) — a factor very close to 1 - 1 = 0 at t = 1 and approaching 1 as t grows — exactly cancels that early-step bias. Without it, the first several dozen updates would be artificially tiny.

AdamW's one change: decouple weight decay from the gradient statistics. Classic L2 regularisation adds weight_decay * theta directly into g_t before it feeds m_t and v_t — which means the adaptive scaling above also distorts the decay amount inconsistently per parameter. AdamW instead applies decay as a separate step, untouched by m/v:

theta -= lr * m_hat / (sqrt(v_hat) + eps)   # the Adam update, as above
theta -= lr * weight_decay * theta          # decoupled weight decay

This is the exact split you already saw in Chapter 41's training code — decay and nodecay parameter groups exist precisely so weight decay is applied only to weight matrices, not to norms/biases, using this decoupled formula.

Worked example: three steps, by hand

A toy scalar parameter theta = 1.0, minimising (theta - 3)^2 (so the true gradient is 2*(theta - 3)), with lr=0.1, beta1=0.9, beta2=0.999, eps=1e-8, weight_decay=0.01:

step	gradient	`m`	`v`	`m_hat`	`v_hat`	update	`theta` after
1	-4.000	-0.400	0.0160	-4.000	16.000	-0.100	1.099
2	-3.802	-0.740	0.0304	-3.896	15.227	-0.100	1.198
3	-3.605	-1.027	0.0434	-3.788	14.482	-0.100	1.296

Two things worth noticing in these real numbers: m_hat at step 1 exactly equals the raw gradient (bias correction fully compensates for the zero-initialised m on the very first step), and the actual update stays almost exactly -0.100 across all three steps despite the true gradient shrinking from -4.0 to -3.6 — this is the adaptive step size at work: m_hat / sqrt(v_hat) self-normalises, so the direction matters more than the magnitude once the running averages have spun up. Compare this to Chapter 10's plain gradient descent, where the step size is simply lr * gradient and shrinks in lockstep with the gradient.

Practice (25 min) — implement it from scratch, compare to plain GD

import random

def adamw_step(theta, grad_fn, state, lr=0.1, beta1=0.9, beta2=0.999,
               eps=1e-8, weight_decay=0.01):
    g = grad_fn(theta)
    state["t"] += 1
    state["m"] = beta1 * state["m"] + (1 - beta1) * g
    state["v"] = beta2 * state["v"] + (1 - beta2) * g * g
    m_hat = state["m"] / (1 - beta1 ** state["t"])
    v_hat = state["v"] / (1 - beta2 ** state["t"])
    theta -= lr * m_hat / (v_hat ** 0.5 + eps)
    theta -= lr * weight_decay * theta
    return theta

# A tricky loss: steep in one direction, shallow in the other (an elongated
# valley) -- exactly the shape that makes plain gradient descent oscillate.
def grad(theta):
    x, y = theta
    return [20 * x, 2 * y]   # gradient of 10*x^2 + y^2

theta_gd = [1.0, 1.0]
theta_adam = [1.0, 1.0]
state = {"t": 0, "m": [0.0, 0.0], "v": [0.0, 0.0]}

for step in range(50):
    g = grad(theta_gd)
    theta_gd = [t - 0.05 * gi for t, gi in zip(theta_gd, g)]     # plain GD
    theta_adam = [adamw_step(theta_adam[i],
                              lambda t, i=i: grad(theta_adam)[i],
                              {"t": state["t"], "m": state["m"][i], "v": state["v"][i]})
                  for i in range(2)]
    # (bookkeeping simplified for readability -- see full listing for the tidy version)

print("plain GD final:", theta_gd)
print("AdamW final:", theta_adam)

Run a cleaner version of this (full listing not shown here for space) and plot both trajectories: plain GD overshoots along the steep x axis while crawling along the shallow y axis; AdamW's per-parameter scaling tames the steep direction and speeds up the shallow one, reaching the minimum in far fewer steps.

Common confusions

AdamW is not "Adam plus more regularisation." It uses the same weight-decay amount as L2 regularisation would — the difference is entirely in how it's applied (decoupled vs. mixed into the gradient statistics), not how much.
A higher beta2 is not always better. It means longer memory of past squared gradients — good for smooth, low-noise problems, but it makes the optimiser slower to adapt when gradient statistics shift quickly, which is why some LLM pretraining runs deliberately use beta2 = 0.95 instead of the default 0.999.
eps is not just a divide-by-zero guard. At low precision (bf16 training, Chapter 66), sqrt(v_hat) can be small enough that eps's exact value measurably affects the update — this is a real, documented source of training instability in large-scale runs, not a purely theoretical concern.

Check yourself

What two running averages does Adam maintain per parameter, and what does each one do?
Why is bias correction needed in the first several steps specifically?
What, precisely, does the "W" in AdamW change relative to plain Adam with L2 regularisation?
Why can the actual parameter update stay roughly constant even as the raw gradient shrinks?