Residual Koopman Spectral Profiling for Predicting and Preventing Transformer Training Instability

📄 arXiv: 2602.22988 📥 PDF

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

分类: cs.LG, cs.AI

发布日期: 2026-02-28

💡 一句话要点

提出残差Koopman谱分析方法，用于预测和预防Transformer训练中的不稳定性

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: Transformer 训练稳定性 Koopman谱分析 动态模式分解 残差网络 模型初始化 谱整形

📋 核心要点

  1. Transformer训练不稳定导致计算资源浪费，现有方法难以在训练前预测和预防。
  2. 论文提出残差Koopman谱分析（RKSP），通过分析层间残差的Koopman谱特征来估计训练发散的风险。
  3. 实验表明，RKSP能有效预测发散，AUROC达0.995，并可通过Koopman谱整形（KSS）降低发散率，提升学习率。

📝 摘要（中文）

Transformer中的训练发散浪费计算资源，但从业者只有在昂贵的训练开始后才能发现不稳定性。因此，他们需要在训练开始前获得Transformer失败的预期概率。我们对残差Koopman谱分析（RKSP）的研究提供了这样的估计。通过在初始化时进行一次前向传递，RKSP通过将白化的动态模式分解应用于逐层残差快照来提取Koopman谱特征。我们的核心诊断指标，即近单位谱质量，量化了集中在单位圆附近的模式的比例，从而捕捉了不稳定性风险。为了预测各种配置中的发散，该估计器实现了0.995的AUROC，优于最佳梯度基线。我们进一步通过Koopman谱整形（KSS）使该诊断可操作，KSS在训练期间重塑谱。我们通过实验验证了我们的方法在实践中有效：RKSP预测初始化时的发散，并且当RKSP标记高风险时，开启KSS成功地防止了发散。在没有归一化层的具有挑战性的高学习率方案中，KSS将发散率从66.7%降低到12.5%，并实现了高50%到150%的学习率。这些发现推广到WikiText-103语言建模、CIFAR-10上的视觉Transformer以及预训练语言模型，包括高达7B的GPT-2和LLaMA-2，以及新兴架构，如MoE、Mamba风格的SSM和KAN。

🔬 方法详解

问题定义：Transformer模型在训练过程中容易出现不稳定性，导致训练发散，浪费大量计算资源。现有的方法通常依赖于训练过程中的梯度信息来诊断不稳定性，但这些方法无法在训练开始前预测潜在的风险，也无法有效地预防训练发散。

核心思路：论文的核心思路是利用Koopman算子的谱分析来捕捉Transformer模型在初始化状态下的动态特性，从而预测训练过程中的不稳定性。通过分析层间残差的Koopman谱，可以识别出接近单位圆的模式，这些模式与训练发散的风险密切相关。

技术框架：RKSP方法主要包含以下几个阶段：1) 在初始化状态下，对Transformer模型进行一次前向传递，记录每一层的残差；2) 对层间残差进行白化处理，以消除数据中的冗余信息；3) 使用动态模式分解（DMD）方法提取Koopman谱特征；4) 计算近单位谱质量，作为衡量不稳定性风险的指标；5) 如果RKSP检测到高风险，则启用Koopman谱整形（KSS）方法，在训练过程中重塑谱，以防止发散。

关键创新：RKSP的关键创新在于将Koopman算子的谱分析应用于Transformer模型的训练稳定性预测。与传统的基于梯度的方法相比，RKSP可以在训练开始前预测潜在的风险，并且能够通过KSS方法主动地防止训练发散。此外，RKSP方法具有通用性，可以应用于各种Transformer架构，包括新兴的MoE、Mamba风格的SSM和KAN等。

关键设计：RKSP的关键设计包括：1) 使用白化的动态模式分解（DMD）方法提取Koopman谱特征，以提高谱分析的准确性；2) 定义近单位谱质量作为衡量不稳定性风险的指标，该指标能够有效地捕捉接近单位圆的模式；3) 提出Koopman谱整形（KSS）方法，通过调整谱的形状来防止训练发散。KSS的具体实现细节（如调整哪些频率分量，调整的幅度等）可能需要根据具体的模型和数据集进行调整。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

RKSP在预测Transformer训练发散方面表现出色，AUROC达到0.995，优于现有梯度基线方法。在具有挑战性的高学习率场景下，KSS将发散率从66.7%降低到12.5%，并允许使用高50%到150%的学习率。该方法成功应用于WikiText-103、CIFAR-10以及包括GPT-2和LLaMA-2在内的多个模型。

🎯 应用场景

该研究成果可广泛应用于Transformer模型的训练优化，尤其是在资源受限或需要快速迭代的场景下。通过RKSP，开发者可以在训练前评估模型的稳定性，避免不必要的计算资源浪费。KSS则能有效提高训练的稳定性和效率，支持更大的学习率，加速模型收敛。该方法对预训练语言模型和新兴架构具有重要意义。

📄 摘要（原文）

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.