LLM事前学習における重み再スケーリングによる分散制御

Variance Control via Weight Rescaling in LLM Pre-training

March 21, 2025

著者: Louis Owen, Abhay Kumar, Nilabhra Roy Chowdhury, Fabian Güra

cs.AI

要旨

大規模言語モデル（LLM）の事前学習の結果は、重み初期化と分散制御戦略に大きく依存します。一般的なニューラルネットワークにおいて初期分散制御の重要性は十分に文書化されていますが、特にLLMの事前学習における初期化とその成長管理に関する文献はやや乏しい状況です。本論文では、レイヤーインデックス再スケーリング（LIR）重み初期化スキームと、ターゲット分散再スケーリング（TVR）分散制御戦略を紹介します。1BパラメータのLLaMAモデルを用いた実験により、これらの技術を用いたより良い分散管理が下流タスクのパフォーマンスを大幅に向上させ（一般的な事前学習ベンチマークで最大4.6%の改善）、極端な活性化値を減少させることで、量子化や低精度学習に関連する課題を軽減することが実証されました。私たちのコードは以下で公開されています：https://github.com/bluorion-com/weight_rescaling。

English

The outcome of Large Language Model (LLM) pre-training strongly depends on weight initialization and variance control strategies. Although the importance of initial variance control has been well documented in neural networks in general, the literature on initialization and management of its growth during LLM pre-training, specifically, is somewhat sparse. In this paper, we introduce the Layer Index Rescaling (LIR) weight initialization scheme, and the Target Variance Rescaling (TVR) variance control strategy. Experiments on a 1B parameter LLaMA model demonstrate that better variance management using these techniques yields substantial improvements in downstream task performance (up to 4.6% on common pre-training benchmarks) and reduces extreme activation values, thus mitigating challenges associated with quantization and low-precision training. Our code is available at: https://github.com/bluorion-com/weight_rescaling.