Residual Koopman Spectral Profiling for Predicting and Preventing Transformer Training Instability

Bum Jun Kim; Shohei Taniguchi; Makoto Kawano; Yusuke Iwasawa; Yutaka Matsuo

cs.LGFeb 26, 2026

Residual Koopman Spectral Profiling for Predicting and Preventing Transformer Training Instability

Authors: Bum Jun Kim, Shohei Taniguchi, Makoto Kawano, Yusuke Iwasawa, Yutaka Matsuo

Organizations: Graduate School of Engineering, The University of Tokyo, Japan

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.

Figures & tables

Figure 1: Scatter plot of DMD eigenvalues across layers in a pre-layer normalization transformer. The color gradient indicates layer depth; blue is early and red is late. Early layers cluster near the unit circle; late layers exhibit an increased spectral radius.

Norm Type	$n$	Div.% (lower)	$M_{\approx 1}$	$M_{>1}$	$\rho$	Acc.% (higher)
Pre-LN	25	$\mathbf{0.0}$	$0.16\pm 0.01$	$0.54\pm 0.01$	$2.29\pm 0.04$	$7.5\pm 8.7$
Post-LN	69	$\mathbf{0.0}$	$0.66\pm 0.02$	$0.31\pm 0.01$	$7.48\pm 0.12$	$1.4\pm 5.8$
RMSNorm	12	$\mathbf{0.0}$	$0.16\pm 0.01$	$0.54\pm 0.01$	$2.28\pm 0.06$	$13.1\pm 8.6$
DeepNorm	12	$\mathbf{0.0}$	$0.00\pm 0.00$	$1.00\pm 0.00$	$3.94\pm 0.21$	$\mathbf{20.0\pm 9.4}$
SubLN	12	$\mathbf{0.0}$	$0.13\pm 0.02$	$0.54\pm 0.01$	$2.23\pm 0.03$	$9.0\pm 8.1$
No-Norm	84	$96.4$	$0.80\pm 0.02$	$0.19\pm 0.01$	$1.11\pm 0.01$	$54.2\pm 24.4^{*}$

Table 1: Comprehensive normalization comparison across different setups. RKSP metrics reveal distinct spectral signatures explaining stability differences. Accuracy is computed from only 3 converged runs out of 84. Statistical significance: No-Norm divergence compared with others,

p<10^{-50}

via Fisher’s exact test. AUROC for divergence prediction using

M_{\approx 1}

: 0.995 [95% CI: 0.986 to 1.00]. Measured on the associative-recall task.

Predictor	AUROC (higher)	95% CI	Timing
$M_{\approx 1}$ at initialization	$\mathbf{0.995}$	$[0.986,1.000]$	Initialization
$M_{\approx 1}\times\log_{10}(\kappa(\mathbf{V}))$	$0.873$	$[0.841,0.905]$	Initialization
Spectral radius $\rho$ at init	$0.845$	$[0.808,0.882]$	Initialization
Eigenvector condition $\kappa(\mathbf{V})$ at init	$0.687$	$[0.638,0.736]$	Initialization
Gradient-Based Baselines
Initial gradient norm	$0.621$	$[0.568,0.674]$	After 1 step

Table 2: AUROC for divergence prediction with bootstrap 95% confidence intervals. The

M_{\approx 1}

CI lower bound of 0.986 exceeds gradient baselines’ upper bounds. Measured on the associative-recall normalization sweep.

Method	Div.% (lower)	Acc.% (higher)	$M_{\approx 1}$	Overhead
No control	$66.7$	$28.5$	$0.85$	—
Gradient clip 0.5	$50.0$	$32.1$	$0.82$	$<1\%$
Gradient clip 1.0	$58.3$	$30.8$	$0.83$	$<1\%$
KSS $\alpha=0.10$	$25.0$	$42.3$	$0.68$	$9.5\%$
KSS $\alpha=0.15$	$\mathbf{12.5}$	$\mathbf{48.2}$	$0.58$	$10.8\%$
KSS $\alpha=0.20$	$8.3$	$46.5$	$0.52$	$11.2\%$

Table 3: Gradient clipping versus KSS in a challenging No-Norm setting with learning rate that ranges from 0.005 to 0.01 and 24 trials each. Measured on the associative-recall task.

Figure 2: KSS regularization effectiveness. (Left) The divergence rate decreases with KSS weight

\alpha

. (Right) A dual axis shows accuracy improvement and

M_{\approx 1}

shifting downward toward the target band. KSS shapes spectral properties, improving both stability and performance. Measured on the associative-recall task.

Method	Div.% (lower)	Acc.% (higher)	Overhead
No stabilization, AdamW	$66.7$	$28.5$	—
Gradient clipping at 1.0	$58.3$	$30.8$	$<1\%$
Spectral normalization	$41.7$	$35.2$	$5.2\%$
Weight normalization	$54.2$	$32.8$	$3.8\%$
Gradient penalty with $\lambda=0.1$	$45.8$	$33.1$	$6.4\%$
Advanced Optimizers

Table 4: Extended baseline comparison including SAM, spectral normalization, and the Lion optimizer. Measured on the associative-recall task.

Norm Type	Max LR w/o KSS	Max LR w/ KSS	Increase
Pre-LN	$0.005$	$0.008$	$+60\%$
RMSNorm	$0.003$	$0.005$	$+67\%$
No-Norm	$0.002$	$0.005$	$+150\%$

Table 5: Maximum stable learning rate (LR). KSS enables learning rates that are 50% to 150% higher. The stability criterion is

<

20% divergence across trials. Measured on the associative-recall task.

Method	Div.% (lower)	Acc.% (higher)	$M_{\approx 1}$	Mechanism
No regularization	$66.7$	$28.5$	$0.85$	—
Generic Regularization, about 11% overhead
Random $\ell_{2}$ on residuals	$54.2$	$31.2$	$0.81$	Magnitude damping
Jacobian penalty	$45.8$	$33.8$	$0.78$	Gradient smoothness
Activation variance reg.	$50.0$	$32.5$	$0.80$	Distribution control
Spectral-Specific Regularization, about 11% overhead

Table 6: Mechanism ablation: KSS versus random regularization. Same overhead at about 11%, different mechanisms. Results use the No-Norm setting, with learning rates ranging from 0.005 to 0.01 across 24 trials each. All methods are matched to about 11% computational overhead. Measured on the associative-recall task.

Model	Params	Method	Div. of 24 (lower)	PPL (lower)	Overhead
Extra Large, $d=1024$ , $L=24$	350M	Baseline	6 of 24	$52.1$	—
Extra Large, $d=1024$ , $L=24$	350M	KSS	3 of 24	$\mathbf{46.8}$	$11.8\%$

Table 7: Scale-up KSS training results up to 350M parameters. Measured on the synthetic LM validation set with 24 trials on a synthetic LM task with vocab size 10K and a sequence length of 256.

Model	Params	Method	Div. of 24 (lower)	PPL (lower)	Max LR	Steps
Small, $d=256$ , $L=6$	25M	Baseline	2 of 24	$58.4$	$3\mathrm{e}{-4}$	10K
Small, $d=256$ , $L=6$	25M	KSS	0 of 24	$\mathbf{52.1}$	$5\mathrm{e}{-4}$	10K
Medium, $d=512$ , $L=12$	85M	Baseline	4 of 24	$42.3$	$2\mathrm{e}{-4}$	10K
Medium, $d=512$ , $L=12$	85M	KSS	1 of 24	$\mathbf{38.5}$	$4\mathrm{e}{-4}$	10K
Large, $d=768$ , $L=12$	125M	Baseline	5 of 24	$48.7$	$1.5\mathrm{e}{-4}$	10K
Large, $d=768$ , $L=12$	125M	KSS	2 of 24	$\mathbf{43.2}$	$3\mathrm{e}{-4}$	10K

Table 8: KSS improves stability and PPL. Pre-LN normalization with 24 trials per configuration. Measured on WikiText-103 and OpenWebText language-modeling tasks.

Model	Norm	Method	Div. of 24 (lower)	Acc.% (higher)	$M_{\approx 1}$
ViT-Tiny, $d=192$ , $L=6$	Pre-LN	Baseline	1 of 24	$72.3$	$0.42$
ViT-Tiny, $d=192$ , $L=6$	Pre-LN	KSS	0 of 24	$\mathbf{75.8}$	$0.38$
ViT-Tiny, $d=192$ , $L=6$	RMSNorm	Baseline	2 of 24	$70.1$	$0.48$
ViT-Tiny, $d=192$ , $L=6$	RMSNorm	KSS	0 of 24	$\mathbf{74.2}$	$0.41$

Table 9: ViT RKSP analysis and KSS training. Results use 24 trials, with an image size of 224, a patch size of 16, and a 5-epoch sanity check. Measured on CIFAR-10.

Layer Group	$\eta_{\mathrm{nl}}$	$M_{\approx 1}$	$\kappa(\mathbf{V})$	Pattern
Early 0 to 10	$0.45\pm 0.05$	$0.74$	$8.2$	Lower $\eta_{\mathrm{nl}}$
Middle 11 to 21	$0.58\pm 0.08$	$0.68$	$9.5$	Mid $\eta_{\mathrm{nl}}$
Late 22 to 31	$0.72\pm 0.10$	$0.61$	$10.8$	Higher $\eta_{\mathrm{nl}}$

Table 10: LLaMA-2-7B analysis. The pattern persists at the 7B scale with RMSNorm. Computed from residual-stream activations on a fixed set of short prompt sentences. Measured on a fixed short-prompt set.

Figure 3: Scaling law for spectral properties. (Left) The near-unit mass

M_{\approx 1}

decreases with model scale, and larger models have more contractive dynamics, implying reduced memory and weaker near-isometric propagation. (Right) The normalized linear-fit error

\eta_{\mathrm{nl}}

increases with scale, indicating a less reliable linear approximation at scale. Log-linear fits are shown. Computed from residual-stream activations on a fixed set of short prompt sentences.

Architecture	$\eta_{\mathrm{nl}}$	$M_{\approx 1}$	$\rho$	Stability Character
Transformer, Pre-LN	$0.52$	$0.16$	$2.29$	Balanced, tunable
Transformer, No-Norm	$0.48$	$0.79$	$1.11$	High memory, risky
MoE, top-2, 8 experts	$0.55$	$0.58$	$2.12$	Routing-dependent
Mamba (SSM)	$0.42$	$0.12$	$0.95$	Contractive, short-memory
KAN	$0.78$	$0.52$	$2.15$	Higher $\eta_{\mathrm{nl}}$

Table 11: Cross-architecture spectral comparison. Transformer rows use the synthetic associative-recall task with seq len 64, vocab 256, and

n_{\mathrm{pairs}}=4

; MoE, Mamba, and KAN rows use the synthetic LM task with random-token next-token prediction.

Appendix figures & tables11 assets

Supplementary material from the paper’s appendix.

Appendix

Symbol	Meaning
Core sizes and indices
$L$	The number of layers.
$\ell$	Layer index.
$d$	Hidden dimension.
$N$	The number of snapshots.
$r$	Randomized DMD rank, the number of eigenvalues used in KSS.

Appendix

Table 12: List of notations used in the main paper and supplementary material.

Method	Div.% (lower)	Acc.% (higher)	$M_{\approx 1}$	Transfer	Overhead
Standard Init, AdamW	$66.7$	$28.5$	$0.85$	✗	—
$\mu$ P, width transfer	$41.7$	$35.8$	$0.72$	✓	$<1\%$
$\mu$ P + higher learning rate	$50.0$	$38.2$	$0.78$	✓	$<1\%$
KSS, $\alpha=0.15$	$\mathbf{12.5}$	$\mathbf{48.2}$	$0.58$	✗	$10.8\%$
$\mu$ P + KSS	$\mathbf{8.3}$	$\mathbf{51.5}$	$0.52$	✓	$11.2\%$

Appendix

Table 13: Comparison of

\mu

P, standard parameterization, and KSS. Results use the No-Norm setting with learning rate ranges from 0.005 to 0.01 across 24 trials. Measured on the associative-recall task.

Optimizer	Div.% (lower)	Acc.% (higher)	$M_{\approx 1}$	Best LR	Overhead
AdamW	$66.7$	$28.5$	$0.85$	$0.003$	—
AdamW + grad clip	$58.3$	$30.8$	$0.83$	$0.005$	$<1\%$
LAMB	$45.8$	$36.2$	$0.75$	$0.01$	about 5%
LAMB + warmup	$37.5$	$40.8$	$0.68$	$0.02$	about 5%
Lion	$45.8$	$36.5$	$0.78$	$0.001$	about 15%
KSS, $\alpha=0.15$	$\mathbf{12.5}$	$\mathbf{48.2}$	$0.58$	$0.008$	$10.8\%$

Appendix

Table 14: Optimizer comparison among AdamW, LAMB, and Lion. Results use the No-Norm setting, 24 trials each. Measured on the associative-recall task.

Model	Layers	$\rho$	$\kappa(\mathbf{V})$	$\eta_{\mathrm{nl}}$	$M_{\approx 1}$
GPT-2 124M	Early 0 to 3	$1.15$	$8.3$	$0.52$	$0.68$
GPT-2 124M	Middle 4 to 7	$1.23$	$10.1$	$0.61$	$0.62$
GPT-2 124M	Late 8 to 11	$1.31$	$11.2$	$0.71$	$0.58$
GPT-2 355M	Early 0 to 7	$1.12$	$7.5$	$0.48$	$0.72$
GPT-2 355M	Middle 8 to 15	$1.19$	$9.2$	$0.58$	$0.65$
GPT-2 355M	Late 16 to 23	$1.27$	$10.8$	$0.68$	$0.60$

Appendix

Table 15: GPT-2 layer-wise spectral analysis. Start Linear, End Nonlinear pattern, shorthand for increasing

\eta_{\mathrm{nl}}

. Spectral statistics are computed from residual-stream activations on a fixed set of short prompt sentences. Measured on a fixed short-prompt set.

Figure 4: Start Linear, End Nonlinear pattern. Layer-wise normalized linear-fit error

\eta_{\mathrm{nl}}

across four pretrained models. All models exhibit a monotonically increasing

\eta_{\mathrm{nl}}

with depth, suggesting a consistent linear-approximation signature across models. Computed from residual-stream activations on a fixed set of short prompt sentences.

Figure 5: Calibration reliability diagram. (Left) Predicted divergence probability versus observed frequency, with an ECE of 0.283. (Right) Distribution of predictions separated by actual outcome. RKSP provides moderately calibrated probability estimates. Based on associative-recall runs, calibration compares predictions to divergence outcomes from that task.

Configuration	Top- $k$	$M_{\approx 1}$	$\rho$	$\eta_{\mathrm{nl}}$	Div. of 24	Acc.%
MoE-Small, 8 experts	$k=1$	$0.72$	$2.85$	$0.68$	6 of 24	$32.4$
MoE-Small, 8 experts	$k=2$	$0.58$	$2.12$	$0.55$	2 of 24	$41.7$
MoE-Small, 8 experts	$k=4$	$0.45$	$1.78$	$0.48$	1 of 24	$38.2$
MoE-Medium, 16 experts	$k=2$	$0.65$	$2.45$	$0.61$	4 of 24	$38.9$
MoE-Medium, 16 experts	$k=2$ + KSS	$0.48$	$1.92$	$0.58$	1 of 24	$\mathbf{44.5}$

Appendix

Table 16: MoE transformer with RKSP analysis. Routing instability revealed via spectral signatures. Results use

d=256

L=6

, and 24 trials. Load balancing loss

\lambda=0.01

. Measured on the synthetic LM task with random-token next-token prediction.

Model	$M_{\approx 1}$	$M_{<1}$	$\rho$	Div. of 24	Acc.%
Mamba-Small, $d=256$	$0.12$	$0.85$	$0.95$	0 of 24	$48.2$
Mamba-Medium, $d=512$	$0.15$	$0.82$	$0.97$	0 of 24	$52.6$
Transformer, Pre-LN, comparable	$0.42$	$0.38$	$1.85$	1 of 24	$45.8$

Appendix

Table 17: Mamba with RKSP analysis. Inherently stable spectral structure. Results use

L=6

, 24 trials. Measured on the synthetic LM task with random-token next-token prediction.

Model	$\eta_{\mathrm{nl}}$	$M_{\approx 1}$	$\rho$	Div. of 24	DMD reliability
KAN-Transformer, $B=4$	$0.78$	$0.52$	$2.15$	3 of 24	Marginal, 68%
KAN-Transformer, $B=8$	$0.82$	$0.48$	$2.35$	4 of 24	Low, 52%
KAN-Transformer + KSS	$0.75$	$0.42$	$1.92$	1 of 24	Improved, 78%

Appendix

Table 18: KAN transformer with RKSP analysis. The B-spline nonlinearity challenges linear approximation. The B-spline order is

B

. We use

d=256

L=6

, and 24 trials. Measured on the synthetic LM task with random-token next-token prediction.

Figure 6: Cross-architecture spectral radar chart. Comparison of five architectures across five normalized metrics. Mamba exhibits strong contraction with low

M_{\approx 1}

and short memory; stability is maintained via suppressed unstable modes, while in near-normal transformer regimes, higher

M_{\approx 1}

aligns with more unstable, near-isometric propagation. The No-Norm transformer shows high memory capacity but poor stability. KAN exhibits high

\eta_{\mathrm{nl}}

. Metrics are derived from Table 11 .

Panel of Figure 7 — Figure 7: Decision flowchart for when to use RKSP and KSS in practice.

Transformer

Transformers are neural architectures that use attention to model relationships among sequence elements in language, vision, and scientific data.

See all 540 papers

Mar 14, 2026

Spectral Edge Dynamics of Training Trajectories: Signal--Noise Geometry Across Scales

Despite hundreds of millions of parameters, transformer training trajectories evolve within only a few coherent directions. We introduce Spectral Edge Dynamics (SED) to quantify this structure: a rolling-window SVD of parameter updates reveals a sharp boundary -- the spectral edge -- between coherent optimization directions and stochastic noise, identified via the maximum consecutive singular value ratio

σ_k / σ_{k+1}

. Across a 51M-parameter TinyStories model (4 seeds) and GPT-2 124M under distribution shift, the spectral edge exhibits a universal three-phase pattern (rise, plateau, collapse). The effective signal rank adapts to task complexity (

k^* = 2

at 51M,

k^* = 3

at 124M), and the directional coupling between spectral geometry and validation loss reverses with window size -- a lag flip reflecting the timescale of trajectory integration. Johnson--Lindenstrauss projection to

d = 10W

dimensions (e.g.,

d = 100

for

W = 10

) preserves the spectral gap within

5.7\%

, making the framework applicable to models of arbitrary scale. In companion work, the same spectral geometry provides early-warning signals of grokking -- predicting generalization

600

1{,}700

steps before it occurs across modular arithmetic, Dyck languages, and the SCAN benchmark.

Yongzhong Xu

Transformer Training Dynamics Neural Network Training Dynamics Gradient Descent Dynamics Spectral Methods

Ask PDF

Apr 3, 2026

The Spectral Lifecycle of Transformer Training: Transient Compression Waves, Persistent Spectral Gradients, and the Q/K--V Asymmetry

We present the first systematic study of weight matrix singular value spectra during transformer pretraining, tracking full SVD decompositions of every weight matrix at 25-step intervals across three model scales (30M--285M parameters). We discover three phenomena: (1)~Transient Compression Waves: stable rank compression propagates as a traveling wave from early to late layers, creating a dramatic gradient that peaks early then reverses -- late layers eventually over-compress past early layers. (2)~Persistent Spectral Gradients: the power-law exponent~

α

develops a permanent depth gradient forming a non-monotonic inverted-U in deeper models, with peaks shifting toward earlier layers as depth increases. (3)~Q/K--V Functional Asymmetry: value/output projections compress uniformly while query/key projections carry the full depth-dependent dynamics. The dissociation between transient compression and persistent spectral shape reveals that rank and spectral shape encode fundamentally different information about training. We formalize this as a two-timescale dynamical model and derive scaling laws (

Δα\propto L^{0.26}

R^2{=}0.99

). We validate on nine models across three families (custom, GPT-2, Pythia; 30M--1B parameters; 8--36 layers), demonstrate that

α

predicts layer importance (

ρ{=}0.69

0.84

p{<}0.02

), and show that spectral-guided pruning outperforms Last-N heuristics by

1.1{\times}

3.6{\times}

across seven models in two families (GPT-2 124M--774M, Pythia 160M--1B), with worst-vs-best gaps up to

23.7{\times}

confirming the causal role of spectral structure.

Yi Liu

Transformer Training Dynamics Transformer Neural Network Training Dynamics Spectral Methods

Ask PDF

Jun 15, 2026

Taming Curvature: Architecture Warm-Up for Stable Transformer Training

Training billion-parameter Transformers is often brittle, with transient loss spikes and divergence that waste compute. Even though the recently developed Edge of Stability (EoS) theory provides a powerful tool to understand and control the stability of optimization methods via the (preconditioned) curvature, these curvature-controlling methods are not popular in large-scale Transformer training due to the complexity of curvature estimation. To this end, we first introduce a fast online estimator of the largest (preconditioned) Hessian eigenvalue (i.e., curvature) based on a warm-started variant for power iteration with Hessian-vector products. We show theoretically, and verify empirically, that the proposed method makes per-iteration curvature tracking feasible at billion parameter scale while being more accurate. Using this tool, we find that training instabilities coincide with surges in preconditioned curvature and that curvature grows with depth. Motivated by these observations, we propose architecture warm-up: progressively growing network depth to carefully control the preconditioned Hessian and stabilize training. Experiments on large Transformers validate that our approach enables efficient curvature tracking and reduces instabilities compared to existing state-of-the-art stabilization techniques without slowing down convergence.

Sameera Ramasinghe, Ajanthan Thalaiyasingam, Hadi Mohaghegh Dolatabadi +6

Pluralis Research

Edge of Stability Transformer Second-Order Optimization Neural Network Training Dynamics Transformer Training Dynamics

Ask PDF

Abstract

Transformer

Explore similar work

Spectral Edge Dynamics of Training Trajectories: Signal--Noise Geometry Across Scales

The Spectral Lifecycle of Transformer Training: Transient Compression Waves, Persistent Spectral Gradients, and the Q/K--V Asymmetry

Taming Curvature: Architecture Warm-Up for Stable Transformer Training