Weight-norm Criticality: A Mechanism for Loss Spikes Induced by the Normalization and Weight Decay

Training instability, often manifesting as sudden spikes in loss, is typically attributed to the Edge of Stability—a phenomenon where the learning rate pushes the optimization process into regions of the loss landscape that are too sharp for the current step size. This paper introduces a complementary mechanism called weight-norm criticality. The authors demonstrate that in networks containing scale-invariant components (such as BatchNorm or LayerNorm), the interaction between these layers and weight decay creates a specific type of instability. As weight decay persistently shrinks the norms of scale-invariant weights, the local curvature of the loss landscape increases significantly, eventually destabilizing the optimization process.

The researchers show that for scale-invariant components, the loss function is invariant to the scaling of weights, but the Hessian (which measures curvature) is not. Specifically, they prove that as the norm of these weights approaches zero, the maximum eigenvalue of the Hessian—a proxy for landscape sharpness—grows at a rate proportional to the inverse square of the weight norm. This creates a feedback loop: stronger weight decay drives weights toward smaller norms, which in turn makes the landscape sharper, leading to the abrupt loss spikes observed in practice.

Through experiments on Transformers, ResNet-50, and fully connected networks on MNIST, the authors confirm that loss spikes are not merely a result of normalization or weight decay in isolation, but a consequence of their interaction. By visualizing training trajectories using PCA, they show that models under higher weight decay converge toward regions of the parameter space where scale-invariant weights are near zero and curvature is extreme. This framework allows researchers to localize instability to specific layers, providing a clear rationale for why weight decay cannot be increased indefinitely without compromising training stability.