arXiv:2602.22988cs.LGcs.AI2026-02

通过初始化快照预测Transformer训练崩溃,提前防范算力浪费。

Residual Koopman Spectral Profiling for Predicting and Preventing Transformer Training Instability

  • 从初始化前向传播中提取残差谱特征,量化接近单位圆的模式比例。
  • 预测准确率AUROC达0.995,显著优于梯度基线方法。
  • 可指导动态调整训练策略,适用于大模型与新架构。

Transformer训练发散会浪费大量算力,但问题常在训练开始后才被发现。为此,本文提出残差Koopman谱分析(RKSP),仅需一次初始化前向传播,通过白化动态模态分解提取逐层残差快照的Koopman谱特征。核心诊断指标——近单位谱质量,衡量集中在单位圆附近的模式占比,有效反映训练不稳定性风险。该估计器在多种配置下预测发散的AUROC达0.995,优于最优梯度基线。进一步提出谱重塑方法(KSS),可在训练中动态调整谱分布。实验证明:初始阶段即可准确预测发散;当RKSP预警高风险时开启KSS,可成功防止崩溃。在无归一化层的高学习率场景下,发散率由66.7%降至12.5%,支持学习率提升50%至150%。该方法在WikiText-103语言建模、CIFAR-10视觉变压器、GPT-2及LLaMA-2(最大7B)以及MoE、Mamba-style SSMs、KAN等新兴架构上均表现良好。

原文摘要 · Abstract (English)

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.

Transformer训练稳定谱分析预防崩溃

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。