Residual Koopman Spectral Profiling for Predicting and Preventing Transformer Training Instability

TL;DR

This paper introduces Residual Koopman Spectral Profiling (RKSP), a method to predict and prevent training divergence in transformers by analyzing spectral features at initialization, significantly improving stability and training efficiency.

Contribution

The paper presents RKSP, a novel spectral diagnostic that predicts transformer instability at initialization and introduces Koopman Spectral Shaping (KSS) to prevent divergence during training.

Findings

RKSP achieves an AUROC of 0.995 in predicting divergence.

KSS reduces divergence rate from 66.7% to 12.5%.

Method generalizes across various models and tasks.

Abstract

Training divergence in transformers wastes compute, yet practitioners discover instability only after expensive runs begin. They therefore need an expected probability of failure for a transformer before training starts. Our study of Residual Koopman Spectral Profiling (RKSP) provides such an estimate. From a single forward pass at initialization, RKSP extracts Koopman spectral features by applying whitened dynamic mode decomposition to layer-wise residual snapshots. Our central diagnostic, the near-unit spectral mass, quantifies the fraction of modes concentrated near the unit circle, which captures instability risk. For predicting divergence across extensive configurations, this estimator achieves an AUROC of 0.995, outperforming the best gradient baseline. We further make this diagnostic actionable through Koopman Spectral Shaping (KSS), which reshapes spectra during training. We empirically validate that our method works in practice: RKSP predicts divergence at initialization, and when RKSP flags high risk, turning on KSS successfully prevents divergence. In the challenging high learning rate regime without normalization layers, KSS reduces the divergence rate from 66.7% to 12.5% and enables learning rates that are 50% to 150% higher. These findings generalize to WikiText-103 language modeling, vision transformers on CIFAR-10, and pretrained language models, including GPT-2 and LLaMA-2 up to 7B, as well as emerging architectures such as MoE, Mamba-style SSMs, and KAN.