Affinity Is Not Enough: Recovering the Free Energy Principle in Mixture-of-Experts

Man Yung (Russell) Wong による論文「Affinity Is Not Enough: Recovering the Free Energy Principle in Mixture-of-Experts」の詳細な技術的サマリーを以下に示す。

1. 問題定義

スパースなミクスチャー・オブ・エキスパート（MoE）モデル（例：DeepSeek-V3、Mixtral）は、現在ステートレスなアフィニティベースのルーティングに依存している。ルーティングゲートは、現在のトークン埋め込み（$x_t$）のみに基づいてエキスパートを選択し、時間的コンテキスト、エキスパートの信頼性の履歴、および将来の要件を無視している。

この論文は、ドメイン遷移という致命的な失敗モードを特定している。シーケンスが一つの分布（ドメイン A）から別の分布（ドメイン B）へ移行する際、遷移ステップ（$t$）のトークンは依然としてドメイン A に属するが、次のステップ（$t+1$）における最適なエキスパートはドメイン B に属する。

• 問題点: 標準的なルーターは $x_t$ しか見ていないため、遷移ステップにおいて正しいエキスパートにほぼゼロの確率を割り当てる。
• 実証的証拠: 制御された実験において、標準的なルーティングは遷移ステップにおいて正しいエキスパートに0.006 ± 0.001の確率質量しか割り当てない。
• 結果: このような低い信頼度で 99% のルーティングカバレッジを達成するには、モデルは非現実的な数のエキスパート（$K \approx 765$）を必要とし、MoE の効率性の利点を無効化してしまう。

2. 手法：自由エネルギー原理（FEP）フレームワーク

著者は、MoE ルーティングをFriston の自由エネルギー原理を通じて再定義すべきであると提案している。この原理によれば、不確実性下での最適システムは以下のことを要求される：

1. 状態を再帰的に更新する（ゼロから始めない）。
2. 信号を精度（誤差の逆分散）で重み付けする。
3. 将来の期待自由エネルギーを最小化する（現在の観測だけでなく、予測に基づいて行動する）。

これを実現するために、論文は**スパイクニューラルネットワーク（SNN）**のダイナミクス（具体的にはリーキー・インテグレート・アンド・ファイア、LIF）に基づき、標準的なルーティングゲートに 3 つの軽量な修正を導入する：

A. 時間的ルーティングメモリ（$\beta$）

• メカニズム: ステートレスな入力 $x_t$ を、LIF ダイナミクスを通じて更新される隠れ状態 $h_t$ に置き換える：
  $$h_t = \beta \odot h_{t-1} + x_t$$
  ここで $\beta$ は学習可能な減衰ベクトル（初期値 0.9）である。
• 機能: トークン全体にわたるコンテキストを蓄積する。これにより、ルーターは「システムが長い間ドメイン A にあった」ことを「記憶」し、遷移トークンが現れる前にも、差し迫ったスイッチを示す「飽和」信号を生成できる。
• 時間スケール: トークンごと（高速）。

B. 精度重み付けゲート（$\Pi$）

• メカニズム: 各エキスパートの平均二乗誤差（MSE）のオンライン非パラメトリック追跡器を維持する。ルーティングゲートは、これらの誤差の逆分散（精度）によって変調される：
  $$g_t = \text{softmax}(W \cdot x_t \odot \Pi_t)$$
  ここで $\Pi_i = 1 / (\hat{\sigma}^2_i + \epsilon)$ である。
• 機能: 現在信頼性が低い（誤差が大きい）エキスパートを動的にダウンウェイトし、信頼性の高いものをアップウェイトする。標準的な負荷分散（頻度を追跡する）とは異なり、これは精度を追跡する。
• 時間スケール: バッチごと（中速）。

C. 予期的ルーティング

• メカニズム: 予測ネットワーク $f(x_t, h_t)$ を使用して、次のトークン埋め込み $\hat{x}_{t+1}$ を推定する。ルーティング決定はこの予測に基づいて条件付けられる：
  $$g_t = \text{softmax}(W \cdot h_t + W_{pred} \cdot \hat{x}_{t+1})$$
• 機能: 現在の状態ではなく、予測された将来の状態に基づいてルーティングを行う。これにより、ルーターは新しいドメインのトークンが入力に現れる前に、新しいドメインのエキスパートへ切り替えることができる。
• 決定的な洞察: ステートレスな予測器（$f(x_t)$）は失敗する。なぜなら、遷移前のトークンは分布論的に同一だからである。予測器は、遷移を示す飽和信号を検出するために、$\beta$ によって蓄積された状態 $h_t$ に条件付けられなければならない。

3. 主要な貢献

1. 理論的統合: MoE ルーティングを自由エネルギー原理と結びつけ、標準的な MoE が再帰的な状態推定、精度重み付け、および将来状態の最小化を欠いていることを実証する。
2. 3 つの軽量メカニズム: $\beta$（メモリ）、$\Pi$（精度）、および Anticipation（予測）を必要な構成要素として提案する。これらは最小限のパラメータオーバーヘッド（合計約 200 行のコード）で実装可能である。
3. 超加和性の発見: この論文の最も重要な発見は、$\beta \times \text{Ant}$ の相互作用である。
   • $\beta$ 単独：+0.295 の改善。
   • Anticipation 単独：+0.000 の改善（メモリなしでは失敗）。
   • 組み合わせ: +0.741 の改善。
   • 相互作用効果: 組み合わせによる改善は、個々の部分の合計を**+0.446**上回り、メモリが予期のための前提条件であることを証明する。
4. DeepSeek-V3 に対する構造的説明: DeepSeek-V3 が推論時にマルチトークン予測（MTP）モジュールを破棄する理由を説明する：ステートレスなMTP ヘッドはドメイン遷移を検出できないからである。論文は、MTP が時間的メモリ（$\beta$）に条件付けられれば、単なるトレーニングの補助ではなく強力な予期的ルーティングメカニズムとなり得ると主張する。

4. 実験結果

実験は、制御されたドメインスイッチタスクと、文字レベルの MoE 言語モデル（5 つのシード）で行われた。

指標	標準 MoE	$\beta$-MoE	$\beta$ + Ant MoE	Oracle
遷移精度	0.006	0.301	0.748	0.994
遷移時の BPC	6.56	4.01	4.19	N/A
99% カバレッジに必要なエキスパート数（$K$）	~765（非現実的）	~6.0	2.4	N/A

• 遷移精度: 組み合わせたメカニズムは、遷移ステップにおける正しいエキスパートの選択確率を**0.6% から 74.8%**に引き上げ、オラクルとのギャップの 75% を埋める。
• 効率性: 99% のカバレッジを達成するために必要なエキスパート数（$K$）は、非現実的な約 765 から、遷移ステップにおいて小さな定数（2.4）に低下する。
• 言語モデルの検証: 文字レベルの LM において、$\beta$ ルーティングは遷移ステップでの文字あたりビット数（BPC）を 6.56 から 4.01 に削減した。予期を追加するとルーティング決定が安定化（分散の減少）したが、この小規模スケールでは BPC をさらに削減しなかった。これは、主な利点がルーティングの安定性と、プリフェッチによるレイテンシ削減にあることを示唆している。

5. 意義と含意

• ナビゲーション対分類: この論文は、ルーティングを静的な分類問題（「どのエキスパートがこのトークンに適合するか？」）から、動的なナビゲーション問題（「システムはどこへ向かっているのか、どの経路が信頼できるのか？」）へと再定義する。
• レイテンシ削減: 1 ステップ先でドメインスイッチを予期することで、システムは現在のステップが実行されている間に、メモリ（HBM）から計算ユニット（SRAM）へエキスパート重みをプリフェッチでき、重要な遷移点におけるメモリレイテンシのストールを排除できる。
• MTP の再評価: この研究は、DeepSeek-V3 などのモデルにおけるマルチトークン予測モジュールが過小利用されていることを示唆する。時間的状態に条件付けられれば、これらは単なるトレーニングの補助ではなく、強力な予期的ルーティングメカニズムとして機能し得る。
• 生物学的妥当性: これらのメカニズムは生物学的プロセス（LIF ニューロン、予測符号化における精度重み付け、能動的推論）を反映しており、現在の MoE アーキテクチャが知的なリソース配分の基本原理を見落としていることを示唆している。

結論として、この論文は、アフィニティだけでは堅牢な MoE ルーティングには不十分であることを実証する。時間的メモリ、精度重み付け、および予期を統合することで、MoE モデルはドメイン遷移においてオラクルに近いルーティング性能を達成し、高精度を維持するための計算コストを劇的に削減できる。