🤖 machine learning

MESH: Memory-Efficient Sinkhorn Optimization for Mixture-of-Experts Training

本論文は、ステートレスなSinkhorn手法を失敗させる原因となるルーティングされたエキスパート勾配の時次的不安定性に対処することで、Mixture-of-Expertsモデルの低メモリ学習を成功させるために、隠れ運動量更新とオプションのブロック・プリコンディショニングを組み込んだメモリ効率の高いSinkhorn最適化手法であるMESHを導入するものである。

原著者： Masato Fujitake

公開日 2026-08-06

📖 1 分で読めます☕ さくっと読める

CC BY 4.0

原著者： Masato Fujitake

原論文は CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) でライセンスされています。 ✨ これは以下の論文のAI生成解説です。著者が執筆または承認したものではありません。技術的な正確性については原論文を参照してください。免責事項の全文を読む

あなたは、巨大で非常に賢いロボットに物語の書き方を教えようとしているところだと想像してください。このロボットを本当に優れたものにするためには、何十億もの小さなスイッチで作られた巨大な脳が必要です。しかし、ここには落とし穴があります。ロボットの脳があまりに巨大すぎて、コンピュータのメモリにほとんど収まりきらないのです。それはまるで、小さなバッテリーを搭載したスマートフォンで、超大作映画を再生しようとしているようなものです。スマートフォンは熱くなり、動作が遅くなり、クラッシュしてしまうかもしれません。

これを解決するために、科学者たちは「オプティマイザ」と呼ばれる特別な「トレーナー（訓練士）」を使用します。最も有名なトレーナーである「AdamW」は、非常に丁寧なコーチのような存在です。彼は、ロボットの脳にある一つ一つのスイッチに対して詳細なノートをつけ、次にどのように微調整すべきかを正確に記憶しています。これは素晴らしい働きをしますが、そのノートは膨大なスペースを占有します。そこで研究者たちは、より軽量なトレーナーである「Sinkhorn」を考案しました。Sinkhornは、ノートを一切持たないコーチのような存在です。彼は現在のミスだけを見て、記憶に頼らない素早い調整を行います。これにより、メモリ消費量を大幅に節沢できますが、結果として、「Mixture-of-Experts（MoE）」と呼ばれる特定のタイプのロボットの脳に対しては、この記憶を持たないコーチは混乱してしまい、ロボットがうまく学習できなくなることが判明しました。科学者たちが問いかけているのは、「記憶を持たないコーチのメリットを維持しつつ、ロボットの学習能力を失わない方法はないか？」ということです。

この論文は、まさにこの問題を解決するための、MESH（Memory-Efficient Sinkhorn for Experts）と呼ばれる新しい手法を紹介しています。研究者たちは、素早いコーチが失敗する理由を、「エキスパート（専門家）」が常に働いているわけではないからだと突き止めました。Mixture-of-Expertsモデルでは、ロボットは各文章を処理するために、特定の数個のエキスパートだけを選びます。これは、特定の注文に対して、数人のシェフだけがキッチンに呼ばれるレストランのようなものです。シェフは注文ごとに変わるため、素早いコーチは、バラバラでノイズの多い信号を受け取り、迷子になってしまうのです。

論文は、解決策はコーチに再びフルセットのノートを与えることではなく、「隠れた記憶」を与えることだと提案しています。コーチは単に現在の注文を見るのではなく、決定を下す前に、直近の数回の注文の「平均」を記憶するのです。著者はこれを「隠れた慣性（hidden momentum）」と呼んでいます。彼らはこれを1億1000万パラメータの小さなモデル（「ナノホエール」）でテストし、このシンプルなトリック――調整を行う前にノイズを滑らかにすること――によって、ロボットの学習能力が回復することを発見しました。

以下は、彼らが発見したこと、および否定したことです：

主な解決策： 論文は、失敗の原因が「ルーテッド・エキスパート（時々しか働かないシェフ）」の信号を、記憶を持たない調整を行う前に時間的に滑らかにする必要があることにあると示しています。隠れたバッファ（恒久的なメモリとしてはカウントされないもの）を使用して最近の履歴を平均化することで、新しいMESHメソッドははるかに優れた成果を上げます。
うまくいかなかったこと： 研究者たちは、問題を解決するために多くの他のアイデアを試しましたが、それらを否定しました。単に調整を小さくしたり、エキスパートのグループ化の方法を変えたり、あるいは「符号（エラーが増えたか減ったかという情報のみ）」だけを使用したりしても、問題は解決しないことがわかりました。鍵となったのは、単なる追加データではなく、具体的には「時間の平滑化（time-smoothing）」でした。
どの程度確信しているか？： 結果は、特定の小型モデルに基づいています。これらのテストにおいて、新しい手法はオプティマイザの状態に必要なメモリを62.5%削減し（0.883GBから0.331GBへ減少）、ピーク時のコンピュータメモリ使用量を約12.6%低下させました。しかし、ロボットの最終的なテストスコア（評価損失）は、重いノートを持つコーチ（AdamW）よりもわずかに劣っており、新しい手法のスコアは約3.64、旧手法は約3.59でした。著者は、この小さな差は、ロボットの脳の他の部分（語彙など）がいまだに重いコーチを必要としているためである可能性を示唆していますが、「隠れた慣性」こそがエキスパートにとって欠けていたピースであると確信しています。

要約すると、この論文は、これらの特別なロボットの脳を訓練するためにフルセットのノートは必要ではなく、単に書き留めることなく直近の過去を記憶する方法が必要であると主張しています。この「隠れた慣性」アプローチであるMESHは、メモリを大量に消費することなく、ロボットが十分に賢く学習できるようにしながら、その大部分を実現するものです。

技術要約: MESH – Mixture-of-Experts トレーニングのためのメモリ効率の高い Sinkhorn 最適化

1. 問題提起

本論文は、Mixture-of-Experts (MoE) 言語モデルをメモリ効率よく最適化するという課題に取り組んでいる。AdamW は、時間的な一次モーメントの平滑化と座標ごとの適応的スケーリングの組み合わせにより、事前学習における標準的なオプティマイザとなっているが、パラメータごとに2つのフル状態テンソルを保持するため、高いメモリコストを伴う。

SAGE や Sinkhorn ハイブリッドのような、行列のモーメント状態を削除することでコストを削減する Sinkhorn スタイルの行列オプティマイザやメモリ効率の高い代替手法が存在する。しかし、著者は、これらのステートレス（状態を持たない）な手法をルーティングされた MoE エキスパートに直接適用すると失敗することを示している。制御された 110M パラメータの「nanowhale」DeepSeek スタイルの MoE 設定において、標準的な SAGE/Sinkhorn ハイブリッドは、オプティマイザ状態を大幅に削減したものの（0.883GB から 0.331GB へ）、評価損失において深刻な劣化（AdamW ベースラインの 3.58–3.64 に対し 3.8265）を招いた。

核心的な失敗の原因は、ルーティングされたエキスパートに特有の順序の問題であると特定された：

ルーティングされた勾配は条件的である: 密な MLP の勾配とは異なり、ルーティングされたエキスパートの勾配は、ルーターの決定とゲートの重みに基づく条件的サンプルであり、これらは時間的に変動する。
Sinkhorn の非線形性: Sinkhorn 正規化は非線形演算である。瞬時で高分散なルーティング勾配 ( $Sinkhorn(G_t)$ ) に Sinkhorn を適用することは、時間的に平均化された信号 ( $Sinkhorn(E[G_t])$ ) に適用することとは挙動が異なる。
ノイズの増幅: ステートレスな Sinkhorn は、毎ステップでルーティングに起因するノイズを正規化するが、ステートレスな手法の失敗は、行列正規化が行われる前に時間的な平滑化が行われていないことに起因する。

2. 手法: MESH および MESH-B

著者は、MoE アーキテクチャ向けに特別に設計されたハイブリッド・オプティマイザである MESH (Memory-Efficient Sinkhorn for Experts) を提案している。

コアメカニズム: 隠れモーメンタム (Hidden Momentum)

MESH は、MoE エキスパートに対するステートレスな Sinkhorn 更新を、隠れモーメンタムアプローチに置き換える。明示的な一次モーメントテンソルをオプティマイザの状態として保持する代わりに（これはメモリ節約の目的を損なうことになる）、MESH は勾配バッファのライフサイクル内に一次モーメントを保持する：

蓄積 (Accumulation): 非正規化されたエキスパート・モーメンタム蓄積器が保持される: $H_{e,t} = \beta_1 H_{e,t-1} + G_{e,t}$ 。
バッファのライフサイクル (Buffer Lifecycle): 蓄積器 $H_{e,t}$ は、オプティマイザの状態辞書ではなく、勾配バッファに格納される。
更新サイクル (Update Cycle): バックプロパースの前に、バッファが勾配として再アタッチされる ( $p.grad \leftarrow \beta_1 H_{e,t}$ )。バックプロパースによって現在の勾配が加算され、実質的に次のステップのための蓄積器が更新される。
正規化 (Normalization): 平滑化された信号に対して Sinkhorn 更新が計算される: $U_{e,t} = Sinkhorn(H_{e,t})$ 。

このアプローチは、エキスパート用の明示的なオプティサイザ状態テンソルを保持することなく、ルーティングされたエキスパートに必要な時間的な一次モーメント信号を復元する。

オプションの拡張: MESH-B

また、本論文では、Sinkhorn の方向に粗いブロック/ニューロンの逆 RMS 乗数を加えた MESH-B についても評価している。これはブロックレベルの前処理（プリコンディショナ）として機能する。実装の単純さを維持するため、ブロック統計量は、個別の生の勾配ストリームからではなく、平滑化されたソース ( $H_{e,t}$ ) から計算される。

ロール割り当て戦略

MESH は、異なるパラメータクラスに対してハイブリッドな戦略を維持する：

埋め込み (Embeddings) / LM Head: SAGE によって処理される（ライトな状態更新）。
非エキスパート密行列 (Non-Expert Dense Matrices): ステートレスな Sinkhorn によって処理される。
ルーティングされた/共有されたエキスパート (Routed/Shared Experts): MESH（隠れモーメンタム Sinkhorn）によって処理される。
正規化層 / バイアス (Norms / Biases): SAGE によって処理される。

3. 主な貢献

失敗点の特定: 本論文は、ルーティングされた Mo 向きのエキスパート行列が、直接的な SAGE/Sinkhorn 最適化を事前学習に適用した際の具体的な失敗点であることを特定し、密な行列や埋め込みとは区別している。
MESH アルゴリズム: 勾配バッファのライフサイクルを利用して、明示的なオプティマイザ状態を追加することなくエキスパートのための一次モーメント信号を復元する、隠れモーメンタム Sinkhorn 更新を導入した。
アブレーション分析: 二次モーメントのみ、符号/モーメンタムのみ、信頼度のみ、およびスケールのみの介入では、性能の差を説明できないことを示す系統的な負のアブレーション分析を行った。主要な因果的要素は、行列正規化の前の時間的平滑化であると特定された。
MESH-B の評価: ブロック前処理を加えたバリアントを評価し、特定のシードにおいてメモリと品質のフロンティアを改善できることは示されたが、MESH が提供する時間的平滑化と比較して、普遍的に必要であるとは結論づけられていない。

4. 実験結果

実験は、FineWeb-Edu 上で 5,000 ステップ学習された 110M パラメータの DeepSeek スタイルの MoE モデルを用いて行われた。

メモリ効率:
- AdamW: 0.883GB のオプティマイザ状態。
- MESH / MESH-B: 0.331GB のオプティマイザ状態（AdamW に対して 62.5% の削減）。
- ピーク CUDA 割り当て: 約 12.6% 減少（~4.37GB から ~3.82GB へ）。
パフォーマンス (評価損失):
- AdamW ベースライン: ~3.5877 (シード 100/512 の平均)。
- ステートレス SAGE/Sinkhorn: 3.8265 (大幅な劣化)。
- MESH: 3.6430 (AdamW に近いが、わずかに高い)。
- MESH-B: 3.6365 (MESH よりわずかに良いが、依然として AdamW には及ばない)。
診断結果 (シード 42):
- 明示的に一次モーメントを保持するフル状態バリアント (MESH-FULL, MESH-B-FULL) は、AdamW に近い品質（例：V(H) に達した MESH-B-FULL は 3.6102）を、AdamW の約半分のオプティマイザ状態で回復した。
- これにより、時間的平滑化が決定的な欠けていた要素であり、隠れモーメンタックスの実装が、フルな状態コストをかけることなくこれをうまく近似していることが確認された。

5. 意義と主張

本論文は、MESH が MoE トレーニングにおけるメモリと品質のフロンティアにおいて、大幅に低い状態量（ステート量）の地点を提供すると控えめに主張している。

主要な結論: MoE エキスパートは、Sinkhorn スタイルのオプティマイザを効果的に機能させるために、行列正規化の前の時間的平滑化を必要とする。この平滑化が、AdamW との性能差を埋めるための支配的な因果要因である。
二次的な結論: ブロック/ニューロンの前処理 (MESH-B) は、特定の条件下で結果を改善できるが、それが普遍的に必要な構成要素であるとは確立されていない。
限界: 著者は、MESH/MESH-B と AdamW の間の残留する性能差を認めている。これは、彼らのセットアップにおいて、語彙行列や非エキスパート密行列などの非エキスパートの役割が SAGE またはステートレスな Sinkhorn に残っているためであると仮定している（これらは AdamW ではフルモーメントで適応される）。
範囲: 本研究は意図的に限定的であり、特定の 110M アーキテクチャと 5,000 ステップの実行に焦点を当てている。著者は、MESH がすべての指標で AdamW を上回るとは主張しておらず、ステートレスな手法によって失われた品質の大部分を回復しつつ、メモリオーバーヘッドを劇的に削減することを目指している。

要約すると、MESH は、MoE トレーニングにおけるメモリ効率の高いオプティマイザの失敗が、MoE 構造そのものに起因するのではなく、ルーティングされた勾配に対する時間的平滑化の欠如に起因することを実証している。「隠れモーメンタム」を実装することで、MESH はこの順序の問題を解決し、メモリ効率の高い MoE 事前学習への実行可能な道筋を提示している。