Why gradients rapidly increase near the end of training: weight decay meets LR decay, says Defazio

aaron_defazio · x · 2026-11-09

Aaron Defazio's short paper (arXiv:2506.02285) explains why gradient norms spike near the end of long LLM training runs: an unintended interaction between weight decay, normalization layers, and the learning rate schedule.

Gradient updates push the weight norm up while weight decay pulls it down; the two reach equilibrium where the norm stays roughly constant.
As the LR schedule decays near the end of training, weight decay becomes relatively stronger, shrinking the weight norm.
In scale-invariant networks f(cw)=f(w), gradients are inversely proportional to the weights (from wᵀ∇L=0 it follows that ∇L(cw)=(1/c)∇L(w)), so smaller weights at the end of training produce larger gradients.
The fix is simple: proportionally decrease weight decay as LR shrinks (the AdamC approach; a cited derivation suggests scaling per-step weight decay by the square of (current lr / peak lr)). This correction also yields lower loss throughout training.

Original post →

Why gradients rapidly increase near the end of training: weight decay meets LR decay, says Defazio

More from Research