Residual Koopman Spectral Profiling for Predicting and Preventing Transformer Training Instability

Abstract

Training divergence in transformers wastes compute, yet practitioners discover instability only after expensive runs begin. They therefore need an initialization-time signal for identifying risky training configurations. We study Residual Koopman Spectral Profiling (RKSP), a fixed, single-forward-pass spectral risk score. 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 score achieves an AUROC of 0.995, outperforming the best gradient baseline. When calibrated probabilities are required for a target configuration distribution, a held-out one-dimensional calibrator maps the same scalar score to probabilities. We further make the diagnostic actionable through Koopman Spectral Shaping (KSS), which reshapes spectra during training. In the challenging high learning rate regime without normalization layers, KSS reduces the divergence rate from 66.7