Train and stabilize a 20-layer neural network

This is the vanishing-gradients chapter's exact failure mode, just twenty layers deep instead of a handful: with weights initialized too small, the gradient reaching the input has almost nothing left of it. Flip on RMSNorm and watch every bar past the first few layers jump from invisible to a healthy, readable size -- the same rescue the RMSNorm chapter demonstrated on a single vector, now stacked twenty times in a row.

Four independent problems, four independent fixes -- each one exactly the function its own chapter built:

• $h_l$ -- layer $l$'s output; RMSNorm (previous chapter) renormalizes it after every layer, so its scale can't compound across depth.
• $g$, $s$ -- a raw gradient and the loss-scale factor; dynamic loss scaling (mixed-precision-training) multiplies by $s$ before casting to fp16, then divides back out, so a small $g$ survives storage instead of underflowing to exactly $0$.
• $\eta(t)$, $t_{\text{warmup}}$ -- the learning rate at step $t$ and the warmup length; LR warmup (learning-rate-schedules) ramps $\eta$ up from $0$ instead of applying the full rate to the very first, least-trustworthy gradient.
• $w$, $\lambda$ -- a weight and the decay strength; decoupled weight decay (weight-decay-vs-l2) shrinks $w$ by a fixed fraction every step, independent of any optimizer's per-parameter gradient history.

1. These four failure modes don't cure each other

   Normalizing every layer fixes vanishing magnitude, but does nothing about:

   • A raw fp16 cast underflowing to zero.
   • An unbounded first update.
   • Weights drifting upward over training.

   Each fix repairs exactly one mechanism.

2. Some of them only matter once another is in place

   Loss scaling rescues a vanishing gradient from becoming exactly $0$ in storage -- but a gradient that's merely tiny, not zero, is still too small to be useful. RMSNorm is what actually lifts the magnitude; loss scaling only stops that magnitude from being destroyed a second time on the way into fp16.

3. A healthy network needs all of its own separate fixes

   Below, toggling on only one or two of the four almost never clears every failure mode at once -- the demo's health check is a logical AND across gradient magnitude, first-update size, and weight growth, exactly like a real training run failing for any one of several unrelated reasons.

step-0 update ≈ 0.000 (unsafe above 0.5) · weight after 50 steps ≈ 11.47 (bound 5)

✗ still unstable

With every fix off, the gradient reaching the input is stuck at essentially the same rounding-floor magnitude as the earlier vanishing-gradients chapter -- and the fp16 cast crushes it to exactly $0$ on top of that. Turn on RMSNorm and the gradient jumps to a healthy scale, but the very first update (still at the full base learning rate) blows past the safe threshold. Add warmup and that first step is safe -- but fifty steps in, an unchecked weight has drifted well past its bound. Only with all three of normalize, warmup, and weight decay on together does the readout turn green.

1. Start from the broken network

   With every fix off, the gradient reaching the input is about $1.05\times10^{-7}$ in full precision -- already far below the $0.3$ healthy floor. Cast to the toy fp16 format with no loss scaling, it stores as exactly $0$: the update for that layer's weights is now not just tiny, but identically zero.

2. Loss scaling alone doesn't fix the magnitude

   Turn on loss scaling by itself: that same gradient now survives storage as $\approx1.05\times10^{-7}$ instead of $0$ -- rescued from total information loss, but still nowhere near the $0.3$ floor. RMSNorm, not loss scaling, is the fix for magnitude.

3. RMSNorm fixes the magnitude, but exposes the next problem

   Turn on RMSNorm instead: the gradient reaching the input jumps to $\approx0.738$, comfortably healthy. But at the base learning rate ($1.2$, unwarmed), step $0$'s update is $1.2\times0.738\approx0.886$ -- past the $0.5$ unsafe threshold before training has learned anything.

4. Warmup and decoupled decay close the other two gaps

   Warmup starts the learning rate at exactly $0$, so step $0$'s update is $0$ regardless of gradient size -- safe by construction.

   For the weight bound, each step first grows the weight by $5\%$, then -- with decay on -- shrinks that grown value back down by another $5\%$: a net per-step factor of $1.05\times0.95=0.9975$.

   • Without decay, 50 steps compounds the raw $5\%$ growth alone: $1\times1.05^{50}\approx11.47$
   • With decay, 50 steps compounds the net $0.9975$ factor instead: $1\times0.9975^{50}\approx0.88$, comfortably under the bound of $5$