Optimizers

Every optimizer runs the same loop: feel the slope, turn it into a step, take it, repeat. This level builds each rule on a one-number valley you can check by hand, then lets them race across a valley that is not round.

You're on a hillside in thick fog and want to reach the lowest point in the valley. You can't see the valley; you can only feel the slope under your feet. So you feel which way is downhill, take a step that way, and repeat. That is gradient descent. The hillside is the loss (how wrong the model is, for every possible setting of its weights), your position is the current weights, and the slope under your feet is the gradient: the list of slopes of the loss, one per weight. (See primer.ml.neural_net for how backprop measures it, and primer.notation for the symbols.) An optimizer is your rule for turning "the slope here" into "the step I take".

Worked example on the simplest possible valley, the bowl f(w) = w², whose slope at w is 2w. Start at w = 1 with step size 0.1:

Each step keeps 80% of w, sliding smoothly toward the bottom at 0.

flowchart LR
  W[Current weights] --> G[Feel the slope<br/>compute gradient]
  G --> R{Optimizer rule}
  R --> S[Step]
  S --> W2[New weights]
  W2 -->|repeat| G
  R -.uses.-> H[Its own memory:<br/>velocity, averages]

w, eta = 1.0, 0.1
# ∇L for the bowl L = w²
def grad_L(w): return 2 * w
for t in range(2):
    # w_(t+1) = w_t - η ∇L(w_t)
    w = w - eta * grad_L(w)
    print(round(w, 2))  # → 0.8 0.64

"Stochastic" gradient descent (SGD) means the slope is estimated from a small random batch of examples instead of the whole dataset: noisier, but thousands of times cheaper per step. descend_bowl runs the table above; SGD is the general version.

Why it matters every model you've heard of was trained by a descendant of this one line. The variants below exist because real loss landscapes are not round bowls.

In the fog, stride length is everything. Tiny shuffling steps are safe but you'll be walking all night. Giant leaps overshoot the valley floor and land you higher up the opposite slope; keep leaping and you climb out of the valley altogether.

Worked example on the bowl w², where each step multiplies w by (1 − 2η):

def three_steps(eta, w=1.0):
    out = []
    for t in range(3):
        # w - η·2w = (1 - 2η) w
        w = (1 - 2 * eta) * w
        out.append(round(w, 3))
    return out
three_steps(1.1)  # → [-1.2, 1.44, -1.728]
three_steps(0.1)  # → [0.8, 0.64, 0.512]
# η = 0.001 after 10 steps: barely moved
round(0.998 ** 10, 3)  # → 0.98

In code: descend_bowl is the same loop with the learning rate as an argument; call it with each rate in the table to reproduce every row.

Why it matters the learning rate is the single most important hyperparameter. Too high and training diverges or bounces (loss spikes, NaNs); too low and it takes forever or settles somewhere poor. The steepest direction of the landscape sets the ceiling: on a bowl with slope 2w, any η above 1 diverges.

Now imagine a long, narrow valley: steep walls on both sides, a gentle slope along the floor. A cautious hiker who only reads the local slope zig-zags from wall to wall and barely moves along the floor. A heavy ball rolling down the same valley behaves differently: its sideways bouncing cancels out, while the gentle downhill pull along the floor keeps adding up, so it builds speed exactly where you want it. That accumulated speed is momentum.

Worked example on the bowl w², learning rate 0.1, momentum β = 0.9:

After two steps plain descent is at 0.64; momentum is already at 0.46.

flowchart LR
  G[Slope now g_t] --> V["Velocity v_t = β·v_(t−1) + g_t<br/>(remember 90% of the old speed)"]
  VO["Old velocity v_(t−1)"] --> V
  V --> S["Step: w − η·v_t"]
  S --> VO

w, v, beta, eta = 1.0, 0.0, 0.9, 0.1
for t in range(2):
    # the slope of w² here
    g = 2 * w
    # v_t = β v_(t-1) + g_t
    v = beta * v + g
    # w_(t+1) = w_t - η v_t
    w = w - eta * v
    print(round(v, 2), round(w, 2))  # → 2.0 0.8 3.4 0.46

In code: momentum_on_bowl runs the two-row table; SGD with a nonzero momentum keeps its velocity between calls to SGD.step. narrow_valley is the valley in the figure (rosenbrock is a harder, banana-shaped one), and run walks any optimizer across a landscape and records its path.

Why it matters real loss surfaces are full of narrow valleys. On the valley above, 100 steps of momentum reach a loss over 10,000× lower than 100 steps of plain SGD at the same learning rate. Momentum is still the default for training CNNs.

Back in the narrow valley, what you'd really like is short steps across the steep walls and long strides along the gentle floor. Adam does exactly that: it keeps, for every single weight, a running average of the slope (the direction, like momentum) and a running average of the squared slope (how big that weight's slopes typically are), then divides the first by the square root of the second. Every weight ends up taking steps of roughly the same size, the learning rate, whatever the scale of its gradient.

Worked example: on its very first step, Adam moves a weight by exactly the learning rate (0.01 here), no matter whether that weight's gradient is 1000, 1 or 0.001. (With bias correction, m̂ = g and v̂ = g², so the step is 0.01 × g / |g| = 0.01.)

flowchart LR
  G[Gradient g] --> M["m: average of g<br/>(direction)"]
  G --> V["v: average of g²<br/>(typical size)"]
  M --> MC["m̂ = m / (1 − β1^t)"]
  V --> VC["v̂ = v / (1 − β2^t)"]
  MC --> D["step = η · m̂ / (√v̂ + ε)"]
  VC --> D
  D --> W[w − step]

import math
beta_1, beta_2, eta, eps, t = 0.9, 0.999, 0.01, 1e-8, 1
def first_step(g):
    # m_1, starting from m_0 = 0
    m = beta_1 * 0 + (1 - beta_1) * g
    # v_1, starting from v_0 = 0
    v = beta_2 * 0 + (1 - beta_2) * g ** 2
    # undo the pull toward zero
    m_hat = m / (1 - beta_1 ** t)
    v_hat = v / (1 - beta_2 ** t)
    return eta * m_hat / (math.sqrt(v_hat) + eps)
# m_1, v_1
round((1 - beta_1) * 1000, 6), round((1 - beta_2) * 1000 ** 2, 6)  # → (100.0, 1000.0)
# the same step every time
[round(first_step(g), 6) for g in (1000, 1, 0.001)]  # → [0.01, 0.01, 0.01]

In code: Adam keeps the two running averages and the step count for every weight and applies the five formulas in Adam.step; adam_first_step shows the first step is always the learning rate.

Why it matters Adam is forgiving: one learning rate works across weights whose gradients differ by orders of magnitude, which is the norm in transformers (embeddings, attention, layer norms all behave differently). On the valley above it reaches a loss below 10⁻¹⁰ in 300 steps while plain SGD is still around 10⁻³.

Weight decay is a gentle leash that pulls every weight a little toward zero each step, so the model prefers small, smooth weights over large, spiky ones (a form of regularization; see primer.ml.regularization). The classic way to add it was an L2 penalty: add λw to the gradient. With plain SGD that is the same as shrinking the weight. With Adam it isn't, because Adam divides the whole gradient, penalty included, by its typical size. The leash's strength gets rescaled away.

Worked example: one step with zero loss-gradient, so only the decay acts. w = 1, learning rate 0.1.

With L2 inside Adam, a 100× weaker penalty shrinks the weight exactly as much: λ has stopped meaning what it says.

flowchart TB
  subgraph L2["Adam + L2 penalty"]
    g1[loss gradient] --> add["+ λ·w"] --> ad1[Adam rescaling<br/>÷ √v̂] --> s1[step]
  end
  subgraph AW["AdamW (decoupled)"]
    g2[loss gradient] --> ad2[Adam rescaling<br/>÷ √v̂] --> s2[step]
    w2[weights] --> dec["shrink: w − η·λ·w"] --> s2
  end

# no loss gradient: Adam's step is 0
w, eta, adam_step = 1.0, 0.1, 0.0
for lam in (0.1, 0.001):
    # AdamW: shrink by η λ w
    print(round(w - eta * adam_step - eta * lam * w, 4))  # → 0.99 0.9999
for lam in (0.1, 0.001):
    g = lam * w
    print(round(w - eta * g / abs(g), 4))  # → 0.9 0.9

In code: Adam implements both recipes: with decoupled decay it is AdamW, otherwise it adds the L2 penalty to the gradient. one_decay_step runs the table.

Why it matters AdamW is the default optimizer for transformers. The fix was a one-line change that made weight decay behave predictably and improved generalization.

Think of driving an unfamiliar car. You ease onto the accelerator at first, because you don't yet know how it responds. You cruise at speed for most of the trip. Near the destination you slow gradually and glide into the parking spot. Transformers are trained the same way: the learning rate warms up linearly from 0 to its peak, then decays along a cosine curve toward a small floor.

Worked example with peak 0.001, 100 warmup steps, 1,000 steps total:

import math
eta_max, eta_min, T_w, T = 0.001, 0.0, 100, 1000
def eta(t):
    if t < T_w:
        # the straight ramp
        return eta_max * t / T_w
    # 0 to 1 through the decay
    progress = (t - T_w) / (T - T_w)
    return eta_min + (eta_max - eta_min) * (1 + math.cos(math.pi * progress)) / 2
[round(eta(t), 6) for t in (50, 100, 550, 1000)]  # → [0.0005, 0.001, 0.0005, 0.0]

In code: warmup_cosine returns the learning rate for any step: the straight ramp during warmup, then the half cosine down to the floor.

Why it matters at the very start, Adam's averages are unreliable and the weights are random, so a full-size step can wreck them; warmup avoids early divergence. The slow finish lets the model settle into a good minimum instead of bouncing around it.

Occasionally one bad batch produces an enormous gradient, a sudden cliff in the fog. Taking a full step along it could throw the weights far from anywhere useful. Clipping is a speed limiter: if the step would be longer than a set limit, shorten it to the limit, keeping its direction.

Worked example: the gradient (3, 4) has length √(3² + 4²) = 5. With a limit of 1, scale it by 1/5 to get (0.6, 0.8), length 1, same direction. A gradient of (0.3, 0.4) has length 0.5, under the limit, so it's left alone.

flowchart LR
  G[All gradients] --> N["Global length<br/>‖g‖ = √(sum of every squared entry)"]
  N --> C{"‖g‖ > limit?"}
  C -->|no| K[Use as is]
  C -->|yes| S["Multiply every gradient<br/>by limit / ‖g‖"]
  S --> K2[Same direction,<br/>length = limit]

import math
def clip(g, c):
    # ‖g‖ = √(Σ g_i²)
    norm = math.sqrt(sum(g_i ** 2 for g_i in g))
    # shrink only if too long
    scale = min(1, c / norm)
    return [round(g_i * scale, 2) for g_i in g]
clip([3, 4], c=1)  # → [0.6, 0.8]
# length 0.5: under the limit, left alone
clip([0.3, 0.4], c=1)  # → [0.3, 0.4]

In code: clip_by_global_norm measures the length of all gradients together and scales every one by the same factor when it exceeds the limit.

Why it matters large-model training runs almost always clip (a limit of 1.0 is common). It turns rare loss spikes from run-ending disasters into harmless blips. It's measured globally, across all layers together, so the update's direction is preserved (see primer.ml.deep_nets for exploding gradients).

Answer each one out loud or on paper before you open it. If you can explain it, you know it.