スケール分布分離：大規模言語モデルの安定かつ効果的な学習を実現する

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

February 21, 2025

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

cs.AI

要旨

大規模言語モデル（LLM）の事前学習における訓練の安定性は、特に勾配爆発や勾配消失が起こりやすいPost-Norm Transformerなどのアーキテクチャにおいて、継続的な課題となっています。本論文では、全結合層における重み行列のスケールと分布を明示的に分離することで訓練を安定化する新しいアプローチ、Scale-Distribution Decoupling（SDD）を提案します。SDDは、活性化を調整するための正規化メカニズムと、良好な勾配状態を維持するための学習可能なスケーリングベクトルを適用し、勾配爆発と勾配消失を効果的に防ぎます。この分離により、特に深層ネットワークにおいて、安定した勾配伝播を確保することで最適化効率が向上します。実験結果は、本手法が様々なLLMアーキテクチャにおいて訓練を安定化し、異なる正規化設定において既存の手法を上回ることを示しています。さらに、提案手法は軽量で既存のフレームワークと互換性があり、LLM訓練の安定化における実用的なソリューションとなっています。コードはhttps://github.com/kaihemo/SDDで公開されています。

English

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 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.