NB2 §5 — Normalization, weight decay, and instabilities

Scale invariance (§5.1).

The weight-decay equilibrium (§5.2).

Loss spikes (§5.3).

Clipping (§5.4).

Also: Li, Lyu & Arora, the intrinsic learning rate (names \eta\lambda, conjectures equilibrium in O(1/\eta\lambda) steps; 5.2); Lobacheva et al., periodic behavior with BatchNorm and weight decay (mechanism 1); D’Angelo et al., why we need weight decay (argues it is never useful as an explicit regularizer; acts through \eta/\|w\| under AdamW); Wan et al. 2021, spherical motion dynamics (the per-step SGD equilibrium angle under explicit noise assumptions).

Take one block of weights w, with norm r=\|w\| and direction u=w/r, that feeds a LayerNorm, RMSNorm or BatchNorm. Then the loss obeys L(cw)=L(w) for every c>0: it is a function on the unit sphere. Its gradient is orthogonal to w and shrinks like 1/r, its sharpness (top Hessian eigenvalue) shrinks like 1/r^2, and one SGD step with learning rate \eta and weight decay \lambda is exactly a gradient step on the sphere (then renormalized) with the effective learning rate \eta_{\rm eff}=\eta/((1-\eta\lambda)r^2).

Paper: van Laarhoven, L2 regularization versus batch and weight normalization (see Readings).

How far this reaches. In a BatchNorm convnet almost every weight is scale-invariant (BatchNorm’s invariance is per output channel, i.e. per row of W); the exceptions are the final classifier, biases and the norm layers’ gains. In a pre-LN transformer few matrices are exactly invariant: scaling an attention or MLP output projection changes that block’s share of the residual stream. QK-layernorm (5.3) makes each head’s W_Q,W_K invariant. Everything below is exact for scale-invariant blocks and a model elsewhere.

Derivations and walkthrough. Setup: LayerNorm and RMSNorm divide a vector by its own scale, so \mathrm{LN}(c\,h)=\mathrm{LN}(h) for c>0, exactly when the stabilizer \epsilon is zero and to excellent approximation when h’s variance (for RMSNorm, mean square) is \gg\epsilon. If h=Wx feeds it, the loss is unchanged by W\to cW. Notation: S(w)=\lambda_{\max}(\nabla^2L(w)) is the sharpness; \lambda is always the weight decay. A reference like “5.1.2.3” means Exercise 5.1.2, part 3.

Exercise 5.1.1 — five identities. Let L be scale-invariant and twice differentiable, with u=w/\|w\|. Show:

1. \langle\nabla L(w),w\rangle=0: the gradient is tangent to the sphere through w.
2. \nabla L(cw)=\nabla L(w)/c.
3. \nabla^2L(cw)=\nabla^2L(w)/c^2. Hence S(w)=S(u)/\|w\|^2: shrinking the weights makes the loss sharper.
4. \nabla^2L(w)\,w=-\nabla L(w). In particular w is a null direction of the Hessian at every critical point.
5. Under gradient flow \dot w=-\nabla L(w), \|w\| is conserved.

(The Mysteries day’s MDL4 §11, Exercise 11.2, derives 1–4 for batch normalization.)

The next cell checks 1–4 numerically, by autograd, on a two-layer net whose first-layer weights feed a LayerNorm (with \epsilon=0, so the invariance is exact).

torch.manual_seed(0)
X, Y = torch.randn(32, 8, dtype=torch.float64), torch.randn(32, 1, dtype=torch.float64)
W2 = torch.randn(1, 16, dtype=torch.float64) / 4

def layer_norm(h):                      # LayerNorm with eps = 0 and no gain: exactly scale-invariant
    h = h - h.mean(1, keepdim=True)
    return h / h.pow(2).mean(1, keepdim=True).sqrt()

def loss(W1_flat):                      # L(W1) for a 8 -> 16 -> LN -> tanh -> 1 network
    h = layer_norm(X @ W1_flat.view(16, 8).T)
    return 0.5 * ((torch.tanh(h) @ W2.T - Y) ** 2).mean()

w = torch.randn(128, dtype=torch.float64)
grad = lambda v: torch.autograd.functional.jacobian(loss, v)
hess = lambda v: torch.autograd.functional.hessian(loss, v)
g, H = grad(w), hess(w)
sharp = lambda M: torch.linalg.eigvalsh(M)[-1].item()

print(f"L(3w) - L(w)                      = {(loss(3 * w) - loss(w)).item():.1e}")
print(f"cos(grad L(w), w)                 = {(g @ w / g.norm() / w.norm()).item():.1e}")
print(f"|3 grad L(3w) - grad L(w)|        = {(3 * grad(3 * w) - g).norm().item():.1e}")
print(f"|H(w) w + grad L(w)|              = {(H @ w + g).norm().item():.1e}")
print(f"S(w) / S(2w)  (expect 4)          = {sharp(H) / sharp(hess(2 * w)):.6f}")

L(3w) - L(w)                      = 0.0e+00
cos(grad L(w), w)                 = 2.1e-17
|3 grad L(3w) - grad L(w)|        = 8.3e-17
|H(w) w + grad L(w)|              = 1.0e-16
S(w) / S(2w)  (expect 4)          = 4.000000

Exercise 5.1.2 — one step, exactly. What does one step of SGD with weight decay do to the direction and to the norm? It is GD on the unit sphere with a step size set by the norm, exactly, not just for small steps; the exactness is what makes 5.2 work.

Write w_k=r_ku_k with r_k=\|w_k\|, \|u_k\|=1, and take one step of SGD with learning rate \eta and weight decay \lambda, w_{k+1}=(1-\eta\lambda)w_k-\eta\nabla L(w_k), with 0\le\eta\lambda<1. (Read \nabla L as a minibatch gradient; each minibatch loss is itself scale-invariant.) Write G_k=\|\nabla L(u_k)\| for the gradient norm at the unit-norm point, so that \|\nabla L(w_k)\|=G_k/r_k. Recall that GD on a quadratic of curvature h is stable iff \eta h<2.

1. Show that the new direction is exactly u_{k+1}=\frac{u_k-\eta_{\mathrm{eff},k}\nabla L(u_k)}{\big\|u_k-\eta_{\mathrm{eff},k}\nabla L(u_k)\big\|}, \qquad \eta_{\mathrm{eff},k}=\frac{\eta}{(1-\eta\lambda)\,r_k^2}. This is a gradient step on the sphere followed by renormalization, with learning rate \eta_{\mathrm{eff}}. Without weight decay, \eta_{\mathrm{eff}}=\eta/r_k^2.
2. Show that the norm obeys the exact recursion r_{k+1}^2=(1-\eta\lambda)^2r_k^2+\frac{\eta^2G_k^2}{r_k^2}.
3. Let \theta_k be the angle between w_k and w_{k+1}. Show that \cos\theta_k=(1-\eta\lambda)\,r_k/r_{k+1} and \tan\theta_k=\eta_{\mathrm{eff},k}G_k.
4. Let u^* be a minimum of L on the sphere. Show that the condition for the direction dynamics to be linearly stable at u^*, \eta_{\mathrm{eff}}\,S(u^*)<2, is the same as the Euclidean condition \eta\,S(w)<2 (up to the factor 1-\eta\lambda): “effective learning rate” and “sharpness of w” are two ways of saying the same thing.

On a scale-invariant block (notation of 5.1: r=\|w\|, learning rate \eta, weight decay \lambda), gradient noise grows the norm and weight decay shrinks it. Without weight decay the norm only grows, so the effective learning rate \eta/r^2 decays by itself, like 1/\sqrt k at step k. With it, the norm settles where each step turns the direction by \theta^*=\arccos(1-\eta\lambda)\approx\sqrt{2\eta\lambda}, independent of the gradient scale: exactly under SGD, approximately under AdamW.

The lead paper is Kosson et al., who call this state rotational equilibrium; the exact SGD statement below is Li and Arora’s (both in the Readings).

The core idea. On these weights weight decay cannot prefer one function over another; its only channel is the norm, hence \eta_{\rm eff} (any effect on generalization goes through that, e.g. through the SGD noise scale, NB2 §2). So under SGD, \eta and \lambda matter only through \eta\lambda, Li, Lyu and Arora’s intrinsic learning rate, and weight decay can be traded exactly for an exponentially growing learning rate. Under AdamW the norm settles at \approx\sqrt{\eta C/(2\lambda)} for a block of C weights, and the step felt by the direction scales as \eta/\|w\|, not \eta/\|w\|^2. Reaching equilibrium takes \sim1/(4\eta\lambda) steps under SGD, often a sizeable part of a run.

Derivations and walkthrough. Setup: SGD with weight decay is w_{k+1}=(1-\eta\lambda)w_k-\eta\nabla L(w_k) (PyTorch’s SGD(weight_decay=λ) without momentum). With r_k=\|w_k\|, G_k the gradient norm at the unit-norm point u_k=w_k/r_k, and \theta_k the angle between w_k and w_{k+1}, Exercise 5.1.2 gives, exactly, r_{k+1}^2=(1-\eta\lambda)^2r_k^2+\frac{\eta^2G_k^2}{r_k^2},\qquad \eta_{\mathrm{eff},k}=\frac{\eta}{(1-\eta\lambda)r_k^2},\qquad \cos\theta_k=(1-\eta\lambda)\frac{r_k}{r_{k+1}},\qquad \tan\theta_k=\eta_{\mathrm{eff},k}G_k.

Exercise 5.2.1 — norm growth without weight decay. With \lambda=0 the norm only grows, so the effective learning rate keeps falling. What schedule is that over a long run? (Arora, Li and Lyu 2019 show that GD with a fixed learning rate on scale-invariant weights still converges at the rate of well-tuned GD.)

1. Show r_{k+1}^4\ge r_k^4+2\eta^2G_k^2, with near-equality when the angular step \tan\theta_k=\eta G_k/r_k^2 is small.
2. Suppose G_k\approx G is roughly constant (as when minibatch noise dominates). Show that for large k, \eta_{\rm eff,k}\approx1/(G\sqrt{2k}) and \theta_k\approx1/\sqrt{2k}, independently of \eta.
3. What does this say about tuning the learning rate of normalized layers trained without weight decay for a long time?

A toy for checking the formulas. Take the scale-invariant loss L(w)=\tfrac12\,u^\top Au with u=w/\|w\| and A=\mathrm{diag}(a_1,\dots,a_d) (a Rayleigh quotient). Minibatch noise is modelled by adding \sigma\xi, \xi standard Gaussian, to the gradient at u and projecting onto the tangent space. (Real minibatch gradients of a scale-invariant network are also orthogonal to w, so only the noise distribution is a model.) The simulator runs several (\eta,\lambda,\sigma) settings in parallel.

Code (36 lines)

d_toy = 20
a_toy = np.linspace(0.0, 1.0, d_toy)          # L(u) = 1/2 u^T diag(a) u on the unit sphere

def sgd_scale_invariant(eta, lam, sigma, T, r0=1.0, seed=0, schedule=None):
    # SGD + weight decay on the noisy scale-invariant toy; one row per (eta, lam, sigma) setting.
    # schedule(k) -> per-step learning rates (overrides eta); used for the exponential-LR check.
    eta, lam, sigma, r0 = (np.array(z, float) for z in np.broadcast_arrays(
        np.atleast_1d(eta), np.atleast_1d(lam), np.atleast_1d(sigma), np.atleast_1d(r0)))
    B = len(eta); rng = np.random.default_rng(seed)
    w = rng.standard_normal((B, d_toy))
    w *= r0[:, None] / np.linalg.norm(w, axis=1, keepdims=True)
    R, COS, G2 = np.empty((T, B)), np.empty((T, B)), np.empty((T, B)); U = np.empty((T, B, d_toy))
    for k in range(T):
        r = np.linalg.norm(w, axis=1, keepdims=True); u = w / r
        gu = a_toy * u + sigma[:, None] * rng.standard_normal(w.shape)
        gu -= (gu * u).sum(1, keepdims=True) * u          # tangent part: the (noisy) gradient at u
        lr = eta if schedule is None else schedule(k)
        w_new = (1 - lr * lam)[:, None] * w - lr[:, None] * gu / r
        R[k], G2[k], U[k] = r[:, 0], (gu**2).sum(1), u
        COS[k] = (w * w_new).sum(1) / r[:, 0] / np.linalg.norm(w_new, axis=1)
        w = w_new
    return dict(r=R, cos=np.clip(COS, -1, 1), G2=G2, u=U)

# --- no weight decay: the effective learning rate forgets eta ---
etas = np.array([0.01, 0.1, 1.0])
out = sgd_scale_invariant(etas, 0.0, 0.5, 20000)
k = np.arange(1, 20001)
fig, ax = plt.subplots(1, 2, figsize=(10, 3.6))
for i, e in enumerate(etas):
    ax[0].loglog(k, e / out["r"][:, i] ** 2, label=f"eta = {e}")
    ax[1].loglog(k, out["r"][:, i], label=f"eta = {e}")
G_rms = np.sqrt(out["G2"].mean())
ax[0].loglog(k, 1 / (G_rms * np.sqrt(2 * k)), "k--", label=r"$1/(G\sqrt{2k})$")
ax[0].set(xlabel="step k", ylabel=r"$\eta/\|w\|^2$", title="effective LR, no weight decay")
ax[1].set(xlabel="step k", ylabel=r"$\|w\|$", title="the norm grows like $k^{1/4}$")
ax[0].legend(); ax[1].legend(); plt.tight_layout(); plt.show()

After a transient (longer for smaller \eta), all three effective learning rates fall onto the dashed line 1/(G\sqrt{2k}), G the RMS gradient norm: the base learning rate has been forgotten, as 5.2.1.2 predicts.

Exercise 5.2.2 — the weight-decay equilibrium. With weight decay on, where does the norm settle, and how far does the direction turn per step once it is there? Use the recursions in the setup above. Take \lambda>0, and first take G_k=G constant.

1. Find the fixed point of the norm recursion and show r_*^4=\frac{\eta G^2}{\lambda(2-\eta\lambda)}\approx\frac{\eta G^2}{2\lambda}.
2. Show that at the fixed point \cos\theta^*=1-\eta\lambda exactly, so \theta^*=\arccos(1-\eta\lambda)=\sqrt{2\eta\lambda}\,\big(1+O(\eta\lambda)\big), independent of G. Show also that \eta_{\rm eff}^*=\tan\theta^*/G\approx\sqrt{2\eta\lambda}/G.
3. Linearize the recursion for s=r^2 about the fixed point and show that deviations shrink by the factor 2(1-\eta\lambda)^2-1\approx1-4\eta\lambda per step. How many steps does it take to reach equilibrium with \eta=0.1, \lambda=5\times10^{-4}?
4. Noise. Now let G_k be random (minibatch noise), with \mathbb E[G_k^2\mid w_k]=\bar G^2. Using 5.1.2.3, show that whenever \log r_k is stationary (its distribution no longer changes with k), \mathbb E\big[\log\cos\theta_k\big]=\log(1-\eta\lambda)\qquad\text{exactly.} Conclude that \mathbb E\,\theta_k^2\le-2\log(1-\eta\lambda), with near-equality, so that the RMS angle is \approx\sqrt{2\eta\lambda}, when the per-step angles are all small. What about the mean angle?
5. Still with noise, show that \mathbb E[r_{k+1}^2\mid w_k]=(1-\eta\lambda)^2r_k^2+\eta^2\bar G^2/r_k^2 exactly, and hence, if r fluctuates little, r_*^4\approx\eta\bar G^2/(\lambda(2-\eta\lambda)) with \bar G the RMS gradient norm at the unit-norm point. (The next cell compares this with the measured norm.)

Code (27 lines)

# Weight-decay equilibrium: five settings, including two pairs with the same eta*lam.
eta = np.array([0.1, 0.1, 0.4, 0.05, 0.2])
lam = np.array([1e-2, 1e-2, 2.5e-3, 2e-2, 2e-2])
sig = np.array([0.3, 1.0, 1.0, 0.3, 0.1])
T = 8000
out = sgd_scale_invariant(eta, lam, sig, T)
theta = np.arccos(out["cos"]); burn = T // 4

fig, ax = plt.subplots(1, 2, figsize=(10, 3.6))
win = 400
for i in range(len(eta)):
    rms = np.sqrt(np.convolve(theta[:, i] ** 2, np.ones(win) / win, mode="valid"))
    ax[0].plot(np.arange(len(rms)) + win, rms, label=f"eta={eta[i]}, lam={lam[i]}, sigma={sig[i]}")
    ax[1].plot(out["r"][:, i])
for el in np.unique(eta * lam):
    ax[0].axhline(np.arccos(1 - el), color="k", ls="--", lw=1)
ax[0].set(xlabel="step", ylabel="RMS angle per step (window 400)", title=r"dashed: $\arccos(1-\eta\lambda)$")
ax[1].set(xlabel="step", ylabel=r"$\|w\|$", title=r"norm relaxes on a $1/(4\eta\lambda)$ timescale")
ax[0].legend(fontsize=7); plt.tight_layout(); plt.show()

print(" eta     lam     sigma | RMS angle  mean angle  arccos(1-el)  sqrt(2el) | E[-log cos]  -log(1-el) | <r>     r* pred")
for i in range(len(eta)):
    th, el = theta[burn:, i], eta[i] * lam[i]
    r_pred = (eta[i] * out["G2"][burn:, i].mean() / (lam[i] * (2 - el))) ** 0.25
    print(f"{eta[i]:5.2f} {lam[i]:7.4f} {sig[i]:6.2f}   | {np.sqrt((th**2).mean()):.5f}    {th.mean():.5f}"
          f"     {np.arccos(1 - el):.5f}      {np.sqrt(2 * el):.5f} | {-np.log(out['cos'][burn:, i]).mean():.6f}"
          f"     {-np.log(1 - el):.6f}  | {out['r'][burn:, i].mean():.3f}   {r_pred:.3f}")

 eta     lam     sigma | RMS angle  mean angle  arccos(1-el)  sqrt(2el) | E[-log cos]  -log(1-el) | <r>     r* pred
 0.10  0.0100   0.30   | 0.04471    0.04414     0.04473      0.04472 | 0.001000     0.001001  | 1.717   1.717
 0.10  0.0100   1.00   | 0.04470    0.04411     0.04473      0.04472 | 0.000999     0.001001  | 3.117   3.117
 0.40  0.0025   1.00   | 0.04473    0.04415     0.04473      0.04472 | 0.001001     0.001001  | 6.248   6.248
 0.05  0.0200   0.30   | 0.04476    0.04418     0.04473      0.04472 | 0.001002     0.001001  | 1.212   1.213
 0.20  0.0200   0.10   | 0.08946    0.08830     0.08947      0.08944 | 0.004007     0.004008  | 1.001   1.001

(In the table, el is \eta\lambda.) Four of the five settings share \eta\lambda=10^{-3}, with three different (\eta,\lambda) pairs and two noise levels, and rotate by the same angle. The noise level \sigma (equivalently the gradient scale) moves only the norm. The stationarity identity \mathbb E[-\log\cos\theta]=-\log(1-\eta\lambda) holds to about 0.2%.

Exercise 5.2.3 — only \eta\lambda matters. On scale-invariant weights, do \eta and \lambda matter separately, or only through \eta\lambda (Li, Lyu and Arora’s intrinsic learning rate)? With the setup’s recursions, define the rescaled squared norm \tilde s_k=r_k^2/\eta.

1. Show that \tilde s_{k+1}=(1-\eta\lambda)^2\tilde s_k+G_k^2/\tilde s_k and \eta_{\mathrm{eff},k}=1/\big((1-\eta\lambda)\tilde s_k\big).
2. Conclude: two runs with the same product \eta\lambda, the same initial direction, the same initial r_0^2/\eta and the same minibatch sequence produce identical direction trajectories u_k, hence identical network functions on scale-invariant layers, even if \eta and \lambda differ individually. With different r_0^2/\eta, what can you say?
3. Suppose you multiply \eta by 10 and divide \lambda by 10 on a scale-invariant layer. What changes?

Exercise 5.2.4 — weight decay is an exponentially increasing learning rate. Weight decay acts here only by shrinking the norm, which raises \eta_{\rm eff}. Can it be traded exactly for a learning-rate schedule? This is Li and Arora’s theorem, without momentum. Let L be scale-invariant (the paper freezes the last layer and the norm layers’ affine parameters to arrange this), and run SGD with weight decay, w_{k+1}=(1-\eta\lambda)w_k-\eta\nabla L_k(w_k), from w_0, where L_k is the minibatch loss at step k.

1. Define \tilde w_k=(1-\eta\lambda)^{-k}w_k. Show that \tilde w_{k+1}=\tilde w_k-\tilde\eta_k\nabla L_k(\tilde w_k) with \tilde\eta_k=\eta\,(1-\eta\lambda)^{-2k-1}, i.e. plain SGD, no weight decay, with an exponentially increasing learning rate, from the same \tilde w_0=w_0. Hint: divide the update by (1-\eta\lambda)^{k+1}, then use \nabla L(cw)=\nabla L(w)/c (5.1.1.2) to rewrite \nabla L_k(w_k) in terms of \tilde w_k.
2. In what sense are the two runs “equivalent”? Are the iterates equal?
3. With \eta\lambda=5\times10^{-5} (e.g. \eta=0.1, \lambda=5\times10^{-4}), by what factor does \tilde\eta grow per epoch of 391 steps (CIFAR-10 at batch 128)? Why do the two pictures (shrinking weights at fixed \eta, versus growing \eta at growing weights) give the same effective learning rate?

The momentum version (stated, not derived): with heavy-ball momentum \gamma, the same holds with \tilde\eta_k=\eta\,\alpha^{-2k-1}, where \alpha is the larger root (it \to1 as \eta\lambda\to0) of x^2-(1+\gamma-\eta\lambda)x+\gamma=0 (real if \eta\lambda\le(1-\sqrt\gamma)^2), plus a matching initial momentum buffer. (If you read the paper’s Theorem 2.9: the rescaling should be \tilde\theta_t=\alpha^{-t}\theta_t, as its own initial condition and Theorem 2.1 require; the printed \alpha^t is a typo.)

Code exercise — check 5.2.3 and 5.2.4. sgd_scale_invariant (above) draws the same noise for the same seed and accepts a per-step learning-rate schedule. (a) Run SGD with weight decay and the exponential-learning-rate run without it from the same w_0; compare their directions u_k. (b) Compare (\eta,\lambda,r_0)=(0.1,10^{-2},1) with (0.4,2.5\times10^{-3},2), which share \eta\lambda and r_0^2/\eta.

# (a) Li-Arora: SGD + WD versus plain SGD with lr_k = eta * (1 - eta*lam)**(-2k-1), same seed.
# (b) same eta*lam and same r0^2/eta: are the direction trajectories identical?
# Hint: sgd_scale_invariant(..., schedule=lambda k: np.array([...])) and compare out["u"].
...

T, e, l = 400, 0.1, 1e-2
wd = sgd_scale_invariant(e, l, 0.3, T, seed=7)
exp_lr = sgd_scale_invariant(e, 0.0, 0.3, T, seed=7,
                             schedule=lambda k: np.array([e * (1 - e * l) ** (-2 * k - 1)]))
print("(a) max |u_WD - u_expLR| over 400 steps:", np.abs(wd["u"] - exp_lr["u"]).max())
print("    norm ratio r_expLR / r_WD at k = 0, 200, 399:",
      (exp_lr["r"] / wd["r"])[[0, 200, 399], 0], " vs (1-eta*lam)^-k:", (1 - e * l) ** -np.array([0, 200, 399]))

A = sgd_scale_invariant(0.1, 1e-2, 0.3, 2000, r0=1.0, seed=3)
B = sgd_scale_invariant(0.4, 2.5e-3, 0.3, 2000, r0=2.0, seed=3)
print("(b) max |u_A - u_B| over 2000 steps:", np.abs(A["u"] - B["u"]).max())
print("    r^2/eta at the end:", A["r"][-1, 0] ** 2 / 0.1, B["r"][-1, 0] ** 2 / 0.4)

(a) max |u_WD - u_expLR| over 400 steps: 1.1102230246251565e-15
    norm ratio r_expLR / r_WD at k = 0, 200, 399: [1.         1.22152499 1.49063117]  vs (1-eta*lam)^-k: [1.         1.22152499 1.49063117]
(b) max |u_A - u_B| over 2000 steps: 0.0
    r^2/eta at the end: 29.32583798154646 29.32583798154646

Both differences are at roundoff level (roundoff in the exponential run grows slowly with its norm). The two runs in (b) have learning rates four times apart and compute the same function at every step.

AdamW. The SGD analysis rested on \|\nabla L(w)\|=G/r. Adam’s update does not scale like 1/r, because dividing by \sqrt{\hat v} cancels the gradient’s scale, so the norm recursion changes. Notation: AdamW with gradient g_k keeps m_k=\beta_1m_{k-1}+(1-\beta_1)g_k and v_k=\beta_2v_{k-1}+(1-\beta_2)g_k^2 (coordinatewise), bias-corrects them to \hat m_k=m_k/(1-\beta_1^k), \hat v_k=v_k/(1-\beta_2^k), and on one block of C weights steps w_{k+1}=(1-\eta\lambda)w_k-\eta U_k with U_k=\hat m_k/(\sqrt{\hat v_k}+\epsilon).

What comes out (derived for \beta_1=0 in Exercise 5.2.5; the momentum factors are Kosson et al.’s), assuming noise-dominated gradients, small relative updates, and \eta\lambda\ll1-\beta_1 with \hat v tracking \mathbb Eg^2:

• At equilibrium the norm is r_*\approx\sqrt{\eta C/(2\lambda)}, independent of the gradient scale, and the RMS angle per step is \approx\sqrt{2\eta\lambda\,(1-\beta_1)/(1+\beta_1)}.
• At any norm r the RMS angle per step is \approx(\eta/r)\sqrt{C(1-\beta_1)/(1+\beta_1)}: the effective learning rate scales as \eta/\|w\|, not \eta/\|w\|^2.
• Approach: the norm relaxes on a timescale \sim1/(2\eta\lambda) steps. Unrolled, w_k=(1-\eta\lambda)^kw_0-\eta\sum_{j<k}(1-\eta\lambda)^{k-1-j}U_j: the weights are an exponential average of past updates with memory \sim1/(\eta\lambda) steps.

With Adam plus an L^2 penalty instead, the decay is normalized along with the gradient, so the rotation rate depends on each layer’s gradient scale: Kosson et al. measure a roughly 30\times spread across layers, against nearly uniform rotation for AdamW, and hypothesize that this balance is AdamW’s main benefit.

Exercise 5.2.5 — the AdamW equilibrium. Take \beta_1=0 (so \hat m_k=g_k), \epsilon\approx0, and r_k=\|w_k\|. Assume (as Kosson et al. do) a noise-dominated “random walk”: \mathbb E\langle U_k,w_k\rangle=0 and \mathbb E\|U_k\|^2=C, the latter because each coordinate of U_k=g_k/\sqrt{\hat v_k} has second moment \approx1 when \hat v tracks \mathbb Eg^2.

1. Show that the expected squared norm obeys \mathbb E r_{k+1}^2=(1-\eta\lambda)^2\,\mathbb Er_k^2+\eta^2C, with equilibrium r_*^2=\eta C/(\lambda(2-\eta\lambda))\approx\eta C/(2\lambda), approached at rate \approx2\eta\lambda per step.
2. Show that the RMS angle per step at equilibrium is again \theta^*\approx\sqrt{2\eta\lambda}, i.e. \cos\theta^*\approx1-\eta\lambda.
3. Show that the angle per step is \approx\eta\sqrt C/r. In what sense is the effective learning rate of AdamW \propto\eta/\|w\| rather than \eta/\|w\|^2?

With momentum (Kosson et al., Table 1). For AdamW the equilibrium norm is unchanged but the RMS angle per step shrinks to \sqrt{2\eta\lambda(1-\beta_1)/(1+\beta_1)}: momentum spreads each gradient over \sim1/(1-\beta_1) correlated steps, so each step rotates less. For SGD with heavy-ball momentum \gamma their prediction is \sqrt{2\eta\lambda/(1+\gamma)}. D’Angelo et al. (Readings) report, for sign-SGD, that matching \eta/\|w\| (not \eta/\|w\|^2) reproduces AdamW’s training curves without weight decay in float32, while the bfloat16 runs diverge. The next cell tests the random-walk assumption.

def adamw_scale_invariant(eta, lam, sigma, T, b1=0.9, b2=0.999, eps=1e-12, d=64, seed=0):
    # AdamW on the noisy Rayleigh toy in d dimensions; returns per-step angles and norms.
    a = np.linspace(0, 1, d); rng = np.random.default_rng(seed)
    w = rng.standard_normal(d); m = np.zeros(d); v = np.zeros(d); ang, rs, rad = [], [], []
    for k in range(1, T + 1):
        r = np.linalg.norm(w); u = w / r
        q = a * u + sigma * rng.standard_normal(d); q -= (q @ u) * u
        g = q / r
        m = b1 * m + (1 - b1) * g; v = b2 * v + (1 - b2) * g * g
        w_new = (1 - eta * lam) * w - eta * (m / (1 - b1**k)) / (np.sqrt(v / (1 - b2**k)) + eps)
        U = (m / (1 - b1**k)) / (np.sqrt(v / (1 - b2**k)) + eps)
        ang.append(np.arccos(np.clip(w @ w_new / r / np.linalg.norm(w_new), -1, 1))); rs.append(r)
        rad.append(u @ U / np.linalg.norm(U))        # cosine between the Adam update and w
        w = w_new
    return np.array(ang), np.array(rs), np.array(rad)

print(" sigma  beta1 | RMS angle  Kosson sqrt(2 el (1-b1)/(1+b1)) | <r>    sqrt(eta C/(2 lam)) | <cos(u, U)>")
eta_a, lam_a, C = 1e-2, 0.1, 64
for sigma in [1.0, 0.1]:
    for b1 in [0.0, 0.9]:
        T = int(10 / (eta_a * lam_a)); th, rs, rad = adamw_scale_invariant(eta_a, lam_a, sigma, T, b1=b1)
        print(f"  {sigma:4.1f}   {b1:3.1f}  |  {np.sqrt((th[T//2:]**2).mean()):.4f}       "
              f"{np.sqrt(2*eta_a*lam_a*(1-b1)/(1+b1)):.4f}                     | {rs[T//2:].mean():.2f}   "
              f"{np.sqrt(eta_a*C/(2*lam_a)):.2f}               | {rad[T//2:].mean():+.3f}")

 sigma  beta1 | RMS angle  Kosson sqrt(2 el (1-b1)/(1+b1)) | <r>    sqrt(eta C/(2 lam)) | <cos(u, U)>
   1.0   0.0  |  0.0427       0.0447                     | 1.91   1.79               | -0.002
   1.0   0.9  |  0.0098       0.0103                     | 1.91   1.79               | -0.096
   0.1   0.0  |  0.0218       0.0447                     | 3.47   1.79               | -0.036
   0.1   0.9  |  0.0050       0.0103                     | 3.47   1.79               | -0.201

On the toy in C=64 dimensions (\sigma the noise level), the predictions hold to about 5–10% when noise dominates (\sigma=1). When signal is comparable to noise (\sigma=0.1) they fail: the norm settles about twice as large and rotates about half as fast. The last column shows why: Adam’s coordinatewise normalization breaks orthogonality to w. At \beta_1=0 the cosine of u with U is only about -0.04, but this radial term enters the norm recursion at first order in \eta (-2\eta(1-\eta\lambda)\langle w,U\rangle\approx0.02 per step) and rotation only at second order (\eta^2C\approx0.006). (With \beta_1=0.9 the cosine is larger even when the predictions hold, because momentum points partly along past tangent directions.) The SGD identity \mathbb E\log\cos\theta=\log(1-\eta\lambda), by contrast, needs only scale invariance and stationarity.

What this means for setting \eta and \lambda on normalized weights:

• With a fixed initialization, changing \eta and \lambda at fixed product changes the transient; that the difference is later forgotten is conjectured, not proved (5.2.3).
• Decaying \eta at fixed \lambda shrinks the equilibrium rotation only like \sqrt\eta, and only once the norm has re-equilibrated, which takes \sim1/(4\eta\lambda) steps (SGD) or 1/(2\eta\lambda) (AdamW). If the schedule moves faster, the equilibrium is a target, not the state.
• With little gradient noise the equilibrium can be unstable, and the dynamics turn periodic (5.3, mechanism 1).
• Parameters that are not scale-invariant (gains, biases, embeddings, the output layer) behave in the ordinary way: weight decay can change their function directly.
• A pre-LN transformer is not exactly scale-invariant, but Kosson et al. find the predicted rotation holding very well in a GPT-2 (124M) under AdamW with weight standardization, and “often close” without it.

Why do runs sometimes spike (a sudden rise in training loss that recovers) or slowly diverge (a drift upward that does not)? There is no single cause. Below are four documented mechanisms, each asked whether it is an edge-of-stability phenomenon: an instability that appears when some curvature, measured in the optimizer’s own geometry, crosses its stability threshold (for GD with learning rate \eta and sharpness S, \eta S=2; NB2 §3).

The lead papers (see Readings) are Wortsman et al. for 3 and 4 (they study slow divergences, not fast spikes) and Bai et al. for 2.

Mechanism 1 is the equilibrium of 5.2 going unstable when gradient noise is small. The walkthrough reproduces 1, 2 and the drift of 4 on a laptop in seconds.

Derivations and walkthrough. Setup: GD on \tfrac12hx^2 gives x_{k+1}=(1-\eta h)x_k, stable iff 0<\eta h<2; near a minimum with Hessian H the condition is \eta S<2, S=\lambda_{\max}(H). For a scale-invariant block w=ru (5.1), the direction moves with \eta_{\rm eff}=\eta/((1-\eta\lambda)r^2), is stable near a minimum iff \eta_{\rm eff}S(u)<2 (S(u) the sharpness at the unit-norm point), and the norm obeys r_{k+1}^2=(1-\eta\lambda)^2r_k^2+\eta^2G_k^2/r_k^2, G_k the gradient norm at u_k.

Mechanism 1: weight decay drives a normalized layer to the edge of stability. Near a minimum the gradient is small, so the norm recursion is dominated by decay: r shrinks and \eta_{\rm eff} rises until \eta_{\rm eff}S(u) crosses 2. The direction goes unstable, the loss jumps, the large gradients inflate r, and training re-converges; then the decay starts again. Lobacheva et al. document these periodic destabilizations in long constant-learning-rate BatchNorm training, and Lyu, Li and Arora prove (locally) that GD with weight decay reaches \eta_{\rm eff}\approx2/S(u).

Exercise 5.3.1 — the cycle. How long does weight decay take to push a converged normalized layer to instability, and what ends each spike? Let u^* be a minimum of a scale-invariant L on the sphere, with sharpness S(u^*) in tangent directions, and run full-batch GD with weight decay from near u^*.

1. Using 5.1.2.4, show that the direction is linearly stable at u^* iff r^2>r_c^2=\eta S(u^*)/(2(1-\eta\lambda)).
2. While u sits at u^*, the gradient vanishes. Show that r_k^2 then decays like (1-\eta\lambda)^{2k}, and estimate the number of steps it takes to go from r^2 to r_c^2.
3. Describe the cycle that follows. Which term in the norm recursion ends each spike?
4. In the toy L(x,y)=x^2/(x^2+y^2) (used by Li, Lyu and Arora and by Lobacheva et al.), the minimum is u^*=(0,\pm1). Compute S(u^*) and r_c for \eta=0.1.

Code (17 lines)

def gd_wd_2d(eta, lam, T, noise=0.0, w0=(1.0, 0.5), seed=0):
    # GD + weight decay on L(x, y) = x^2 / (x^2 + y^2); 'noise' adds a tiny isotropic kick per step.
    rng = np.random.default_rng(seed); w = np.array(w0); rec = []
    for _ in range(T):
        x, y = w; r2 = x * x + y * y
        g = np.array([2 * x * y * y, -2 * x * x * y]) / r2**2
        rec.append((x * x / r2, eta * 2.0 / ((1 - eta * lam) * r2), math.sqrt(r2)))  # loss, eta_eff*S(u*), r
        w = (1 - eta * lam) * w - eta * g + noise * rng.standard_normal(2)
    return np.array(rec)

eta, lam, T = 0.1, 0.01, 6000
quiet = gd_wd_2d(eta, lam, T, noise=1e-6)          # noise ~ roundoff: periodic edge-of-stability cycles
fig, ax = plt.subplots(1, 3, figsize=(12, 3.3))
ax[0].semilogy(quiet[:, 0]); ax[0].set(title="loss", xlabel="step")
ax[1].plot(quiet[:, 1]); ax[1].axhline(2, color="k", ls="--"); ax[1].set(title=r"$\eta_{\rm eff}\,S(u^*)$", xlabel="step")
ax[2].plot(quiet[:, 2]); ax[2].axhline(math.sqrt(eta / (1 - eta * lam)), color="k", ls="--")
ax[2].set(title=r"$\|w\|$ (dashed: $r_c$)", xlabel="step"); plt.tight_layout(); plt.show()

The loss axis is logarithmic: each spike jumps many orders of magnitude out of a converged state, then re-converges. Exercise 5.3.1.2 predicts the first crossing of r_c after \log(r_0^2/r_c^2)/(2\eta\lambda)\approx\log(1.25/0.1)/0.002\approx1260 steps (r_0^2=1.25 for w_0=(1,0.5), r_c^2\approx\eta=0.1). The loss starts rising at about step 1340 (the extra delay is the oscillation growing from its tiny seed), and the cycle then repeats every \approx130 steps. The tiny noise stands in for roundoff and minibatch noise: without it the instability never visibly triggers within 6000 steps (the unstable component underflows) and \|w\| shrinks to about 0.003. With much larger noise you are back in 5.2’s smooth equilibrium.

Mechanism 2: Adam’s second moment goes stale. Adam (with gradient g) divides each coordinate’s step \eta\hat m by \sqrt{\hat v}+\epsilon, where \hat m and \hat v are running averages of g and g^2 with decay rates \beta_1,\beta_2 (5.2). Late in training, freeze the preconditioner P=\mathrm{diag}(\sqrt{\hat v}+\epsilon) and define the preconditioned sharpness S_P=\lambda_{\max}(P^{-1}H), H the Hessian. On a quadratic the positive-curvature modes are then stable iff \eta S_P<(2+2\beta_1)/(1-\beta_1), which is 38 at \beta_1=0.9 (NB2 §3.5; Cohen et al. 2022), and full-batch Adam drives \eta S_P up to this threshold and holds it there.

In a quiet phase where some coordinates’ gradients fall well below \sqrt{\hat v}, \hat v decays instead of tracking g^2, so \eta S_P grows past the threshold; the oscillation produces large gradients, \hat v catches up on its 1/(1-\beta_2) timescale, and the loss recovers. That is Bai et al.’s account. Not every crossing spikes: they find the better predictor is the preconditioned curvature along the gradient, \eta\,g^\top P^{-1}Hg/\|g\|^2, crossing the threshold. (Molybog et al. give a related account for 65B–546B-parameter models in which \epsilon, not \hat v’s decay, sets the scale.)

Exercise 5.3.2 — Adam’s quiet-phase overshoot, in one dimension. How fast does a quiet phase push Adam past its threshold, and how large a gradient does it take to recover? Take L=\tfrac12hx^2 and Adam with \beta_1=0 (i.e. RMSprop), \epsilon=0, late in training (ignore bias correction): g_k=hx_k, v_k=\beta_2v_{k-1}+(1-\beta_2)g_k^2, x_{k+1}=x_k-\eta g_k/\sqrt{v_k}. Write q_k=\eta h/\sqrt{v_k}; this is \eta S_P in one dimension, and at \beta_1=0 the threshold (2+2\beta_1)/(1-\beta_1) is 2.

1. With v frozen, show x_{k+1}=(1-q)x_k, stable iff q<2.
2. Quiet phase. Suppose g_k^2\ll v_k for a stretch of steps. Show v_k\approx\beta_2^kv_0, so q grows by a factor F in k_F=2\log F/\log(1/\beta_2)\approx2\log F/(1-\beta_2) steps. Evaluate k_2 for \beta_2=0.999 and 0.9.
3. Overshoot. Suppose q has drifted to 2(1+\delta) with small \delta>0. Show that the oscillation amplitude then grows by a factor \approx1+2\delta per step, and that v must grow by a factor (1+\delta)^2\approx1+2\delta to bring q back to 2. Writing c_k=g_k^2/v_{k-1}, show v_k/v_{k-1}=\beta_2+(1-\beta_2)c_k. Conclude that correcting the overshoot within m steps requires c\gtrsim1+2\delta/\big(m(1-\beta_2)\big) on those steps. Evaluate for \delta=0.1, m=10, with \beta_2=0.999 and \beta_2=0.9.
4. Why do spikes of this kind need a quiet phase, and why are they rare in the noisy early phase of training?

The next cell reproduces Bai et al.’s two-layer contrast: a width-20 tanh network fitting \sin x+\sin4x with full-batch Adam (\eta=0.02, \beta_1=0.9), at \beta_2=0.999 and \beta_2=0.9. Every 20 steps it computes the exact Hessian (61 parameters) and \eta S_P=\eta\,\lambda_{\max}\big(P^{-1/2}HP^{-1/2}\big) (same eigenvalues as P^{-1}H).

Code (33 lines)

x_fit = torch.linspace(-math.pi, math.pi, 64, dtype=torch.float64)
y_fit = torch.sin(x_fit) + torch.sin(4 * x_fit)
WIDTH = 20

def fit_loss(theta):                    # theta = (W1, b1, W2, b2) for a 1 -> 20 -> 1 tanh network
    W1, c1, W2, c2 = torch.split(theta, [WIDTH, WIDTH, WIDTH, 1])
    pred = torch.tanh(x_fit[:, None] * W1 + c1) @ W2 + c2
    return 0.5 * ((pred - y_fit) ** 2).mean()

def adam_full_batch(eta, b2, T, b1=0.9, eps=1e-8, every=20, seed=0):
    gen = torch.Generator().manual_seed(seed)
    theta = torch.cat([torch.randn(WIDTH, generator=gen, dtype=torch.float64), torch.zeros(WIDTH, dtype=torch.float64),
                       torch.randn(WIDTH, generator=gen, dtype=torch.float64) / math.sqrt(WIDTH),
                       torch.zeros(1, dtype=torch.float64)])
    m, v = torch.zeros_like(theta), torch.zeros_like(theta); losses, psharp = [], []
    for k in range(1, T + 1):
        th = theta.clone().requires_grad_(True); L = fit_loss(th); g, = torch.autograd.grad(L, th)
        m = b1 * m + (1 - b1) * g; v = b2 * v + (1 - b2) * g * g
        mh, vh = m / (1 - b1**k), v / (1 - b2**k); losses.append(L.item())
        if k % every == 0:
            H = torch.autograd.functional.hessian(fit_loss, theta, vectorize=True)
            s = (vh.sqrt() + eps).rsqrt()
            psharp.append((k, eta * torch.linalg.eigvalsh(s[:, None] * H * s[None, :])[-1].item()))
        theta = theta - eta * mh / (vh.sqrt() + eps)
    return np.array(losses), np.array(psharp)

fig, ax = plt.subplots(2, 2, figsize=(11, 5), sharex=True)
for j, b2 in enumerate([0.999, 0.9]):
    losses, ps = adam_full_batch(0.02, b2, 5000)
    ax[0, j].semilogy(losses); ax[0, j].set(title=rf"full-batch Adam, $\beta_2={b2}$", ylabel="loss")
    ax[1, j].plot(ps[:, 0], ps[:, 1]); ax[1, j].axhline(38, color="k", ls="--")
    ax[1, j].set(xlabel="step", ylabel=r"$\eta S_P$", ylim=(0, 50))
plt.tight_layout(); plt.show()

With \beta_2=0.999, \eta S_P climbs to 38 and then saws around it (each slow rise is \hat v decaying), and the loss spikes by one to two orders of magnitude, though not at every crossing. With \beta_2=0.9, \eta S_P overshoots once early (to about 45), then sits a little below 38, and the loss never spikes; but it ends near 10^{-3}, against about 2\times10^{-5} for \beta_2=0.999 despite its spikes. Here lowering \beta_2 trades spikes for slower progress. That this is what happens in language models rests on Bai et al.’s larger experiments, not on this cell.

Mechanism 3: attention-logit growth. An attention head with query and key matrices W_Q,W_K has logits z_{ij}=\langle W_Qx_i,W_Kx_j\rangle/\sqrt{d_h} (x_i the inputs, d_h the head dimension). They are bilinear in W_Q,W_K, so they grow quadratically with the parameter scale; when they get large, attention collapses to near one-hot and training diverges. Dehghani et al. saw this in ViTs of about 8B parameters (logits past 5\times10^4). Wortsman et al. reproduce it in small models at high learning rate, find the growth is in \|q\|,\|k\| rather than their alignment, show it persists with a pointwise (non-softmax) attention variant, and remove it with QK-layernorm. Nobody has shown it to be an edge-of-stability crossing. The exercise shows why logits grow with the weights and why a collapsed head stays collapsed, not the divergence itself.

Exercise 5.3.3 — attention logits and QK-layernorm. What sets the size of attention logits, and how does QK-layernorm cap it? Let x_i\in\mathbb R^d be the (normalized) inputs to an attention head, q_i=W_Qx_i, k_j=W_Kx_j\in\mathbb R^{d_h}, logits z_{ij}=\langle q_i,k_j\rangle/\sqrt{d_h} and weights p_{ij}=\mathrm{softmax}_j(z_{ij}).

1. Show z(cW_Q,c'W_K)=cc'\,z(W_Q,W_K) and |z_{ij}|\le\|W_Q\|_2\|W_K\|_2\|x_i\|\|x_j\|/\sqrt{d_h}. If both matrices grow in norm by a factor \rho (as can happen at high learning rate without weight decay), how do the logits scale?
2. The Jacobian of p_{i\cdot} with respect to z_{i\cdot} is \mathrm{diag}(p)-pp^\top. Show that if the largest attention probability in the row, say p_{i1}, is 1-\delta, every entry of this Jacobian is O(\delta). What does that do to the gradient signal reaching the losing keys?
3. QK-layernorm replaces q and k by \hat q=\mathrm{LN}(q), \hat k=\mathrm{LN}(k) (no gain). Show |z_{ij}|\le\sqrt{d_h}. With elementwise gains \gamma_q,\gamma_k, what is the bound? What does QK-layernorm do to W_Q and W_K in the sense of 5.1?

Mechanism 4: output-logit divergence. For logits z\in\mathbb R^K, write \log Z=\log\sum_ie^{z_i}; the cross-entropy for label y is \log Z-z_y. It is invariant under shifting all logits by a constant, so the overall level of the logits is a flat direction of the loss. Wortsman et al. observe, in models trained without weight decay, that late in training the output logits drift to very negative values and the loss diverges. z-loss (from PaLM), an auxiliary term 10^{-4}\log^2Z, removes it, and weight decay also helps at larger scales. They do not say why the drift eventually breaks training; the exercise shows why it happens at all under Adam. The drift runs along a direction of zero curvature, so it is not an edge-of-stability crossing; whether what finally diverges is one is not established.

Exercise 5.3.4 — a flat direction that Adam does not respect. Why does Adam push the output logits along a direction the loss cannot see, when SGD does not? Take a final layer z=Wh+b, W\in\mathbb R^{K\times n}, on features h\in\mathbb R^n, with softmax cross-entropy: p=\mathrm{softmax}(z), \log Z=\log\sum_ie^{z_i}, and loss \mathrm{CE}=\log Z-z_y for label y, whose gradient is \partial\,\mathrm{CE}/\partial z=p-e_y (e_y the one-hot vector of y). Write \bar z=\frac1K\sum_iz_i for the mean logit.

1. Show that the cross-entropy is invariant under b\to b+c\mathbf 1 and under W\to W+\mathbf 1v^\top for any c,v. Show that the gradient is orthogonal to these directions, so under SGD (any batch size, no weight decay) the column sums \mathbf 1^\top W and the sum \mathbf 1^\top b are exactly conserved.
2. Now take sign descent (Adam with \beta_1=\beta_2=\epsilon=0) with a batch of one example with label y, and assume 0<p_i<1. Show that the logits of that example (holding h fixed) change by \Delta z_i=-\eta\,s_i(\|h\|_1+1), with s_i=+1 for i\ne y and s_y=-1. Deduce that the mean logit falls by \eta(\|h\|_1+1)(K-2)/K per step, and that to first order \Delta\log Z=-\eta(\|h\|_1+1)(1-2p_y). When does \log Z fall?
3. Show that z-loss, \alpha_z(\log Z)^2 with coefficient \alpha_z>0, has gradient 2\alpha_z\log Z\,p with respect to z, whose component along \mathbf 1 is 2\alpha_z\log Z. What does it pin, and what does \log Z=0 mean for the logits?

Code (37 lines)

# A small classifier on overlapping Gaussian classes (loss stays ~2.5 nats, like an LM).
gen = torch.Generator().manual_seed(0)
K_cls, d_in, N, hid = 32, 16, 8192, 32
means = torch.randn(K_cls, d_in, generator=gen)
labels = torch.randint(0, K_cls, (N,), generator=gen)
inputs = means[labels] + 3 * torch.randn(N, d_in, generator=gen)

def train_classifier(opt_name, lr, z_coef, steps=2000, batch=64):
    torch.manual_seed(1)
    net = torch.nn.Sequential(torch.nn.Linear(d_in, hid), torch.nn.LayerNorm(hid), torch.nn.ReLU(),
                              torch.nn.Linear(hid, K_cls))
    with torch.no_grad():                      # start near uniform predictions: log Z ~ 0 at init
        net[-1].weight.mul_(0.1); net[-1].bias.fill_(-math.log(K_cls))
    opt = (torch.optim.Adam(net.parameters(), lr=lr, betas=(0.9, 0.95)) if opt_name == "adam"
           else torch.optim.SGD(net.parameters(), lr=lr))
    hist = []
    for t in range(steps):
        idx = torch.randint(0, N, (batch,))
        z = net(inputs[idx]); logZ = torch.logsumexp(z, 1)
        ce = (logZ - z.gather(1, labels[idx, None])[:, 0]).mean()
        opt.zero_grad(); (ce + z_coef * (logZ ** 2).mean()).backward(); opt.step()
        hist.append((ce.item(), logZ.mean().item(), z.mean().item()))
    return np.array(hist)

runs = {"SGD (lr 0.3)": ("sgd", 0.3, 0.0), "Adam (lr 1e-2)": ("adam", 1e-2, 0.0),
        "Adam + z-loss 1e-2": ("adam", 1e-2, 1e-2)}
fig, ax = plt.subplots(1, 3, figsize=(12, 3.3))
for name, args in runs.items():
    h = train_classifier(*args)
    smooth = lambda s: np.convolve(s, np.ones(50) / 50, mode="valid")
    for j in range(3):
        ax[j].plot(smooth(h[:, j]), label=name)
    print(f"{name:20s} final CE {h[-200:, 0].mean():.3f}   log Z {h[-200:, 1].mean():6.2f}   mean logit {h[-200:, 2].mean():6.2f}")
for j, t in enumerate(["cross-entropy", r"batch mean of $\log Z$", "mean logit"]):
    ax[j].set(title=t, xlabel="step")
ax[1].axhline(0, color="k", ls="--", lw=1)
ax[0].legend(); plt.tight_layout(); plt.show()

SGD (lr 0.3)         final CE 2.659   log Z   0.94   mean logit  -3.46
Adam (lr 1e-2)       final CE 2.659   log Z  -2.93   mean logit  -7.40
Adam + z-loss 1e-2   final CE 2.660   log Z  -0.54   mean logit  -5.00

All three runs start with \log Z\approx0 (uniform predictions). Under SGD the final layer’s column sums and bias sum are exactly conserved (5.3.4.1); the mean logit could still move through the features h, but here it stays put, and \log Z rises only because the logits spread out. Adam reaches the same cross-entropy while steadily dragging the mean logit down and \log Z well below zero. With z-loss, \log Z stops near zero (about -0.5, where its pull balances Adam’s drift). The z-loss coefficient is 10^{-2} so the effect shows within 2000 steps; PaLM and Wortsman et al. use 10^{-4} over 10^5 steps. This cell shows the drift and the pinning, not the divergence: nothing breaks here in float32.

Is gradient clipping (rescaling the gradient when its norm exceeds a threshold) only a guard against rare huge gradients? Not only: when the local curvature grows with the gradient norm, the right step size is \propto1/\|\nabla L\| where the gradient is large, which is what clipping does, and clipped GD can then be arbitrarily faster than GD with any fixed step. Under Adam the step is already bounded, and clipping mainly protects the optimizer state.

Paper: Zhang, He, Sra & Jadbabaie, Why gradient clipping accelerates training (see Readings).

The core idea. L is (L_0,L_1)-smooth if \|\nabla^2L(x)\|\le L_0+L_1\|\nabla L(x)\| (operator norm on the Hessian). This contains ordinary L-smoothness (L_1=0) and also, e.g., every univariate polynomial of degree at least 3. A fixed step must be safe at the worst point of the run and is then too small everywhere else. Zhang et al. see curvature tracking the gradient norm clearly in LSTM language models, but in ResNet-20 on CIFAR-10 only with clipping and a large learning rate; whether clipping speeds up a given run depends on that.

Derivations and walkthrough. Setup: Clipped GD is x_{k+1}=x_k-h_k\nabla L(x_k),\qquad h_k=\min\Big(\eta,\ \frac{\gamma\eta}{\|\nabla L(x_k)\|}\Big): step \eta where the gradient is small, a step of fixed length \gamma\eta where it is large. Exercise 5.4.1 shows on a two-line example why it can beat every fixed step; Exercise 5.4.2 derives the descent lemma behind the general statement.

Exercise 5.4.1 — a two-line example. How much faster can clipped GD be than GD with the best fixed step? Take L(x,y)=\cosh x+\tfrac14y^2, started at (x_0,1) with x_0>0 large. The x direction has curvature growing with the gradient; the y direction is an ordinary quadratic with curvature \tfrac12.

1. Show that L is (1,1)-smooth, i.e. \|\nabla^2L\|\le1+\|\nabla L\|, using \cosh x\le1+|\sinh x|.
2. Show that fixed-step GD with h\sinh x_0>2x_0 overshoots to |x_1|>x_0 and then diverges. Conclude that every fixed step that converges has h<2x_0/\sinh x_0\approx4x_0e^{-x_0}, and that the y coordinate then needs at least \approx2\log(1/\epsilon)/h\approx\sinh(x_0)\log(1/\epsilon)/x_0 steps to reach |y|<\epsilon.
3. Clipped GD on the global gradient, with \eta=1, \gamma=1 (step length at most 1): roughly how many steps?
4. Why is the second coordinate needed? (What happens in one dimension with h=x_0/\sinh x_0?)

Then check it numerically: implement clipped GD, and compare it with fixed-step GD over a grid of step sizes, on the number of steps to reach \|(x,y)\|<10^{-8}.

def clipped_gd_steps(x0, eta=1.0, gamma=1.0, tol=1e-8):
    # TODO: clipped GD on L = cosh(x) + y^2/4 from (x0, 1); return the number of steps to |(x, y)| < tol.
    ...

def fixed_gd_steps(x0, hs, tol=1e-8):
    # TODO: fixed-step GD for each h in the grid hs (vectorize over hs); return the fewest steps and its h.
    ...

# for x0 in [1.7, 3.7, 5.7, 7.7, 9.7, 11.7]: print both, and the lower bound from Exercise 5.4.1.2

def clipped_gd_steps(x0, eta=1.0, gamma=1.0, tol=1e-8):
    x, y = float(x0), 1.0
    for k in range(10**6):
        if math.hypot(x, y) < tol:
            return k
        gx, gy = math.sinh(x), y / 2
        h = min(eta, gamma * eta / math.hypot(gx, gy))
        x, y = x - h * gx, y - h * gy

def fixed_gd_steps(x0, hs, tol=1e-8, max_steps=3 * 10**5):
    x, y = np.full(len(hs), float(x0)), np.ones(len(hs))
    for k in range(max_steps):
        ok = np.hypot(x, y) < tol
        if ok.any():                          # the first step size to converge is the best on the grid
            return k, hs[np.argmax(ok)]
        with np.errstate(over="ignore", invalid="ignore"):
            x = np.where(np.abs(x) < 50, x - hs * np.sinh(np.clip(x, -50, 50)), np.inf)
        y = y - hs * y / 2
    return None, None

print("  x0 | clipped GD | best fixed step on grid (h)  | lower bound for any fixed step")
for x0 in [1.7, 3.7, 5.7, 7.7, 9.7, 11.7]:
    n, h = fixed_gd_steps(x0, np.geomspace(1e-5, 1.9, 400))
    print(f"{x0:5.1f}| {clipped_gd_steps(x0):6d}     | {n:8d}   (h = {h:.2e})     | {math.log(1e8) / -math.log(1 - x0 / math.sinh(x0)):10.0f}")

  x0 | clipped GD | best fixed step on grid (h)  | lower bound for any fixed step
  1.7|     28     |       20   (h = 1.24e+00)     |         18
  3.7|     30     |       95   (h = 3.56e-01)     |         91
  5.7|     32     |      481   (h = 7.52e-02)     |        474
  7.7|     34     |     2688   (h = 1.37e-02)     |       2632
  9.7|     36     |    15774   (h = 2.33e-03)     |      15485
 11.7|     38     |    95214   (h = 3.87e-04)     |      94906

The best fixed step on the grid sits just below 2x_0/\sinh x_0, and its step count sits just above the lower bound, which holds for every fixed step; clipped GD’s grows by about one step per unit of x_0 (\approx x_0+26), so its advantage is clear by x_0\approx3 (at x_0=1.7 the tuned fixed step still wins, 20 steps to 28).

Exercise 5.4.2 — the descent lemma and the optimal step. What step size does (L_0,L_1)-smoothness certify as safe, and how does it compare with clipping? Assume \|\nabla^2L\|\le L_0+L_1\|\nabla L\| everywhere (operator norm on the Hessian). Fix x, write g=\nabla L(x), and let \|d\|\le1/L_1.

1. Let \phi(s)=\|\nabla L(x+sd)\|. Show \phi'(s)\le(L_0+L_1\phi(s))\|d\| (where differentiable), and conclude by Grönwall that L_0+L_1\phi(s)\le(L_0+L_1\|g\|)e^{L_1s\|d\|}.
2. Deduce \|\nabla L(x+sd)-g\|\le(L_0+L_1\|g\|)\big(e^{L_1s\|d\|}-1\big)/L_1, and then L(x+d)\le L(x)+\langle g,d\rangle+(e-2)\,(L_0+L_1\|g\|)\,\|d\|^2 .
3. Take d=-hg. Find the h minimizing the bound, check it satisfies \|d\|\le1/L_1, and show it equals clipped GD (step \min(\eta,\gamma\eta/\|g\|) for suitable \eta,\gamma) up to a factor of 2. How much does each step decrease L, for small and for large \|g\|?
4. Fixed-step GD uses a single h for the whole run. For which h does the descent lemma certify descent at every iterate, and what does that suggest about its speed?

Clipping under Adam. In large-model training, clipping is usually applied to the global gradient norm before the gradient enters Adam (Wortsman et al. clip at norm 1), and the GD theory above does not transfer directly. With Adam’s coordinatewise v_k=\beta_2v_{k-1}+(1-\beta_2)g_k^2, bias correction \hat v_k=v_k/(1-\beta_2^k), \beta_1=0 (so the step is \eta g_k/(\sqrt{\hat v_k}+\epsilon)) and any \epsilon\ge0, we have \hat v_k\ge v_k\ge(1-\beta_2)g_k^2, so |\Delta w_i|=\frac{\eta\,|g_{k,i}|}{\sqrt{\hat v_{k,i}}+\epsilon}\le\frac{\eta}{\sqrt{1-\beta_2}}. So clipping is not needed to keep Adam’s step finite. It still changes the current step (with \beta_2=0.999 and remembered v=1, a gradient of 1000 moves a coordinate by about 31.6\eta; clipped to 1, by \eta), and it protects the optimizer state: an unclipped huge gradient inflates \hat v, which then suppresses those coordinates’ updates until it decays back, which takes \approx\log F/(1-\beta_2) steps to fall by a factor F late in training if later gradients are much smaller. This is a plausible role, not an established theory; we know of no analysis of clipped Adam comparable to the GD story above.