Spectral Alignment as Predictor of Loss Explosion in Neural Network Training

TL;DR

Loss explosions in large-scale neural network training incur huge costs and are difficult to predict with traditional scalar metrics. The authors introduce Spectral Alignment (SA), a distribution-aware predictor that measures the cosine similarity between layer inputs and the top singular vector of each layer’s weights. They provide a theoretical link between pathological alignment and spectral-norm growth, and demonstrate that SA delivers early, unambiguous warnings in language-model training, outperforming conventional indicators while maintaining low overhead. Empirically, SA detects instability thousands of steps before explosions in Flash Attention and FFN failure modes, offering a practical safeguard for expensive training runs.

Abstract

Loss explosions in training deep neural networks can nullify multi-million dollar training runs. Conventional monitoring metrics like weight and gradient norms are often lagging and ambiguous predictors, as their values vary dramatically across different models and even between layers of the same model, making it difficult to establish a unified standard for detecting impending failure. We introduce Spectral Alignment (SA), a novel, theoretically-grounded metric that monitors the distributional alignment between layer inputs and the principal singular vectors of weight matrices. We show that a collapse in the sign diversity of this alignment is a powerful early predictor of representational collapse and training divergence. Empirical results on language models demonstrate that monitoring the SA distribution provides a significantly earlier and clearer warning of loss explosions than traditional scalar metrics. SA's low computational overhead makes it a practical tool for safeguarding model training.

Figures (8)

• Figure 1: Some conventional metrics, such as weight norm and maximal activation. This figure shows two models investigated in our paper, which shows the inconsistent patterns for conventional metrics across different models. This makes it is difficult to establish a universal threshold for loss explosion prediction.
• Figure 2: Conceptual illustration of Spectral Alignment (SA). (a) SA measures the alignment between input features and the principal singular vector. (b) During stable training, inputs show diverse alignment directions. (c) Impending failure is signaled by a collapse in this diversity, with most inputs aligning in the same direction.
• Figure 3: Comparison of stable and exploding training loss caused by FA and FFN.
• Figure 4: Spectral Alignment distributions for stable and unstable training.
• Figure 5: Spectral norm of weights for stable and unstable training in two failed cases.

Theorems & Definitions (14)

• Definition 1: Spectral Alignment
• Theorem 1: Spectral Norm Growth During Training
• Corollary 1
• Definition 2
• Theorem 2: Gradient Expression for Linear Layers
• proof
• Theorem 3: Change in Spectral Norm under Perturbation
• proof
• Theorem 4: Change of the Spectral Norm
• proof