Spike No More: Stabilizing the Pre-training of Large Language Models

📄 arXiv: 2312.16903v4 📥 PDF

作者: Sho Takase, Shun Kiyono, Sosuke Kobayashi, Jun Suzuki

分类: cs.CL, cs.AI

发布日期: 2023-12-28 (更新: 2025-07-25)

备注: COLM 2025

💡 一句话要点

稳定大语言模型预训练：通过控制梯度范数避免损失尖峰

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大语言模型 预训练 损失尖峰 梯度范数 雅可比矩阵

📋 核心要点

  1. 大语言模型预训练中的损失尖峰会严重影响模型性能，甚至导致训练失败，现有方法难以有效避免。
  2. 本文提出通过控制梯度范数来稳定预训练过程，核心思想是减小模型子层规模并增大shortcut连接。
  3. 实验结果表明，满足小子层和大学校正连接条件的方法能够有效抑制预训练过程中的损失尖峰。

📝 摘要（中文）

大语言模型的预训练过程中经常出现损失尖峰，这会降低模型性能，甚至破坏预训练过程。由于预训练需要大量的计算资源，因此应避免此类尖峰。本文假设损失尖峰是由梯度范数的突然增长引起的，通过分析子层雅可比矩阵的谱范数，探索保持梯度范数较小的因素。研究结果表明，稳定预训练过程需要两个条件：小的子层和大的shortcut连接。我们进行了各种实验，以实证验证我们的理论分析。实验结果表明，满足这些条件的方法可以有效地防止预训练期间的损失尖峰。

🔬 方法详解

问题定义：论文旨在解决大语言模型预训练过程中出现的损失尖峰问题。这些尖峰会导致训练不稳定，降低模型性能，甚至完全破坏预训练过程。现有的优化方法和模型结构无法有效避免这些损失尖峰，使得预训练过程耗时且不稳定。

核心思路：论文的核心思路是认为损失尖峰是由梯度范数的突然增长引起的。因此，通过控制和约束梯度范数，可以稳定预训练过程，避免损失尖峰的出现。具体来说，论文通过分析模型子层的雅可比矩阵的谱范数，来研究影响梯度范数的因素。

技术框架：论文主要通过理论分析和实验验证来支持其观点。首先，论文对Transformer模型的子层（如注意力层和前馈网络层）的雅可比矩阵进行谱范数分析，推导出梯度范数与子层规模和shortcut连接之间的关系。然后，论文设计了一系列实验，通过调整模型结构中的子层规模和shortcut连接强度，来验证理论分析的正确性。

关键创新：论文的关键创新在于将损失尖峰与梯度范数联系起来，并通过分析雅可比矩阵的谱范数，找到了影响梯度范数的关键因素：子层规模和shortcut连接。与现有方法不同，论文不是简单地调整优化器参数或使用正则化方法，而是从模型结构本身入手，通过调整子层规模和shortcut连接来稳定预训练过程。

关键设计：论文的关键设计包括：1) 对Transformer模型的子层进行雅可比矩阵的谱范数分析，推导出梯度范数与子层规模和shortcut连接之间的关系；2) 设计实验，通过调整模型结构中的子层规模（例如，减小前馈网络的隐藏层大小）和shortcut连接强度（例如，调整残差连接的权重），来验证理论分析的正确性；3) 实验中使用了标准的Transformer模型结构和预训练数据集，以便与其他方法进行比较。

📊 实验亮点

实验结果表明，减小Transformer模型子层规模并增大shortcut连接强度可以有效抑制预训练过程中的损失尖峰。例如，通过减小前馈网络的隐藏层大小，并增加残差连接的权重，可以显著降低损失尖峰的频率和幅度，从而提高模型的训练稳定性和最终性能。具体性能提升数据未知。

🎯 应用场景

该研究成果可应用于各种基于Transformer的大语言模型的预训练过程，尤其是在计算资源有限的情况下，可以帮助研究人员更稳定、高效地训练模型。通过避免损失尖峰，可以减少训练时间和计算成本，并提高模型的最终性能。该方法对于开发更大规模、更高性能的大语言模型具有重要意义。

📄 摘要（原文）

Loss spikes often occur during pre-training of large language models. The spikes degrade the performance of large language models and sometimes ruin the pre-training. Since the pre-training needs a vast computational budget, we should avoid such spikes. Based on the assumption that the loss spike is caused by the sudden growth of the gradient norm, we explore factors to keep the gradient norm small through an analysis of the spectral norms of the Jacobian matrices for the sub-layers. Our findings suggest that stabilizing the pre-training process requires two conditions: small sub-layers and large shortcut. We conduct various experiments to empirically verify our theoretical analyses. Experimental results demonstrate that methods satisfying the conditions effectively prevent loss spikes during pre-training.