Technical Summary: The Stability of Singular Distribution (SoSD)

Problem Statement

Large Language Model (LLM) pre-training consistently exhibits a characteristic "fast-then-slow" two-phase trajectory: an initial regime of precipitous loss decay followed by a prolonged period of asymptotic, heavy-tailed improvement. While this empirical pattern is well-documented across major technical reports, the underlying theoretical mechanisms dictating the transition from the fast-learning phase to the slow-saturation phase remain fundamentally under-explored. Existing analyses often rely on restricted task formulations (e.g., in-context learning or linear regression) or focus on specific structural phenomena like rank collapse, failing to provide a unified spectral perspective on the macroscopic temporal evolution of the loss landscape in general pre-training scenarios.

Methodology

The authors investigate the optimization dynamics of LLMs by shifting the analytical lens from parameter matrix convergence to the temporal evolution of their spectral properties.

Empirical Approach

The study conducts pre-training experiments on two decoder-only model families:

• GPT-2 (Small and Medium) trained on the FineWeb dataset.
• LLaMA (0.5B and 2B) trained on the C4 dataset.

The authors monitor the cosine similarity between current states and final trained states for both parameter matrices ($W$) and their singular value matrices ($\Sigma$). They define the Singular Distribution (SD) as the trace-normalized spectrum ($\hat{\Sigma}_t = \Sigma_t / \text{tr}(\Sigma_t)$) and measure the Singular Distribution Variation (SD Variation) as the Frobenius norm of the update to this normalized spectrum ($\Delta \Sigma_t = \|\hat{\Sigma}_{t+1} - \hat{\Sigma}_t\|_F$).

Theoretical Framework

To explain the observed phenomena, the authors analyze a simplified single-layer, single-head Transformer under full-batch Gradient Descent with small initialization. They establish a theoretical framework based on:

1. Strictly Increasing Norms: Weight norms grow monotonically during training (in the absence of weight decay).
2. Non-Degeneracy and Gradient Boundedness: Parameters remain full-rank with bounded condition numbers, and gradients are non-vanishing but bounded.
3. Smoothness and Margin Conditions: Assumptions regarding local smoothness of the loss landscape and positive margins in attention and logit distributions during the late training regime.

Key Contributions

1. Identification of the Stability of Singular Distribution (SoSD)

The paper identifies a phenomenon termed Stability of Singular Distribution (SoSD). Empirical results demonstrate that the trace-normalized singular value spectrum ($\hat{\Sigma}$) stabilizes significantly earlier than the parameter matrices ($W$) themselves.

• Spectral Decoupling: While weight entries remain distant from their final configuration, the relative shape of the singular value spectrum saturates near 1.
• Synchronization: The onset of SoSD (where SD Variation drops to a metastable floor) aligns precisely with the transition of the validation loss from the fast-descent phase to the slow-descent (asymptotic) regime. This synchronization is observed across diverse architectures (GPT-2, LLaMA), modules (Attention, MLP), and optimization settings.

2. Theoretical Analysis of SoSD and Two-Phase Dynamics

The authors prove that the emergence of SoSD is inevitable under standard training conditions and that it mechanistically bounds the rate of loss reduction.

• Theorem 4.5 (Emergence of SoSD): Under the assumption of monotonically increasing weight norms, the stability bound for the singular distribution scales as $\varepsilon \propto \eta / \|W\|$. As training progresses and $\|W\|$ grows, the permissible variation in the singular distribution shrinks, forcing the system into a stable state.
• Theorem 4.9 (Two-Phase Dynamics):
  • Phase I (Fast Descent): Before the SoSD threshold ($t < T^*$), moderate parameter norms allow for significant spectral shifts, enabling a constant-order loss decrease ($\Delta L(t) = O(1)$).
  • Phase II (Slow Descent): After the SoSD threshold ($t > T^*$), the singular distribution stabilizes. The loss decrement becomes kinetically constrained and is theoretically bounded by the magnitude of the SD variation ($\Delta L(t) \le O(\varepsilon^p)$). The stabilization of $\hat{\Sigma}$ acts as a bottleneck, forcing the loss into an asymptotic plateau.

3. Interpretation of Pre-training Strategies

The paper interprets standard optimization strategies through the lens of modulating the SoSD stability bound ($\varepsilon \propto \eta / \|W\|$):

• Learning Rate Schedules: Strategies like Step-wise Decay, Warmup-Stable-Decay (WSD), and Cosine Decay function by reducing the learning rate $\eta$. This tightens the stability bound, forcing the system out of its current quasi-equilibrium and allowing for further loss minimization by resolving finer-grained features.
• Weight Decay: By suppressing the growth of the weight norm $\|W\|$, weight decay relaxes the stability constraint (increasing $\varepsilon$). This permits larger updates to the singular distribution, effectively raising the theoretical upper bound for loss reduction and preventing premature saturation.
• Optimizers (Muon vs. Adam): The SoSD phenomenon persists even with the Muon optimizer, confirming it as a fundamental characteristic of Transformer pre-training. However, Muon maintains a significantly lower magnitude of SD variation compared to Adam, suggesting it navigates the spectral landscape with higher efficiency.

Results

• Empirical Validation: Experiments on GPT-2 and LLaMA models confirm that singular distributions stabilize early (often within the first 10-20% of training steps), while parameters continue to evolve. This stabilization coincides with the validation loss entering a heavy-tailed, slow-descent regime.
• Correlation: A tight synchronization is observed between the decay of SD Variation and the transition to the slow-descent phase of the loss curve.
• Strategy Efficacy:
  • Reducing the learning rate induces a sharp decline in SD Variation, synchronizing with renewed loss descent.
  • Weight decay models achieve lower final losses than non-regularized baselines, correlating with higher sustained SD Variation (due to relaxed constraints on $\|W\|$).
  • The Muon optimizer achieves lower loss and faster convergence while maintaining lower SD Variation magnitudes than Adam.

Significance and Claims

The paper claims to offer a novel mechanistic understanding of Transformer optimization trajectories by linking microscopic spectral behaviors to macroscopic training dynamics.

• Unified Perspective: It provides a spectral lens to explain the ubiquitous "fast-then-slow" convergence pattern, positing that the transition is not merely an artifact of specific layers but a global kinetic constraint characterized by the saturation of the singular distribution.
• Theoretical Grounding: The work establishes a theoretical framework proving that the growth of weight norms inevitably precipitates an early SoSD threshold, after which the rate of loss decrease is strictly bounded by the variation in the singular distribution.
• Practical Insight: By framing learning rate schedules and weight decay as mechanisms to modulate the SoSD scale, the paper offers a rationale for why these strategies facilitate efficient pre-training, suggesting that overcoming the spectral barrier is key to achieving lower training losses.

The authors maintain a modest stance, noting that while SoSD acts as a characterization of singular value dynamics and an indicator of the slow-descent regime, a lower SoSD value does not inherently imply better performance (as evidenced by the weight decay experiments where higher SoSD coexists with lower loss). The primary contribution is the identification of the synchronization between spectral stability and optimization saturation.