Memory-Efficient LLM Training with Dynamic Sparsity: From Stability to Practical Scaling

技術要約：動的スパース性を利用したメモリ効率の高いLLM学習

問題提起
動的スパース学習（Dynamic Sparse Training: DST）は、学習を通じてスパースなトポロジーを維持することにより、ディープニューラルネットワークの効率を向上させるパラダイムを提示している。しかし、LLMに対してDSTを適用することは、最適化の不安定性のために困難であることが判明している。著者らは、標準的なAdamベースのオプティマイザをDSTに素朴に適用すると、トポロジーの更新（プルーニングおよび再成長）の直後に深刻な「損失スパイク（loss spikes）」が発生することを特定した。

根本的な原因は、「コールドスタート」効果であると特定されている。新たに再成長したパラメータには、蓄積されたオプティマイザの状態（一次および二次モーメント推定値）が欠如している。Adamにおいて、これらの新しいパラメータに対するバイアス補正項は、成熟した重みと比較して不釣り合いに大きな更新幅をもたらす。これが学習の軌道を乱し、一時的な損失スパイクを引き起こし、特にモデルスケールが大きくなる、あるいはトポロジー更新がよりアグレッシブになる場合に収束を妨げる。さらに、従来のDSTの実装は、パラメータ空間全体に対してフルサイズのオプティマイザ状態を割り当てるため、メモリ節約を実現できないことが多い（実際には一部のパラメータのみがアクティブであるにもかかわらず）。

手法：Sparse Memory-Efficient Training (SMET)
これらの課題に対処するため、著者らは、DSTを安定化させ、LLMのメモリ消費量を削減するように設計されたフレームワークであるSMETを提案する。この手法は、以下の3つのコアコンポーネントで構成される。

1. オプティマイザ状態のウォームアップ（Optimizer-State Warm-up）: コールドスタート効果を緩和するために、SMETは新たに再成長したパラメータに対して2段階のウォームアップ戦略を採用している。

   • 暗黙的ウォームアップ（タイムステップのリセット）: 再成長したパラメータの内部タイムステップカウンタ（$t_i$）をリセットする。これによりAdamのバイアス補正メカニズムが再活性化され、更新規則の分母が大きくなり、初期の更新幅が自然に抑制される。
   • 明示的ウォームアップ（学習率のスケーリング）: 再成長したパラメータに対して、学習率がゼロからターゲットレートまで$W$ステップにわたって線形に増加するローカルな線形ウォームアップを適用する。
   • これらのメカニズムが共同で更新の爆発を防ぎ、トポロジー変更直後の最適化の軌道を滑らかにする。

2. 密度依存の学習率スケーリング（Density-Aware Learning Rate Scaling）: ウォームアップによって導入される保守的な更新を補償し、スパース層の有効なファンイン（fan-in）の減少を考慮するために、SMETは学習率をスケーリングする。Maximal Update Parameterization ($\mu$P) に着想を得て、学習率は現在のスパース密度 $d$ に対して $1/\sqrt{d}$ でスケーリングされる。これにより、更新幅が密な学習（dense training）のレジームと同等に保たれる。

3. メモリ効率の高い最適化（Memory-Efficient Optimization）: SMETは、勾配とオプティマイザの状態（モーメント推定値 $m$ および $v$）をアクティブな（非ゼロの）パラメータに対してのみ保存するスパースストレージ戦略を実装している。

   • 勾配はバックプロパゲーション中にアクティブなインデックスに対してのみ保持される。
   • アクティブなインデックスやタイムステップの追跡に伴うオーバーヘッドを管理するために、パラメータをサイズ $B$ のブロックにグループ化するブロック単位のストレージ戦略を使用し、メタデータのコストを償却する。
   • このアプローチは、状態を圧縮するものの密なモデルを生成してしまう手法とは異なり、学習ライフサイクル全体を通じてスパース性を維持する。

主な貢献

1. 不安定性の分析: 本論文は、DSTを用いたLLMの事前学習における損失スパイク現象を特定・理論的に分析し、それがAdamオプティマイザにおける再成長した重みのコールドスタート効果に起因することを明らかにした。
2. 安定化フレームワーク: 著者らは、トポロジー更新を安定させ、学習性能を向上させるために、オプティマイザ状態のウォームアップと密度依存の学習率スケーリングを統合したSMETを提案している。
3. メモリ効率の高い実装: アクティブなパラメータに対してのみ勾配とオプティマイザの状態を保存するスパース最適化の実装を開発し、非構造化およびブロック単位のスパース性パターンの両方をサポートしつつ、学習を通じてスパース性を維持する。
4. 体系的な評価: 様々なモデルスケールとスパースレベルにわたるDSTの包括的な評価を行い、SMETが既存のDST手法を凌駕し、密な学習の性能に迫ることを実証している。

実験結果
C4およびOpenWebTextデータセットを用いて、LLaMAベースのモデル（60Mから1Bパラメータの範囲）で実験が行われた。

• 性能: SMETは、他のメモリ効率の高いベースライン（LoRA, ReLoRA, GaLore）および既存のDST手法（SET, RigL）を一貫して上回っている。より大きなモデルスケール（例：1Bパラメータ）において、SMETは高いスパースレベル（例：密度 $d=0.1$）であっても、密な学習に近いパープレキシティを達成している。
• 安定性: 本手法は、バニラなDSTで見られる損失スパイクを効果的に排除し、滑らかな学習曲線を実現している。
• メモリ効率: アクティブなパラメータに対してのみ状態を保存することで、SMETは勾配とオプティマイザ状態のメモリフットプリントを大幅に削減する。1Bモデルにおいて密度 $d=0.4$ の場合、大幅なメモリ節約が観察される。
• スケーラビリティ: スパース学習と密な学習の性能差はモデルサイズが大きくなるにつれて縮小しており、これはスパース学習が大規模モデルにおいて特に実行可能であることを示唆している。
• 汎用性: SMETの利点は、異なる再成長戦略（ランダムおよび勾配ベース）および異なるモデルアーキテクチャ（LLaMAおよびQwen）の両方において一般化される。

意義と主張
本論文は、SMETがLLMにとって密な学習に代わる実用的な選択肢としての道を開くものであると主張している。動的スパース学習に固有の最適化の不安定性を解決し、スパースな状態保存を通じて実際のメモリ節約を実現することで、SMETは安定、スケーラブル、かつメモリ効率の高い事前学習を可能にする。著者らは、SMETが学習の安定性を向上させオプティマイザのメモリを削減する一方で、推論の高速化を完全に実現するためには、アクティベーションメモリに関する課題や、細粒度なスパース性に対する効率的なハードウェアサポートの必要性が残っていると指摘している。それにもかかわらず、本研究は、適切な安定化とメモリ管理があれば、動的スパース性が基盤モデルへと成功裏にスケールできることを示している。