PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Lianhai Ren, Yucheng Ding, Xiao Liu, Qianxiao Li, Peng Cheng, Yeyun Gong

arXiv:2602.01734v1cs.LG

TL;DR

LLM pretraining can suffer sudden gradient explosions, and the paper investigates whether stable-rank collapse and adjacent-layer Jacobian alignment explain these failures. It develops MSign, which periodically restores weight stable rank, and reports stable training across models from 5M to 3B parameters with overhead below 7%.

  • Problem

    LLM pretraining can experience sudden, unrecoverable gradient explosions and loss growth that waste substantial computational resources.

  • Method

    The paper identifies stable-rank collapse and Jacobian alignment as a failure mechanism, proves their role in gradient growth, and proposes MSign to periodically restore stable rank.

  • Results

    MSign prevents training failures and enables stable convergence across diverse dense and mixture-of-experts models spanning 5M–3B parameters, with overhead below 7%.

  • Takeaways & Limitations

    Low stable rank combined with Jacobian alignment provides a sufficient mechanistic explanation for training failure in the studied regimes, while MSign breaks this feedback loop.

  • Takeaways & Limitations

    The theoretical feedback mechanism relies on strong structural assumptions that may not universally hold, and its full range of applicability remains open.

Abstract

from arXiv · show

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

1. Introduction

LLM pretraining can fail through sudden gradient explosions, and the paper identifies stable-rank collapse and increasing adjacent-layer Jacobian alignment as precursors. It proposes MSign, which restores stable rank and stabilizes training across model scales with limited overhead.

  • Training failures manifest as sudden, unrecoverable gradient explosions and loss growth that can waste substantial computational resources.
  • A 5M-parameter NanoGPT study identifies stable-rank collapse and increasing adjacent-layer Jacobian alignment as phenomena preceding training collapse.
  • Low stable rank increases layer Jacobian spectral norms, while high alignment makes these norms multiply constructively across layers.
  • MSign periodically applies the matrix sign operation to projection weights, equalizing nonzero singular values and restoring the original Frobenius norm.It is typically applied every P = 100 steps.
  • MSign training converges stably across NanoGPT-5M, Sigma-40M, LLaMA-1B, and LLaMA-MoE-3B configurations, while baseline training fails via gradient explosion.The intervention maintains stable rank above critical thresholds, controls Jacobian alignment, and keeps gradient norms bounded.
  • Attention-layer application, particularly to output projections, is sufficient to prevent failures, whereas MLP-only application does not prevent them.
  • MSign demonstrates effectiveness across dense and mixture-of-experts architectures spanning 5M–3B parameters with minimal overhead.

2. Literature Review

The literature documents training instability, low-rank structure, and Jacobian dynamics as important phenomena in deep networks. This paper connects these strands by relating low-rank structure to LLM training instability through stable rank.

  • Prior LLM projects report loss spikes, manual interventions, checkpoint rollbacks, restarts, hardware failures, and gradient explosions during training.
  • Neural network weights and gradients commonly exhibit low-rank or redundant structure, including reports of substantial parameter predictability and low stable rank during transformer training.
  • The paper extends these research directions by connecting low-rank gradient structure directly to training instability through a stable-rank mechanism.
  • Prior work analyzes Jacobian singular values and Jacobian evolution to understand gradient flow and optimization in deep and residual networks.
  • Stable rank has been studied as a robust relaxation of matrix rank and used in neural-network generalization and optimization analyses.

3. Empirical Observations: Training Failure Phenomena

Experiments establish two recurring phenomena before transformer training collapse: sharp stable-rank decline and increasing alignment between adjacent layer Jacobians. The paper proposes these observations as a destabilizing feedback mechanism leading to failure.

  • 3.1. Experimental Setup: The study constructs a reproducible transformer training-failure scenario using a modified NanoGPT configuration with standard hyperparameters.The corresponding μP configuration uses a 0.02 standard deviation and 6 × 10^-4 learning rate at 100M scale.
  • 3.1. Experimental Setup: The model is a decoder-only transformer with L stacked blocks, residual connections, LayerNorm, self-attention, and position-wise MLP sublayers.
  • 3.1. Experimental Setup: Layer Jacobians are defined between consecutive hidden-state representations, and stable-rank and alignment analyses refer to the transformer’s projection matrices and these Jacobians.
  • 3.2. Observations during training failure: Stable rank of the first several layers drops sharply around step 20000 before gradient explosion, indicating energy concentration in top singular values.
  • 3.2. Observations during training failure: Alignment between adjacent layer Jacobians increases during the failure phase, as their leading singular subspaces become increasingly correlated.
  • 3.2. Observations during training failure: High Jacobian alignment correlates with weight-scale growth, stable-rank decline, and gradient growth during the failure phase.
  • 3.3. Conjecture: Low Stable Rank + Jacobian Alignment Drives Training Failure: The paper conjectures that low weight stable rank combined with high Jacobian alignment creates a destabilizing feedback mechanism leading to training failure.

4. Theoretical Analysis: Understanding the Failure Mechanism

The analysis links low stable rank and increasing Jacobian alignment to large total Jacobian and gradient norms, providing a sufficient mechanism for training instability under stated assumptions. It also analyzes feedback toward stable-rank decline and motivates periodic rank restoration.

  • Failure pathway: Low stable rank combined with Jacobian alignment leads to high total Jacobian norms and large weight gradients, forming the paper’s proposed failure pathway.The analysis treats large gradient norms as an indicator of increased divergence risk under a simplifying assumption.
  • Jacobian alignment: When adjacent Jacobian singular subspaces align, cancellation is suppressed and layer norms multiply constructively across the network.Alignment measures whether the output direction of one Jacobian matches the input direction of the next.
  • Jacobian alignment: When aM > 1, the Jacobian-product lower bound grows exponentially with depth, providing a sufficient condition for large total Jacobian norms.Here, M denotes a layer-Jacobian norm lower bound and a denotes an alignment lower bound.
  • Stable rank and Jacobians: Decreasing stable rank increases operator norms under fixed or approximately fixed Frobenius norm, amplifying Jacobian norms across linear, attention, and MLP layers.For attention projections, low stable rank in any projection matrix amplifies the Jacobian norm; across layer types, the dependence is inverse to the square root of stable rank.
  • Gradient growth: Under the paper’s gradient-alignment assumptions, high total Jacobian norms produce large weight-gradient lower bounds, with empirical validation showing the theoretical bound tracks observed gradient growth.The total-gradient result exhibits exponential growth in depth when aM > 1.
  • Feedback and scope: The analysis identifies low-rank gradient propagation and a highly aligned feedback regime in which gradient descent can decrease stable rank, while emphasizing restricted generality.The authors describe the stable-rank-decline result as an existence result rather than a claim about typical training trajectories.

5. The MSign Optimizer: Breaking the Feedback Loop

MSign periodically restores weight-matrix stable rank using matrix sign operations, while preserving scale and limiting computational cost. Applied selectively and periodically, it prevents the training failures observed across model scales.

  • MSign Design: MSign periodically applies sign(W) = UV^T to projection weights, setting non-zero singular values to 1 and maximizing stable rank for a fixed matrix rank.The reduced SVD preserves the matrix’s row and column spaces and produces the correct m × n shape.
  • MSign Design: Frobenius-norm-preserving rescaling restores the original weight scale after the matrix sign operation.This rescaling is used to maintain training dynamics.
  • MSign Design: When stable rank is sufficiently low, Frobenius-norm rescaling may excessively amplify minor singular values, motivating more principled schemes as future work.This is identified as a limitation of the current implementation.
  • Practical Implementation: Applying MSign every P steps, typically P = 100, reduces the cost relative to applying the operation at every update while maintaining effectiveness.The method also supports selective targeting of critical layers, including attention-only or all 2D parameter tensors.
  • Empirical Validation: Across NanoGPT-5M, Sigma-40M, LLaMA-1B, and LLaMA-MoE-3B, MSign maintains stable convergence while baseline training collapses with loss spikes and exploding gradients.The figure compares training loss in the top row and gradient norms in the bottom row.
  • Practical Implementation: Less than 7.0% computational overhead is reported for MSign, with cost depending on application frequency and scope.The paper analyzes overhead relative to the FLOPs of a standard training step and provides a concrete transformer-layer breakdown.

6. Experiment

Experiments across four model configurations show that MSign prevents training collapse while preserving stable convergence, with modest throughput overhead and robust performance across layer selections and application periods.

  • Main result: MSign prevents training collapse across NanoGPT-5M, Sigma-40M, LLaMA-1B, and LLaMA-MoE-3B configurations.Baseline runs exhibit loss spikes and divergence, whereas MSign maintains stable convergence across dense and sparse architectures.
  • Main result: Baseline gradient norms reach 10^1–10^7 before failure, while MSign keeps them bounded within 10^0 throughout training.The periodic pattern in MSign runs corresponds to the application period P = 100.
  • Throughput analysis: 4.6–6.7% measured throughput overhead occurs on larger models, compared with −2.4% for NanoGPT-5M and a theoretical prediction below 0.1%.Implementation factors include distributed SVD communication, disrupted kernel fusion, and pipeline bubbles.
  • Layer selection ablation: Attention-only MSign stabilizes NanoGPT-5M and Sigma-40M, whereas MLP-only application fails to prevent collapse.Applying MSign to all 2D parameters yields test perplexity 102.6 versus 118.6 on NanoGPT-5M.
  • Application period ablation: P values from 10 to 10,000 prevent collapse, but P = 10000 produces more variable loss and gradient norms than smaller periods.P = 100 is recommended as a conservative default because it maintains stable dynamics with acceptable overhead.

7. Conclusion

The paper identifies a stable-rank-collapse feedback loop underlying LLM training instability and proposes periodic matrix-sign restoration as a low-overhead intervention, while acknowledging limits to the theory’s generality.

  • Conclusion: Low stable rank amplifies layer Jacobian norms, and inter-layer alignment causes these norms to multiply constructively, producing exponential gradient growth.MSign breaks this feedback loop by periodically restoring weight stable rank via the matrix sign operation.
  • Limitations: The theoretical analysis relies on strong structural assumptions, including uniform negative correlation of input and output gradient projections.The paper leaves complete characterization and relaxation of these conditions for future work.
  • Practical solution: The matrix-sign operation has O(d^3) cost per weight matrix, but amortization over P = 100 steps makes its overhead negligible relative to forward and backward passes.The forward and backward pass costs scale as O(BTd^2).

A.4. Throughput Model Analysis

The throughput analysis models MSign’s amortized cost as inversely related to its application period, while measured overhead exceeds FLOPs-based predictions because of distributed-system effects.

  • Throughput model: MSign’s amortized overhead per token is F/P when applied every P steps.The throughput model defines F as additional computation per application.
  • Throughput model: Throughput is modeled as T(P) = T∞ / (1 + r/P), where r = F/f and T∞ is baseline throughput.T∞ denotes asymptotic throughput as P approaches infinity.
  • Model fitting: Least-squares fitting uses Sigma-40M measurements after linearizing the reciprocal-throughput relationship.The fitted model is based on 1/T(P) = (1/T∞)(1 + r/P).
  • Model fitting: The fitted r ≈3.9 substantially exceeds the FLOPs-based prediction r_theory ≪1.The theoretical overhead ratio is R < 0.1% for typical configurations.
  • Implementation effects: The discrepancy is attributed to all-gather synchronization, disrupted FlashAttention fusion and CUDA streams, and distributed-training pipeline bubbles.These latency-dominated effects scale poorly with batch size.

A.5. Application Period Analysis

Application-period experiments show that frequent MSign updates provide the most stable dynamics, while infrequent updates permit transient or intermittent instability.

  • Application periods: P = 10000 shows noticeably higher gradient norms during steps 20000 to 40000, indicating intermittent instability from infrequent applications.All tested periods from P = 10 to P = 10000 eventually converge.
  • Application periods: P = 10 and P = 100 produce smooth training loss and gradient-norm dynamics with minimal variance.Frequent applications maintain stable rank above critical thresholds throughout training.
  • Application periods: P = 1000 remains stable but exhibits slightly greater variance as longer intervals allow transient stable-rank decline.Recovery occurs before instability develops.
  • Application periods: P = 100 is recommended as the default because it provides a safety margin against instability with negligible computational overhead.The recommendation follows the observed stability of P = 100 and its low amortized cost.
  • Mechanism: Low stable rank amplifies attention Jacobian norms, especially through the V-O pathway, whose scale is proportional to (srank(WV)srank(WO))^-1/2.The analysis identifies V and O projections as typically lower-stable-rank contributors than Q and K.

B.4. Proof of Theorem 4.6 (Jacobian Norm Bound: MLP Layer)

The MLP Jacobian analysis bounds layer sensitivity using weight operator norms and activation Lipschitzness, then connects low stable rank to increasing Jacobian norms and gradient growth.

  • MLP Jacobian bound: The two-layer MLP Jacobian is bounded by the product of the two weight operator norms and the activation derivative norm.Submultiplicativity gives ∥W1 · diag(ϕ′(z)) · W2∥2 ≤ ∥W1∥2∥diag(ϕ′(z))∥2∥W2∥2.
  • MLP Jacobian bound: The activation derivative norm is bounded by the Lipschitz constant Lϕ, with GELU ≈1.13 and SiLU ≈1.1.These are the activation constants used in the bound.
  • Stable-rank connection: Expressed through stable rank, the MLP Jacobian bound scales as Lϕ∥W1∥F∥W2∥F divided by the square root of the stable-rank product.The displayed bound is given in equation (65).
  • Stable-rank connection: Low stable ranks can increase the Jacobian norm when Frobenius norms remain approximately constant or grow moderately over short training windows.This increase can contribute to gradient explosion.
  • Unified interpretation: Across linear, attention, and MLP layers, Jacobian norms are inversely related to the square root of stable rank.Combined with Jacobian alignment, stable-rank collapse creates conditions for exponential gradient growth.
  • Gradient propagation: The theoretical argument uses local-Jacobian alignment assumptions to derive bounds on cumulative Jacobians and backpropagated gradients.The assumptions are empirically motivated but stylized.

B.6. Proof of Theorem 4.9 (Total Gradient Norm Lower Bound)

The proof aggregates per-weight gradient lower bounds across layers and uses a geometric-series argument to show exponential total gradient growth when aligned Jacobian amplification exceeds one.

  • Gradient aggregation: The total bound is obtained by aggregating per-weight-vector bounds across all columns and layers.The column-wise gradient contributions are squared and summed before summing over layers.
  • Geometric-series bound: The resulting layer contributions form the geometric series 1 + (aM)^2 + (aM)^4 + ··· + (aM)^{2(L−1)}.The common ratio is (aM)^2.
  • Geometric-series bound: When aM > 1, the series is approximately O((aM)^{2L}) for large depth L.The highest-depth term dominates the sum.
  • Conclusion: The total gradient norm therefore grows exponentially with network depth L.This is the conclusion of the total-gradient lower-bound proof.
  • Connection to observations: Stable-rank decline raises layer Jacobian norms while increasing alignment drives a toward one, together making aM exceed one and triggering gradient explosion.This links the proof to the two empirical observations preceding failure.
  • Attention-gradient structure: Attention gradients inherit rank constraints from the hidden states; if cohidden states have rank at most r, the output-projection gradient has rank at most r.The proof extends this rank argument across all four attention gradients.

B.8. Proof of Theorem 4.12 (Stable Rank Feedback Mechanism)

The proof analyzes gradient descent in the singular-vector basis of a weight matrix and shows that aligned input-output structure reduces stable rank. Under the stated covariance and perturbation assumptions, singular-value dynamics explain this stable-rank feedback mechanism.

  • The weight gradient is expressed as the outer product of the backpropagated output gradient and the layer input.
  • Under aligned input-output covariances, projecting onto the singular bases makes the gradient covariance approximately diagonal.The analysis assumes the projections are concentrated along corresponding top singular directions and uses zero-mean projections for simplicity.
  • The gradient-descent update preserves the singular vectors to first order while modifying the singular values.This conclusion uses the SVD-based update and the assumption that the perturbation is small relative to singular-value gaps.
  • Under the negative-correlation condition, the analyzed term is nonpositive, so stable rank decreases under gradient descent with aligned input-output structure.The proof connects the covariance inequality to the sign of the stable-rank differential.
  • All singular values increase under the same analysis, consistent with the observed growth of the weight norm.
  • The resulting stable-rank decrease is identified as the theorem’s feedback mechanism under the alignment assumptions.
Loading 2602.01734v1…