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

> Weight-norm criticality, driven by weight decay shrinking scale-invariant parameters, causes loss spikes by amplifying Hessian curvature, with explicit boundaries predicting instability in Transformers and ResNets.

- **Source:** [arXiv](https://arxiv.org/abs/2607.21005)
- **Published:** 2026-08-15
- **Permalink:** https://picx.dev/p/UtGNZI
- **Whiteboard:** https://picx.dev/p/UtGNZI/image

## Summary

# 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.

---

_Markdown view of https://picx.dev/p/UtGNZI, served by PicX — AI-generated visual whiteboard summaries of research papers._