Residual Koopman Spectral Profiling for Predicting and Preventing Transformer Training Instability

📄 arXiv: 2602.22988v1 📥 PDF

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

分类: cs.LG, cs.AI

发布日期: 2026-02-26

备注: 23 pages, 7 figures

💡 一句话要点

提出残差Koopman谱分析(RKSP)以预测和预防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模型的训练过程中，由于各种原因（如学习率过高、模型结构不稳定等）容易出现训练发散的问题，导致训练失败，浪费大量计算资源。现有的方法往往需要在训练过程中才能发现这些问题，无法提前预知并采取措施。因此，如何在训练开始前预测Transformer模型的训练稳定性，并采取相应的措施来避免训练发散，是一个重要的研究问题。

核心思路：论文的核心思路是利用Koopman算子的谱分析来捕捉Transformer模型在初始化状态下的动态特性，从而预测其训练稳定性。具体来说，通过对模型各层的残差连接进行动态模式分解，提取Koopman谱特征，并利用这些特征来量化模型的不稳定性风险。这种方法的核心在于，模型的初始化状态包含了其内在的动态特性，而这些特性与模型的训练稳定性密切相关。

技术框架：RKSP方法的整体框架包括以下几个主要步骤：1) 在模型初始化后，进行一次前向传递，记录各层的残差连接的输出；2) 对各层的残差快照应用白化的动态模式分解，提取Koopman谱特征；3) 计算近单位谱质量，作为模型不稳定性的指标；4) 如果RKSP预测模型存在高风险，则启动Koopman谱整形(KSS)，在训练过程中重塑谱，以提高训练稳定性。

关键创新：该论文的关键创新在于：1) 提出了一种新的基于Koopman谱分析的Transformer训练稳定性预测方法，能够在训练开始前预测模型的不稳定性风险；2) 提出了一种Koopman谱整形(KSS)方法，能够在训练过程中动态调整模型的谱特性，从而提高训练稳定性；3) 将该方法成功应用于多种Transformer模型和任务，包括语言建模、视觉Transformer和预训练语言模型。

关键设计：RKSP的关键设计包括：1) 使用白化的动态模式分解来提取Koopman谱特征，以提高特征的鲁棒性；2) 使用近单位谱质量作为模型不稳定性的指标，该指标能够有效地捕捉模型中存在的接近单位圆的模式，这些模式与模型的不稳定性密切相关；3) KSS通过调整模型的权重矩阵来实现谱整形，具体的调整策略需要根据模型的具体结构和任务进行设计。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

RKSP在预测Transformer训练发散方面表现出色，AUROC达到0.995，优于现有梯度基线方法。在具有挑战性的高学习率场景下，KSS能将发散率从66.7%降至12.5%，并允许使用更高的学习率（提升50%-150%）。该方法成功应用于多种模型（GPT-2, LLaMA-2, MoE, Mamba, KAN）和任务（语言建模，图像分类）。

🎯 应用场景

该研究成果可广泛应用于Transformer模型的训练和优化，尤其是在资源受限或需要快速迭代的场景下。通过提前预测训练风险并采取相应措施，可以显著减少计算资源的浪费，加速模型开发周期。此外，该方法还可以用于模型架构搜索和超参数优化，帮助研究人员找到更稳定、更高效的Transformer模型。

📄 摘要（原文）

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.