MSign: An Optimizer Preventing Training Instability in Large Language Models via Stable Rank Restoration

TL;DR

This work identifies a fundamental instability mechanism in large-scale transformer pretraining: simultaneous stable rank collapse and increasing inter-layer Jacobian alignment cause exponential gradient growth as depth increases. It theoretically connects low stable rank and Jacobian alignment to large total Jacobian norms and gradient magnitudes, establishing a causal failure pathway. To break this feedback loop, the authors propose MSign, a matrix-sign-based optimizer that periodically restores weight stable rank by projecting weight matrices to partial isometries while preserving the Frobenius norm, with targeted application to attention projections and a modest overhead. Empirical validation across 5M–3B parameter models (dense and MoE) demonstrates that MSign prevents training failures with less than 7% overhead, maintaining stable rank and bounded gradients; ablations show attention-layer targeting is essential and provide guidance on default application period, delivering practical stability improvements for scalable LLM pretraining.

Abstract

Training instability remains a critical challenge in large language model (LLM) pretraining, often manifesting as sudden gradient explosions that waste significant computational resources. We study training failures in a 5M-parameter NanoGPT model scaled via $μ$P, identifying two key phenomena preceding collapse: (1) rapid decline in weight matrix stable rank (ratio of squared Frobenius norm to squared spectral norm), and (2) increasing alignment between adjacent layer Jacobians. We prove theoretically that these two conditions jointly cause exponential gradient norm growth with network depth. To break this instability mechanism, we propose MSign, a new optimizer that periodically applies matrix sign operations to restore stable rank. Experiments on models from 5M to 3B parameters demonstrate that MSign effectively prevents training failures with a computational overhead of less than 7.0%.

Figures (4)

• Figure 1: Correlation between training failure indicators and gradient norm explosion. Left (Observation 1): Stable rank (geometric mean across early layers) vs. gradient norm over training steps. As stable rank declines sharply around step 20000, gradient norms begin explosive growth. Right (Observation 2): Jacobian alignment (average between adjacent layers) vs. gradient norm. Increasing alignment precedes and accompanies gradient explosion.
• Figure 2: Validation of Theorem \ref{['thm:jacobian-product']}: Jacobian product norm lower bound vs. actual gradient norm. The theoretical bound closely tracks observed gradient growth.
• Figure 3: MSign prevents training failures across model scales. Top row: Training loss comparison between baseline (blue) and MSign (orange). Baseline training collapses with sudden loss spikes, while MSign maintains stable convergence. Bottom row: Corresponding gradient norm dynamics. Baseline runs exhibit exponential gradient explosion preceding collapse, while MSign keeps gradient norms bounded throughout training. Training is terminated after failure to conserve computational resources. Results demonstrate that MSign effectively breaks the stable rank collapse feedback loop identified in our theoretical analysis. From left to right: NanoGPT-5M, Sigma-40M, LLaMA-1B, LLaMA-MoE-3B.
• Figure 4: Training dynamics under different MSign application periods on NanoGPT-5M. Left: Training loss comparison. Right: Gradient norm comparison. While all periods from $P=10$ to $P=10000$ eventually converge, $P=10000$ exhibits noticeably higher gradient norm in step 20000 to 40000, indicating intermittent instability when MSign applications are too infrequent.

Theorems & Definitions (15)

• Definition 3.1: Stable Rank
• Definition 3.2: Matrix Alignment
• Remark 4.1: Simplifying Assumption
• Theorem 4.2: Jacobian Product Norm Lower Bound
• Remark 4.3: Exponential Growth Condition
• Theorem 4.4: Stable Rank Controls Jacobian Norm: Linear Layer
• Theorem 4.5: Jacobian Norm Bound: Attention Layer
• Theorem 4.6: Jacobian Norm Bound: MLP Layer
• Theorem 4.8: Weight Gradient Norm Lower Bound
• Theorem 4.9: Total Gradient Norm Lower Bound