Why Gradients Rapidly Increase Near the End of Training

TL;DR

Long-duration LLM (and ImageNet/ResNet) training runs exhibit an unexplained rapid increase in gradient norm near the end of training. Defazio shows this is a clean three-way interaction between weight decay, normalization layers (LayerNorm/BatchNorm), and the learning-rate schedule: for any layer immediately followed by a normalization, weight decay drives the gradient-to-weight ratio to a steady state $\eta\lambda$ , so when $\eta \to 0$ on the decay tail, that ratio (and hence the gradient norm) must blow up. The proposed fix — AdamC / SGDC — replaces the weight-decay term $\lambda$ with $\lambda \cdot \eta_{\max} / \eta_t$ , exactly cancelling the schedule-induced inflation; on 200B-token Llama-3-style 120M pretraining and ResNet-50/ImageNet it removes the late-training gradient blow-up and also lowers loss/improves accuracy.

Key claims

For any weight matrix followed by a normalization layer, gradients are orthogonal to weights (a known property of normalized layers), giving an exact closed-form recursion for $\|w\|^2$ under SGD with weight decay [§3, Eq. 1].
Weight decay drives the per-layer gradient-to-weight ratio $\|g\|/\|w\|$ to a steady state of $\sqrt{\eta \lambda}$ (SGD) or $\eta \lambda$ (AdamW), and empirically all normalized layers in ResNet-50 converge to this ratio early in training and track it for the remainder of the run [§3, Fig. 2].
Because the steady-state ratio depends on the current $\eta_t$ , a cosine learning-rate schedule with $\eta_t \to 0$ forces the steady state $\to \infty$ at the end of training; this exactly matches the observed “tail blow-up” of the gradient norm in both SGD-momentum ResNet runs and AdamW LLM pretraining [§3.1, Fig. 3].
Weight decay’s role under AdamW is reinterpreted as balancing the per-layer $\|w\|_\infty$ to a uniform value across layers (via a Hölder-inequality upper bound on the AdamW-weighted norm), giving a new explanation for why Adam without decoupled weight decay underperforms AdamW: it lacks this balancing property [§4.1–4.2].
Decoupled weight decay alone (the AdamW form) is not enough to prevent the tail blow-up; the steady-state ratio still depends on $\eta$ [§5].
The proposed correction is to replace the weight-decay term with $\lambda \cdot \eta_{\max} / \eta_t$ on all normalized layers (called SGDC for SGD, AdamC for AdamW); this exactly cancels the $\eta_t$ dependence in the steady state, yielding $\lim \|g\|/\|w\| = \sqrt{\eta_{\max} \lambda}$ (SGD) or $\eta_{\max} \lambda$ (AdamW) — schedule-independent [§5, Algorithm 1].
On a 120M Llama-3-style model trained 200B tokens with batch 256k on FineWeb-Edu, AdamC vs. AdamW (each LR-swept on a power-of-2 grid): the loss-curve shape shifts substantially with lower final loss, the gradient-norm increase in the second half of training is largely removed, and weight norms become much more stable [§6, Fig. 4].
On ResNet-50 / ImageNet with SGD-momentum and cosine LR (5 seeds, swept LR and WD): SGDC eliminates the rapid end-of-training gradient-norm spike, though a slower linear upward trend over training remains unaffected [§6, Fig. 5].
The training-noise regime matters for whether this is observable: classical MNIST/CIFAR setups (low label noise, multi-epoch, gradients go to zero near a minimum) do not exhibit the tail blow-up; modern single-pass LLM training does. Many optimizer benchmarks are simply too short (10B tokens) to surface this [§2, §6].

Method

The analysis is built on a single property of normalized layers — the gradient of a pre-norm linear layer is orthogonal to its weights, so $\langle g_t, w_t \rangle = 0$ . Plugging this into the SGD-with-weight-decay update gives a clean recursion for the squared weight norm; solving for the fixed point (and dropping a second-order $\eta^2$ term) yields the steady-state $\|g\|/\|w\| = \sqrt{\eta \lambda}$ . The same derivation under AdamW — treating the Adam step as multiplication by a diagonal preconditioner $D_t$ and measuring the change in $\|w\|_{D_t^{-1}}^2$ — gives the AdamW steady state $\|g\|/\|w\| \approx \eta \lambda$ , but only because AdamW decouples decay from $D_t$ ; the original Adam (where decay is also normalized by $D_t$ ) breaks the cancellation, which the paper proposes as a fresh explanation for the AdamW>Adam gap.

The correction is a one-line patch: replace $\lambda \to \lambda \cdot \eta_{\max} / \eta_t$ in the weight-decay term on every layer followed by a normalization (in Transformers, the paper applies it to every linear layer except the LM head). The $\eta_{\max}$ factor is purely to keep the magnitude consistent with the uncorrected case. Experiments use the standard cosine-with-warmup schedule and sweep LR/WD on power-of-2 grids for both methods to remove a “you just got lucky with the LR” confound.

Results

120M Llama-3-style LLM, 200B tokens, FineWeb-Edu, AdamC vs AdamW (independent LR sweeps): noticeably lower loss values throughout training, near-elimination of the end-of-training gradient-norm rise, and substantially more stable weight norms (AdamW’s weight norms decrease rapidly; AdamC’s stay flat) [§6, Fig. 4].
ResNet-50/ImageNet, 5 seeds, SGDC vs SGDM with cosine LR (independent LR/WD sweeps): tail-end gradient-norm spike eliminated; the slow linear gradient-norm rise across all of training is not removed (the paper notes this as residual and open) [§6, Fig. 5].
The size of the correction grows with how aggressively the schedule decays $\eta_t$ : under cosine with min-LR = $\eta_{\max}/10$ , the late-training $\lambda \cdot \eta_{\max}/\eta_t$ is up to 10× the nominal $\lambda$ — i.e. AdamC effectively increases weight-decay strength on the tail rather than letting the gradient norm increase to compensate.

Why it’s interesting

This sits directly in the Training stability at scale / Hyperparameter scaling laws arc the wiki has been tracking, but locates the fix in a fourth place — the weight-decay term itself, conditional on the schedule. How to Set the Learning Rate for Large-Scale Pre-training? argues QK-Norm subsumes µP’s role at modern scale; Controlled LLM Training on Spectral Sphere argues the optimizer should retract weights and updates to a spectral sphere; this paper argues a much smaller surgery — multiply $\lambda$ by $\eta_{\max}/\eta_t$ on normalized layers — gives schedule-independent gradient-to-weight ratios and recovers stable weight-norm dynamics with one extra scalar per step. The mechanism (gradient-weight orthogonality on normalized layers ⇒ steady-state $\|g\|/\|w\| = f(\eta\lambda)$ ) is also exactly the rotational-equilibrium picture How to Set the Learning Rate for Large-Scale Pre-training? §6.2 uses to justify why module-specific LRs are unnecessary at scale — so AdamC is the natural “fix the residual schedule artifact” companion to that line.

It is also a candidate cause for the slow linear gradient-norm drift mentioned (but unexplained) in several wiki entries on long-horizon training. AdamC handles the schedule-induced spike but leaves that drift untouched, which is a clean open question Defazio himself flags.