勾配グループ化による学習率スケーリングを用いたLLMの制御

要旨

大規模言語モデル（LLM）の訓練は、その巨大なスケールと異種混在のアーキテクチャにより課題を抱えています。AdamWのような適応型最適化手法は勾配の変動に対処するものの、パラメータごとの学習率推定を効率的かつ効果的に行う点では依然として難があり、訓練の不安定性、収束の遅さ、そしてパラメータ効率的なファインチューニング（PEFT）技術との互換性の低さといった問題を引き起こします。本研究では、勾配グループ化によるスケーリング（SGG）を提案します。これは、動的なグループ化とグループ固有のスケーリングを通じて適応学習率推定を改善する最適化ラッパーです。SGGはまず各層の勾配統計をクラスタにグループ化し、その後クラスタ固有のスケーリングを適用して各パラメータの学習率を調整します。これにより、集団的なグループ単位の制約を課しつつ、精密なパラメータごとの適応を維持します。多様な（M）LLMベンチマークでの実験結果から、SGGは既存の最適化手法とシームレスに統合され、さまざまなモデルサイズにおいて一貫した性能向上とより速い収束を実現することが示されています。バッチサイズや学習率の変動に対する安定性から、SGGはLLM最適化における堅牢な選択肢として確立されています。

English

Training large language models (LLMs) poses challenges due to their massive scale and heterogeneous architectures. While adaptive optimizers like AdamW help address gradient variations, they still struggle with efficient and effective parameter-wise learning rate estimation, resulting in training instability, slow convergence, and poor compatibility with parameter-efficient fine-tuning (PEFT) techniques. This work introduces Scaling with Gradient Grouping (SGG), an optimizer wrapper that improves adaptive learning rate estimation by dynamic grouping and group-specific scaling. SGG first groups gradient statistics in each layer into clusters and then applies cluster-specific scaling to calibrate learning rates for each parameter, thus imposing collective group-wise constraints while maintaining precise per-parameter adaptation. Experiments on diverse (M)LLM benchmarks show that SGG integrates seamlessly with existing optimizers, and offers consistent gains and faster convergence over baselines, with various model sizes. Its stability across varying batch sizes and learning rates establishes SGG as a robust choice for LLM optimization.

勾配グループ化による学習率スケーリングを用いたLLMの制御

Taming LLMs by Scaling Learning Rates with Gradient Grouping

要旨

Support