Scale-Distribution Decoupling: Enabling Stable and Effective Training of Large Language Models

📄 arXiv: 2502.15499v2 📥 PDF

作者: Ya Wang, Zhijian Zhuo, Yutao Zeng, Xun Zhou, Jian Yang, Xiaoqing Li

分类: cs.CL

发布日期: 2025-02-21 (更新: 2025-02-25)

🔗 代码/项目: GITHUB

💡 一句话要点

提出尺度-分布解耦方法，稳定高效地训练大型语言模型

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

关键词: 大型语言模型 训练稳定性 梯度爆炸 梯度消失 尺度-分布解耦 预训练 Transformer

📋 核心要点

  1. 大型语言模型训练中，梯度爆炸和消失是常见问题，尤其是在Post-Norm Transformer结构中。
  2. 论文提出尺度-分布解耦（SDD）方法，通过解耦权重矩阵的尺度和分布来稳定训练过程。
  3. 实验表明，SDD方法在多种LLM架构上稳定了训练，并在不同归一化配置下优于现有技术。

📝 摘要（中文）

大型语言模型（LLM）的预训练面临着训练稳定性这一长期挑战，特别是对于Post-Norm Transformer等架构，它们容易出现梯度爆炸和梯度消失。本文提出了一种新的方法，即尺度-分布解耦（SDD），通过显式地解耦全连接层中权重矩阵的尺度和分布来稳定训练。SDD应用了一种归一化机制来调节激活值，并使用一个可学习的缩放向量来维持良好条件化的梯度，从而有效地防止梯度爆炸和梯度消失。这种分离提高了优化效率，尤其是在深度网络中，通过确保稳定的梯度传播。实验结果表明，我们的方法稳定了各种LLM架构的训练，并在不同的归一化配置中优于现有的技术。此外，该方法轻量级且与现有框架兼容，使其成为稳定LLM训练的实用解决方案。

🔬 方法详解

问题定义：大型语言模型（LLM）的训练，特别是基于Transformer的LLM，经常遇到训练不稳定的问题，例如梯度爆炸和梯度消失。现有的归一化方法虽然可以在一定程度上缓解这些问题，但仍然存在优化效率不高，或者对特定架构不适用的情况。因此，如何设计一种通用的、高效的训练稳定方法是亟待解决的问题。

核心思路：论文的核心思路是将全连接层中权重矩阵的尺度和分布进行解耦。具体来说，就是将权重矩阵的归一化和缩放这两个过程分开处理。通过对权重矩阵进行归一化，可以控制其分布，避免梯度爆炸；通过引入可学习的缩放向量，可以维持良好条件化的梯度，防止梯度消失。这种解耦的设计使得模型能够更有效地学习，从而提高训练的稳定性。

技术框架：SDD方法主要包含两个关键模块：归一化模块和缩放模块。首先，归一化模块对全连接层的权重矩阵进行归一化处理，使其具有稳定的分布。然后，缩放模块引入一个可学习的缩放向量，用于调整归一化后的权重矩阵的尺度。这两个模块共同作用，实现了权重矩阵尺度和分布的解耦。整个流程可以嵌入到现有的LLM训练框架中，无需对整体架构进行大幅修改。

关键创新：SDD方法最重要的创新点在于其尺度-分布解耦的思想。与传统的归一化方法不同，SDD将权重矩阵的归一化和缩放过程显式地分离，从而可以更精细地控制权重矩阵的尺度和分布。这种解耦的设计使得模型能够更好地适应不同的训练场景，从而提高训练的稳定性和效率。

关键设计：SDD方法中的关键设计包括：(1) 归一化模块采用了一种标准化的归一化方法，例如Layer Normalization或RMSNorm，以确保权重矩阵具有稳定的分布。(2) 缩放模块引入了一个可学习的缩放向量，该向量与归一化后的权重矩阵进行逐元素相乘，以调整其尺度。缩放向量的初始化和学习率需要仔细调整，以避免梯度爆炸或消失。(3) SDD方法可以与不同的归一化方法结合使用，例如Layer Normalization、RMSNorm等，从而具有良好的通用性。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，SDD方法在多种LLM架构上都能够稳定训练，并且在不同的归一化配置下优于现有的技术。具体来说，SDD方法可以显著降低训练过程中的梯度方差，从而避免梯度爆炸和梯度消失。此外，SDD方法还能够提高模型的收敛速度和最终性能，例如在perplexity指标上取得明显的提升。

🎯 应用场景

该研究成果可广泛应用于各种大型语言模型的预训练和微调任务中，尤其是在模型规模较大、训练数据复杂的情况下，可以有效提高训练的稳定性和效率。此外，该方法还可以应用于其他深度学习模型，例如计算机视觉和语音识别等领域，具有重要的实际应用价值和广阔的未来发展前景。

📄 摘要（原文）

Training stability is a persistent challenge in the pre-training of large language models (LLMs), particularly for architectures such as Post-Norm Transformers, which are prone to gradient explosion and dissipation. In this paper, we propose Scale-Distribution Decoupling (SDD), a novel approach that stabilizes training by explicitly decoupling the scale and distribution of the weight matrix in fully-connected layers. SDD applies a normalization mechanism to regulate activations and a learnable scaling vector to maintain well-conditioned gradients, effectively preventing $\textbf{gradient explosion and dissipation}$. This separation improves optimization efficiency, particularly in deep networks, by ensuring stable gradient propagation. Experimental results demonstrate that our method stabilizes training across various LLM architectures and outperforms existing techniques in different normalization configurations. Furthermore, the proposed method is lightweight and compatible with existing frameworks, making it a practical solution for stabilizing LLM training. Code is available at https://github.com/kaihemo/SDD.