arXiv:2505.21910cs.LGcs.CV2025-05ICLR被引 7

提出无需学习率预热的训练方法,解决大模型训练崩溃问题。

Taming Transformer Without Using Learning Rate Warmup

  • 通过分析注意力矩阵谱能量集中现象,揭示模型崩溃根源。
  • 设计平滑权重更新策略,防止熵崩溃与训练失败。
  • 适用于ViT/Swin-GPT等模型,可稳定训练且无需学习率预热。

在不使用学习率预热等技巧的情况下,将Transformer模型大规模化极具挑战性,日益受到关注。本文对Transformer训练过程进行理论分析,揭示了训练中模型崩溃的根源——查询与键投影矩阵 ${W_q}^{ op} W_k$ 的谱能量集中现象,该现象导致恶性熵崩溃。为缓解此问题,受Weyl不等式启发,提出一种新型优化策略:若梯度最大奇异值与前步权重最大奇异值之比超过阈值,则自动将学习率限制为 $\frac{σ_{1}(W_{t-1})}{σ_{1}( abla W_t)}$ 的加权倍数,使连续步骤的权重更新平滑。该策略可防止谱能量过度集中于少数方向,从而避免恶性熵崩溃引发的模型崩溃。在ViT、Swin-Transformer和GPT上进行了大量实验,结果表明该方法能有效且稳定地训练这些Transformer模型,无需学习率预热。

原文摘要 · Abstract (English)

Scaling Transformer to a large scale without using some technical tricks such as learning rate warump and using an obviously lower learning rate is an extremely challenging task, and is increasingly gaining more attention. In this paper, we provide a theoretical analysis for the process of training Transformer and reveal the rationale behind the model crash phenomenon in the training process, termed \textit{spectral energy concentration} of ${\bW_q}^{\top} \bW_k$, which is the reason for a malignant entropy collapse, where ${\bW_q}$ and $\bW_k$ are the projection matrices for the query and the key in Transformer, respectively. To remedy this problem, motivated by \textit{Weyl's Inequality}, we present a novel optimization strategy, \ie, making the weight updating in successive steps smooth -- if the ratio $\frac{σ_{1}(\nabla \bW_t)}{σ_{1}(\bW_{t-1})}$ is larger than a threshold, we will automatically bound the learning rate to a weighted multiple of $\frac{σ_{1}(\bW_{t-1})}{σ_{1}(\nabla \bW_t)}$, where $\nabla \bW_t$ is the updating quantity in step $t$. Such an optimization strategy can prevent spectral energy concentration to only a few directions, and thus can avoid malignant entropy collapse which will trigger the model crash. We conduct extensive experiments using ViT, Swin-Transformer and GPT, showing that our optimization strategy can effectively and stably train these Transformers without using learning rate warmup.

Transformer训练稳定优化策略

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