Gefen: Optimized Stochastic Optimizer
本論文は、自動的な二次のモーメントの共有と学習による一次のモーメントの量子化を通じて、AdamWのメモリフットプリントを約8分の1に削減し、同等の性能を維持しながら、より大きなバッチサイズとスループットの向上を可能にするメモリ効率の高いオプティマイザであるGefenを導入する。
原論文は CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) でライセンスされています。 これは以下の論文のAI生成解説です。著者が執筆または承認したものではありません。技術的な正確性については原論文を参照してください。 免責事項の全文を読む
巨大なロボット（ディープラーニングモデル）に新しいスキルを学習させる訓練をしていると想像してみてください。これを行うために、ロボットには、毎回の練習セッションの後に内部設定をどのように調整すべきかを教える「コーチ」（オプティマイザ）が必要です。現在、最も普及しているコーチはAdamWと呼ばれています。
AdamWは優れたコーチですが、重いバックパックを背負っています。ロボットのあらゆる設定に対して、AdamWは過去のミスや変化の速さを記憶するための2つの追加のノート（「第1次および第2次モーメント状態」と呼ばれます）を運びます。もしあなたのロボットに10億の設定がある場合、AdamWのバックパックはさらに20億の設定分の重さを加えます。このバックパックがあまりに重いため、標準的なトレーニング用コンピュータに巨大なロボットを載せることができなかったり、ロボットを非常に小さく、ゆっくりとしたステップで練習させなければならなくなったりします。
この論文では、Gefenという新しいコーチを紹介しています。Gefenは「メモリ効率の良い」オプティマイザであり、学習の速度と品質をAdamWと全く同じに保ったまま、その重いバックパックを8分の1に縮小します。
Gefenがどのように行っているのか、簡単な比喩を用いて説明します。
1. 「グループハグ」戦略（ノートの共有）
問題点: AdamWは、ロボットの脳内にある一つひとつの数字に対して、個別のメモを書いています。
Gefenの解決策: Gefenは、これらの数字の多くが実は「親友」であることを気づいています。彼らは世界に対して非常に似たような反応を示します。これら一つひとつの友人のために個別のメモを書く代わりに、Gef其は彼らをチーム（ブロック）としてグループ化し、チーム全体に一つの共有ノートを与えます。
- どのようにしてグループ化するか？ 通常、誰が誰と友達であるかを知るには、非常に複雑なマップ（ヘシアン行列と呼ばれます）が必要ですが、そのマップは巨大なロボットに対しては大きすぎて描けません。
- Gefenのトリック: Gefenは、ロボットの最初の練習セッションを見ます。どの数字が似たようなパターンで一緒に動いたかを確認します。そして、「もし最初に一緒に動いたのなら、おそらく後でも一緒に動くだろう」と仮定します。こうして、これらの数字を自動的にグループ化します。どのグループを作るかを人間に教える必要はなく、自分自身で判断します。
2. 「スケッチ画家」戦略（量子化）
問題点: たとえ共有ノートを使ったとしても、その中の数字は高い精度（例えば小数点以下10桁など）で書かれているため、スペースを占有します。
Gefenの解決策: Gefenは「スケッチ画家」のアプローチを採用しています。数値を正確に書き留める（例：0.123456789）代わりに、あらかじめ用意されたリスト（コードブック）から最も近い「標準的な値」に丸めます。
- スマートなリスト: 多くのオプティマイザは、汎用的で固定されたリスト（固定された目盛りのある定規のようなもの）を使用します。しかし、Gefenはコーチングしている特定のロボットを観察し、そのロボットに最適な目盛りのリストを学習します。誤差を最小限に抑えるために、動的計画法（Dynamic Programming）というスマートな数学的手法を用いて、データに完璧にフィットするカスタム定規を作成します。
- 安定性: このリストを最初に学習した後は、ずっと同じリストを使用し続けます。毎日定規を引き直す必要はないため、時間を節約でき、安定性も保たれます。
結果：より軽いバックパック、同じパフォーマンス
著者らは、小さな画像分類器から大規模言語モデル（Llama 3など）まで、さまざまなモデルでGefenをテストしました。
- メモリ: Gefenは、オプティマイザに必要なメモリをAdamWと比較して約8倍削減しました。130億パラメータのモデルの場合、メモリを97 GBからわずか1.5 GBに節約できます。
- 速度: バックパックが軽くなったことで、ロボットは一度に大きな「マイクロバッチ」（練習グループ）を運ぶことができます。複数のコンピュータ（FSDPおよびDDP）を使用したテストでは、コンピュータがメモリの解放を待つ必要がなくなったため、トレーニングが56%高速化されました。
- 品質: 高度な圧縮と共有を行っているにもかかわらず、ロボットは重いAdamWのバックパックを使用していた時と同じように学習しました。最終的なパフォーマンス（精度や損失）は同一でした。
なぜこれが重要なのか
これは、旅行の荷造りに例えることができます。AdamWは、靴下の一足一足に個別の箱を用意してスーツケースに詰めるようなものです。Gefenは、すべての靴下を一つのタイトな束にまとめ、一つの小さなポーチに入れることができると気づくようなものです。膨大なスペースを節約できますが、それでも靴下はすべて揃っており、目的地に同じ速さで到達することができます。
論文は、Gefenが「ドロップイン・リプレイスメント（そのまま置き換え可能なもの）」であると主張しています。つまり、設定を変更することなく既存のトレーニングコードに差し替えるだけで、即座にメモリを節約し、トレーニングを高速化できるということです。
自分の分野の論文に埋もれていませんか？
研究キーワードに一致する最新の論文のダイジェストを毎日受け取りましょう——技術要約付き、あなたの言語で。