SkewAdam：MoE 学習のメモリ効率を高める階層型オプティマイザ

導入

Mixture-of-Experts（MoE）は、すべてのパラメータを毎回使わずにモデル規模を広げられる手法として注目されている。しかし、学習時のメモリ消費は依然として重い。重み、勾配、アクティベーションだけでなく、オプティマイザ状態が大きな割合を占めるためだ。

論文「Where Should Optimizer State Live?」が提案する SkewAdam は、この問題を MoE の構造から見直す。報告された 6.78B パラメータの MoE 言語モデルでは、bfloat16 の重みは 12.6GB だが、AdamW が保持する一次・二次モーメントは 50.6GB に達する。つまり、オプティマイザ状態こそが最大級のメモリ項目になっている。

主要ポイント

• MoE の部位ごとに状態を変える。 SkewAdam は、約 5% のパラメータを占める dense backbone には float32 momentum と分解型の二次モーメントを保持する。約 95% を占める experts には分解型二次モーメントのみを使い、0.01% 未満の router には正確な二次モーメントを保持する。
• メモリ削減幅が大きい。 オプティマイザ状態は AdamW の 50.6GB から 1.29GB へ減少し、比率では 2.6% になる。ピーク学習メモリも 81.4GB から 31.3GB に下がり、40GB クラスのアクセラレータ予算内に収まる。
• 検証性能も維持されている。 同一初期化、82M tokens の比較では、SkewAdam の検証 perplexity は 108.4。AdamW の 126.8、Muon の 120.2、Lion の 393.7 を上回る結果として報告されている。
• 効いている要素は分けて見る必要がある。 論文では、階層化は主にメモリを節約し、精度面では momentum を維持することが重要だと分析している。Adafactor は関連する分解推定を使うが momentum を持たず、性能差が残った。

意義と影響

この研究の面白さは、単なる圧縮テクニックではなく、MoE の不均一な構造に合わせてオプティマイザ状態を配置する点にある。MoE では experts がパラメータの大半を占める一方、router は極小だが負荷分散に重要であり、backbone もモデル全体の基盤を担う。これらを同じ設定で処理するのは、必ずしも合理的ではない。

SkewAdam は「状態をどれだけ持つか」だけでなく、「どこに、どの種類の状態を持つか」という問いを前面に出した。もし今後、より大きなモデルや長い学習、異なる MoE 構成でも確認されれば、限られた GPU メモリで MoE を訓練するための実用的な方向性になり得る。

ただし、素材で示されている結果は特定の実験条件に基づく。一般化には追加検証が必要だ。それでも、MoE 学習のメモリ最適化では、オプティマイザ状態の配置設計が重要な研究対象であることを明確に示している。

出典：Hugging Face Daily Papers