技術要約：トークン幾何学とEmberオプティマイザ

問題提起

現在の大規模言語モデル（LLM）は、埋め込みテーブルとLMヘッド行列（総称して「トークン・インターフェース」）をAdamオプティマイザを用いて最適化しています。Adamは優れた性能を発揮しますが、分散学習において重大なエンジニアリング上の負担を強いています。具体的には、Adamはすべてのパラメータに対して第1次および第2次のモーメント・バッファを保持するため、語彙サイズ $V$ と埋め込み次元 $D$ に対して $O(2VD)$ のオプティマイザ状態の複雑さをもたらします。大規模モデルにおいて、この状態は数ギガバイトのVRAMを消費し、モデルの計算の主たるボトルネックではないパラメータに対しても、複雑なシャーディング戦略（ZeRO/FSDPなど）を必要とさせます。

さらに、本論文は、トークン・インターフェースへの標準的なAdamの適用が、これらのパラメータの特定の幾何学的性質を無視していると主張しています。高密度な線形層とは異なり、トークン・インターフェースは疎な勾配（ステップごとにアクティブなトークンの行のみが更新される）を介して更新されます。また、クロスエントロピーにおける基礎となる損失関数は、標準的なオプティマイザでは活用しきれない特有の性質を持っている可能性があります。

手法：Emberオプティマイザ

著者らは、埋め込みおよびLMヘッド行列のために設計された軽量なオプティマイザであるEmberを導入しています。その核心となる手法は、情報幾何学、特に**フィッシャー情報量（Fisher Information metric）**に基づいています。

理論的基礎

本論文は、クロスエントロピー損失において、フィッシャー計量（勾配の外積の期待値）がカルバック・ライブラー（KL）ダイバージェンスの正確な2次項となることを提示しています。最適解の近傍において、フィッシャー計量はヘシアン（Hessian）を近似します。

• 平方根フィッシャー（Square-Root Fisher）： フル・ニュートン・ステップ（ヘシアンの逆行列計算）を実行することは計算コストが高く不安定であるため、著者らは、勾配をフィッシャー計量の平方根で条件付けすることが十分であると主張しています。これは、実質的に勾配を標準偏差によって正規化する（「zスコア」変換を行う）ことに相当します。
• 勾配の疎性（Gradient Sparsity）： トークン・インターフェースにおいて、トークン $i$ の勾配は、そのトークンが選択されたときのみ非ゼロになります。本論文では、行ごとの勾配エネルギーが、トークンの参加確率（$p_i$）の不偏推定量であると観察されています。

アルゴリズム設計

Emberは、第2次モーメントの推定値を因子分解することで、Adamの $O(2VD)$のメモリ・フットプリントを$O(V + D)$ へと置き換えます。

1. 第1次モーメントの除去： 著者らは、Adamの第1次モーメント（モーメンタム）の指数移動平均（EMA）を完全に排除することが、パフォーマンスの向上とメモリ節約の両方に寄与することを発見しました。
2. 因子分解された第2次モーメント： $V \times D$ の第2次モーメント・バッファを保持する代わりに、Emberは2つの1次元ベクトルを保持します。
   • $r_t$: 行方向の第2次モーメント推定値（トークンの頻度/エネルギーを追跡）。
   • $c_t$: 列方向の第2次モーメント推定値（特徴量ごとのエネルギーを追跡）。
3. 前処理（Preconditioning）： 更新ルールは、軽量な外積前処理器を形成します。パラメータ $(i, j)$ に対する実効的な分母は、行と列の推定値の幾何平均（$\sqrt{\hat{r}_i \hat{c}_j}$）から導出されます。
4. バイアス補正： 前処理器を形成する前に、行および列の推定値に対して標準的なバイアス補正が適用されます。

結果として得られる更新ルールは、勾配をフィッシャー計量の平方根によって除算するものであり、これはトークンの出現頻度に依存しないように正規化されつつ、頻度の高いトークンの分散を低減します。

主な貢献

1. メモリ効率： Emberは、トークン・インターフェースにおけるオプティマイザ状態を $O(2VD)$から$O(V + D)$ に削減します。これにより、埋め込みのオプティマイザ状態をシャーディングする必要がなくなり、巨大なモデルであっても、その全状態を単一のGPUに収めることが可能になります。
2. 幾何学的洞察： 本論文は、Transformer向けの競争力のあるオプティマイザ（Adam、Lion、Muonを含む）が、勾配をフィッシャー計量の平方根（またはその行列類似体）で条件付けるという枠組みの下で統一できるという、原理的な視点を提供しています。
3. 最適化軌道の分析： トークンの軌跡に対する特異値分解（SVD）を通じて、著者らは、埋め込みトークンが驚くほど単純な**1Dのレイ（単項基底）**に沿って最適化されることを示しています。これは、ニューラルネットワークの最適化が高度に非凸でカオス的なランドスケープをナビゲートするという従来の考え方に疑問を投げかけるものであり、代わりに、トークン・インターフェースにおいては滑らかで線形な収束パスが存在することを示唆しています。
4. オープンソースの実装： 既存のZeRO/FSDPセットアップと互換性のある分散実装が公開されており、即時の導入を容易にしています。

結果

著者らは、複数のレジームにおいて経験的な証拠を提示しています。

• 教師あり微調整（SFT）： GPT-2およびPythiaモデルにおいて、EmberはAdamWと同等またはわずかに優れた検証損失を実現しつつ、オプティマイザ状態を数ギガバイトからキロバイト単位へと削減しました（例：Pythia-2.8Bにおいて2.06 GB $\to$ 413 KB）。
• 強化学習（RL）： GRPOを用いたQwen2.5-3Bのポストトレーニングにおいて、Emberは様々なロールアウト・バッチサイズにわたってAdamWと同等の性能を達成しながら、4,000倍少ないオプティマイザ状態を使用しました。
• 事前学習： FineWebにおいて、EmberはAdamWと同等の性能を示し、特に低バッチ・レジームにおいてAdafactorを上回りました。これは、ランク1の勾配分解が、疎またはノイズの多い勾配に対する安定性を提供するためです。
• スケーリング則： Emberはモデルサイズおよびバッチサイズに対して効果的にスケールします。EmberとAdamWの性能差はモデルサイズが大きくなるにつれて消失しており、これはAdamのフルな第2次モーメント・バッファが、Emberによるランク1の因子分解によって十分に近似可能であることを示唆しています。
• 安定性： Emberは、AdamWで見られる失敗モードである、小規模バッチでの学習初期の損失スパイクを示しません。

意義と主張

本論文は、Emberがメモリ制約のある学習において、SOTA（State-of-the-Art）の性能を達成しながらメモリ・オーバーヘッドを劇的に削減するという「パレート・フロンティア」の改善を提供すると主張しています。その意義は二重です。

1. 実用的なエンジニアリング： 埋め込みのオプティマイザ状態をシャーディングする必要性を排除することで、Emberは分散学習の設定を簡素化し、Adamの状態で以前は制約となっていた最大7Bパラメータ（およびそれ以上）のモデルであっても、単一デバイスでの高速なイテレーションを可能にします。
2. 理論的理解： トークンの最適化軌道がクリーンな1Dの単項基底に分解されるという発見は、「非凸」な性質が、トークン・インターフェースの固有の性質ではなく、オプティマイザやパラメータ化の副産物である可能性を示唆しています。これは、Transformer向けの最適なオプティマイザの探索空間が、これまで考えられていたよりも狭く、平方根フィッシャー計量の周辺に集中していることを意味しています。

著者らは、Emberがメモリ制約のある学習のための実用的なツールである一方で、その理論的枠組みはLLMがどのように学習するかを理解するための新しいレンズを提供し、特にトークン表現における最適化パスの単純さを浮き彫りにしていると結論付けています。