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

Summary (Overview)

• Novel mechanism for loss spikes: The paper identifies a previously overlooked "weight-norm criticality" — distinct from the well-known learning-rate criticality (Edge of Stability) — that explains training loss spikes in deep neural networks containing normalization layers (e.g., BatchNorm, LayerNorm) combined with weight decay.

• Core insight: Normalization introduces scale-invariant components whose norms are progressively shrunk by weight decay. As these norms approach zero, the Hessian eigenvalues amplify on the order of $\alpha^{-2}$, rapidly increasing local sharpness and destabilizing optimization.

• Theoretical contributions: The authors prove a curvature explosion theorem (Theorem 5.1) showing $\lambda_{\max}(H(\alpha u, v)) \geq \alpha^{-2} \lambda_{\max}(H_{uu}(u, v))$, and derive explicit weight-norm stability boundaries ($c^*$ and $c^*_{\text{spike}}$) that predict when loss spikes will occur.

• Empirical validation: The theory is validated across multiple settings: a 187M-parameter LLaMA-style Transformer on 100B tokens, ResNet-50 on CIFAR-100, FNNs on MNIST, and controlled synthetic regression tasks. The predicted unstable regimes align with observed loss spikes.

• Practical implications: The framework enables layer-wise localization of instability sources, explaining why weight decay improves generalization but cannot be made arbitrarily strong. Disabling weight decay on MLP modules in Transformers substantially reduces loss spikes.

Introduction and Theoretical Foundation

Background and Motivation

Loss spikes — sudden increases in training loss — are a common challenge in deep learning. The paper categorizes existing explanations into three factors:

1. Data: Heterogeneous mini-batches yielding erratic gradient directions (Chowdhery et al., 2023)
2. Landscape: Lower-loss-as-sharper structures making learning rates cross stability thresholds (Li et al., 2025)
3. Optimization: Adam over-amplifying adaptive learning rates when second-momentum decreases trail gradient reductions (Bai et al., 2025)

Most existing work interprets loss spikes through the Edge of Stability (EoS) framework (Cohen et al., 2021), which emphasizes learning-rate criticality. This paper argues for a complementary, often overlooked mechanism: weight-norm criticality.

Key Theoretical Foundation

The central observation is that normalization layers (BN, LN) introduce positive scale invariance in the weights preceding them:

$$L(\alpha u, v) = L(u, v), \quad \forall \alpha > 0$$

where $u$ represents scale-invariant parameters and $v$ represents other parameters. This property means the loss is insensitive to rescaling certain weights, but the local curvature is not — creating a dangerous decoupling.

Related Work Context

The paper builds on prior work showing:

• Van Laarhoven (2017), Hofer et al. (2018), Zhang et al. (2019): interaction between normalization and weight decay influences effective learning rate
• Li & Arora (2019): globally scale-invariant functions are ill-conditioned under canonical optimization
• Li et al. (2020a): improper weight decay can drive certain layer weights to vanish, causing training failure
• Prior analyses focused on globally scale-invariant models; this work extends to architectures mixing scale-invariant and non-scale-invariant parameters, which better reflects practical networks

Methodology

Experimental Settings

The paper employs a systematic progression of experiments:

1. Large-scale LLM pretraining: LLaMA-style 16-layer, 16-head 187M Transformer on 100B-token corpus, sweeping weight decay over {0, 0.5, 1}

2. ResNet-50 on CIFAR-100: Standard architecture with SGD, varying only weight decay

3. Mechanistic probe on MNIST: Fully connected networks with hidden blocks of form:

   $$\text{Linear} \to \text{Norm} \to \text{Tanh}$$

   where Norm ∈ {BN, LN, Identity}, sweeping weight decay

4. Fully controlled synthetic regression: Three-layer FNN with BN mapping $x = [x_1, x_2] \in \mathbb{R}^2$ to $y = x_1 + 2x_2$

PCA Trajectory Visualization

To visualize optimization dynamics, the paper uses PCA on flattened parameter vectors across all epochs and weight decay settings:

$$\boldsymbol{\Theta} = [\boldsymbol{\theta}_1, \boldsymbol{\theta}_2, \ldots, \boldsymbol{\theta}_n]$$

The first two principal components define a shared projection plane, with loss contours computed via inverse PCA transform.

Module-wise Hessian Eigenvector Analysis

For Transformers, the top Hessian eigenvector is tracked during training and decomposed across parameter modules by computing the fraction of squared norm supported on each module.

Empirical Validation / Results

Key Findings

1. Weight decay induces loss spikes only with normalization: In Fig. 2, FNNs without normalization show no spikes under weight decay sweeps, while BN and LN variants exhibit pronounced spikes.

2. Convergence toward singularities: PCA analysis (Figs. 3, 4b) shows that increasing weight decay systematically shifts solutions toward a "zero-limit reference point" where scale-invariant parameters vanish. These reference points lie in regions of extremely large curvature.

3. Curvature explosion verified: Fig. 5 empirically validates Theorem 5.1 — rescaling scale-invariant parameters toward zero causes $\lambda_{\max}$ to increase dramatically, matching the predicted $\alpha^{-2}$ scaling.

4. Scale-invariant block dominates sharpness: Fig. 6 shows that during training, the maximum eigenvalue of the scale-invariant Hessian block rapidly surpasses the non-scale-invariant block and becomes the dominant contribution to global $\lambda_{\max}$.

5. Weight-norm criticality predicts spikes: Fig. 7 shows that when the weight norm $\|W_1\|$ of the first scale-invariant layer persistently falls below the spike boundary $c^*_{\text{spike}}$, pronounced loss spikes occur. The onset and return to stability align with the predicted boundaries.

6. Module-wise localization in Transformers: Fig. 8 shows MLP blocks dominate the leading-curvature direction, with norm contributions increasing during training. Disabling weight decay on MLP parameters (Fig. 9) substantially reduces loss spikes and achieves lower training loss.

Theoretical Results

Theorem 5.1 (Curvature Explosion Induced by Scale Invariance): For a twice continuously differentiable loss $L(u, v)$ positively scale-invariant in $u$:

$$\lambda_{\max}(H(\alpha u, v)) \geq \alpha^{-2} \lambda_{\max}(H_{uu}(u, v))$$

Proposition 5.2 (Weight-Norm Stability Boundary): With intrinsic curvature $\rho := \|u\|^2 \cdot \lambda_{\max}(H_{uu}(u, v))$, the stability boundary is:

$$c^* = \sqrt{\frac{\eta \rho}{2}}$$

Proposition 5.3 (Weight-Norm Spike Boundary): Using gradient-direction curvature:

$$\lambda_{\text{grad}}(H_{uu}(u, v)) := \frac{g_u^\top H_{uu}(u, v) g_u}{\|g_u\|^2}$$

and intrinsic spike curvature $\rho_{\text{grad}} := \|u\|^2 \cdot \lambda_{\text{grad}}(H_{uu}(u, v))$, the spike boundary is:

$$c^*_{\text{spike}} := \sqrt{\frac{\eta \rho_{\text{grad}}}{2}}$$

The derivation follows from combining the EoS stability criterion $\eta \lambda_{\max}(H) \leq 2$ with the scaling property $\lambda_{\max}(H(w, v)) \geq \rho/\|w\|^2$.

Theoretical and Practical Implications

Theoretical Significance

• Extends beyond global scale invariance: Prior analyses focused on globally scale-invariant models; this framework applies to practical architectures containing scale-invariant components alongside non-scale-invariant ones (e.g., residual connections in ResNet).

• Layer-wise decomposition: The critical boundary decomposes across individual scale-invariant components, requiring only the Hessian restricted to that component's parameters. This enables attributing instability to specific layers.

• Complementary to EoS: Weight-norm criticality provides a complementary perspective to learning-rate criticality, collectively explaining loss spike emergence.

Practical Implications

1. Explains the weight decay tension: Weight penalties improve generalization but cannot be made arbitrarily strong — excessive decay pushes scale-invariant weight norms past the critical boundary, destabilizing training.

2. Monitoring tool: The layer-wise thresholds provide monitoring value — when loss spikes arise from weight-norm collapse, instability can be attributed to specific layers.

3. Actionable training strategies: The module-wise analysis suggests practical interventions, e.g., disabling weight decay on MLP modules in Transformers to reduce instability (Fig. 9).

4. Architectural design awareness: The work shows that architectural design choices — alongside data and optimization — can decisively shape optimization dynamics.

Conclusion

This paper identifies weight-norm criticality as a novel mechanism for loss spikes arising from the interaction between normalization and weight decay. The key contributions are:

1. Mechanistic understanding: Weight decay shrinks scale-invariant parameter norms, which amplifies curvature on the order of $\alpha^{-2}$, driving the loss landscape toward sharp regions and destabilizing training.

2. Quantitative boundaries: Explicit weight-norm stability ($c^*$) and spike ($c^*_{\text{spike}}$) boundaries that predict when training becomes unstable.

3. Practical applicability: The framework applies to realistic architectures (Transformers, ResNets) with mixed scale-invariant and non-scale-invariant components, enabling layer-wise localization of instability.

Future Directions

The authors acknowledge open questions:

• The mechanisms underlying loss spikes in general settings remain incompletely understood
• The implications of loss spikes for generalization are not yet well understood (spikes may sometimes coincide with improved test performance)
• For complex models like Transformers, analytical tools are limited, and direct application of theoretical results faces challenges

The work aims to serve as a basis for further investigation into training instability in neural networks, contributing to more reliable and interpretable training methodologies.