研究者らは、Mixture-of-Experts (MoE) モデル用に設計されたオプティマイザであるSkewAdamを発表した。これは異なるパラメータ群に異なる状態タイプを割り当てることで、メモリ使用量を劇的に削減する。

バックボーンにはfloat32モーメンタムと因数分解された第2モーメントを、エキスパートには因数分解された第2モーメントを、ルーターには正確な第2モーメントを割り当てることで、6.78Bパラメータのモデルにおいてオプティマイザ状態が50.6 GBから1.29 GBに削減される。

  • トレーニング中のピークメモリ使用量は81.4 GBから31.3 GBに減少し、40 GBのアクセラレータ予算内に収まる。
  • SkewAdamは82Mトークンで検証パープレキシティ108.4を達成し、AdamW (126.8)、Muon (120.2)、Lion (393.7) を上回る性能を示した。
  • このメモリ削減は精度の低下なしに実現されており、パフォーマンスの向上は階層化割り当てそのものではなく、モーメンタムの保持によるものである。

これらの結果は、オプティマイザ状態の割り当て先がその量と同様に重要であることを示唆しており、限られたメモリを持つハードウェアでの効率的なトレーニングを可能にする。