Kiru Lab  /  Deep Learning  /  Training Dynamics and Failure Modes

Optimizers, Schedules, and Normalization

Everything past plain SGD exists to handle curvature that differs by direction and by time.

Concept  ·  about 30 minutes

Plain gradient descent uses one learning rate for every parameter at every step. Real loss surfaces are anisotropic — steep in some directions, nearly flat in others — and the right step size differs by direction and changes as training proceeds. Optimizers and schedules are the responses.

  • Momentum accumulates a running average of gradients, damping oscillation across steep directions and accelerating along consistent ones.
  • Adam keeps per-parameter running averages of both the gradient and its squared magnitude, effectively giving each parameter its own adaptive step size.
  • Weight decay pulls parameters toward zero; in AdamW it is applied separately from the adaptive scaling, which matters more than it sounds.
  • Warmup starts with a tiny learning rate because early gradients are large and unrepresentative, and one bad early step can be unrecoverable.
  • Cosine decay anneals the rate toward zero so late training refines rather than bounces.

Normalization

Batch normalization standardizes each feature across the batch; layer normalization standardizes across the features of each example. Transformers use layer norm because it does not couple examples in a batch together — which matters enormously at inference, where the batch is often a single sequence.

Layer norm, complete
import numpy as np

def layer_norm(x, gamma, beta, eps=1e-5):
    mu = x.mean(axis=-1, keepdims=True)
    var = x.var(axis=-1, keepdims=True)
    return gamma * (x - mu) / np.sqrt(var + eps) + beta
Note for the interpretability track

Layer norm rescales activations, which means the raw magnitude of a residual-stream vector is not directly comparable across layers. Interpretability work that compares activation norms has to account for this, and a surprising number of confusing results trace back to forgetting it.

Hold on to

  • Adam adapts the step size per parameter using gradient statistics
  • Warmup exists because early gradients are large and unrepresentative
  • Layer norm avoids coupling examples, which is why transformers use it

Work through

Try each one before opening the solution. Getting it wrong first is most of where the learning happens.

  1. Implement momentum and Adam, and compare all three optimizers on the ill-conditioned bowl from track three.
    Hint

    Adam maintains two running averages per parameter: the gradient and its square.

    Solution

    On the 10:1 ill-conditioned bowl, plain SGD zig-zags, momentum straightens the path and roughly halves the steps needed, and Adam is close to insensitive to the conditioning because it normalizes each direction by its own gradient magnitude — the steep direction gets a small effective step and the shallow one a large step, automatically.

    import numpy as np

A = np.array([[10.0, 0.0], [0.0, 1.0]])
grad = lambda p: A @ p

def adam(steps=60, lr=0.1, b1=0.9, b2=0.999, eps=1e-8):
    p, m, v = np.array([1.0, 1.0]), np.zeros(2), np.zeros(2)
    for t in range(1, steps + 1):
        g = grad(p)
        m = b1 * m + (1 - b1) * g
        v = b2 * v + (1 - b2) * g ** 2
        m_hat, v_hat = m / (1 - b1 ** t), v / (1 - b2 ** t)
        p = p - lr * m_hat / (np.sqrt(v_hat) + eps)
    return p

print(np.linalg.norm(adam()))
  2. Train with and without warmup at a high learning rate and show the difference in the first hundred steps.
    Hint

    Ramp the learning rate linearly from near zero over the first few hundred steps.

    Solution

    Without warmup, a high target learning rate applied from step one often produces a large loss spike or immediate divergence, because the initial gradients are large and unrepresentative of the loss surface and one bad early step can be unrecoverable. With warmup the same target rate trains stably. This is why every large-model recipe includes a warmup phase.

    def lr_at(step, target=1e-3, warmup=500):
    if step < warmup:
        return target * (step + 1) / warmup
    return target

# Train the same model twice, with and without warmup, at an aggressive target.
# Plot the first 1000 steps of each loss curve on the same axes.

Sign in to track your progress through the lab.

← Training Dynamics and Failure Modes Reading a Failing Run →