Hierarchical Copula-Gumbel-Top-\texorpdfstring{$K$}{K} Routing: Two-Sided Dependence Control for Frozen Mixture-of-Experts at Fixed Per-Token Routing Laws

技術要約：階層的コピュラ・ガンベル・Top-K ルーティング

問題提起
スパースな Mixture-of-Experts (MoE) モデルにおいて、標準的な確率的ルーティング機構（Gumbel-Top-K）は、各トークンに対して「ルーティング則」を定義します。これは、トークンのロジットに基づいた、順序付けられたエキスパートのリストおよび混合重みに対する特定の確率分布です。既存の研究のほとんどは、このパー・トークン（トークンごとの）の法則を修正すること（例：ロジットの微調整や補助損失によるもの）に焦点を当てていますが、異なるトークン間におけるルーティング選択の結合分布については、ほとんど検討されてきませんでした。本論文は次のような問いを投げかけます。すべての個々のトークンのルーティング則が厳密に固定されている場合（凍結されたモデルの場合）、異なるトークンのルーティング選択の結合分布として、どのようなものが到達可能でしょうか？ 特に、単一のトークンの周辺挙動を変更することなく、トークン間の依存関係を制御することで、エキスパートのコヒーレンス（一貫性）や負荷分散を改善できるでしょうか？

手法：階層的コピュラ・ガンベル・Top-K (H-CGA)
著者は、Top-K 選択の前のガンベル・ノイズ・ベクトルに対して作用する、ルーティング側の依存性レイヤーである H-CGA を提案しています。この手法は、周辺分布と依存構造を分離するスカラ（Sklar）の定理に基づいています。構築は以下の4つの段階で行われます。

1. プリ・ルーティング・グルーピング： トークンは、互いに素なグループ（例：隣接するトークンの固定ウィンドウ）に分割されます。コントローラーは、凍結されたプリ・ルーティング特徴量（例：隠れ状態、ゲート・エントロピー）を読み取り、結合強度を決定します。
2. グループ内正の結合： グループ $g$ 内のトークンに対して、各エキスパート座標のガンベル摂動は、交換可能なガウス・コピュラを介して生成されます。トークンは共通の潜在変数 $\zeta_{ge}$ を共有し、これによりランダムな摂動に正の相関が誘発されます。これにより、関連するトークンが同じエキスパートを選択する可能性が高まります（コヒーレンス）。
3. グループ間アンチセティック（対抗的）結合： 互いに素なグループのペア $(g, g')$ がマッチングされます。調整可能な反対強度 $\alpha_{g,g'} \in [0, 1]$ が導入されます。第2のグループの共有潜在変数は、第1のグループの潜在変数と独立したノイズの線形結合として構成されます。$\alpha_{g,g'} = 1$ のとき、グループは完全に反対の共有プッシュを受け取り、$\alpha_{g,g'} = 0$ のとき、グループは独立しています。これにより、グループ間に負の依存性が導入され、負荷のバースト性を抑制します。
4. 周辺分布の保存： 共有潜在変数からガンベル変数への変換により、任意の単一トークン $t$ について、摂動ベクトル $(\gamma_{t1}, \dots, \gamma_{tE})$ は、独立したルーティングと同様に、i.i.d.（独立同一分布）なガンベル分布に従うことが保証されます。

主な貢献

1. H-CGA アーキテクチャ： ルーターのロジット、エキスパートの重み、またはいかなるパー・トークンのルーティング則も変更することなく、クロス・トークンの依存性を制御する階層的コピュラ・レイヤー。
2. 理論的不変性： 著者は、この階層全体が、すべてのトークンの条件付き順序付き Top-K リストおよび混合重みの分布を保存することを証明しています（定理1）。その結果、条件付き期待エキスパート包含カウント（系1）は、独立したルーティングのそれと同一になります。
3. コヒーレンス–分散のトレードオフ特性化： 本論文は、分散のトレードオフを特性化しています（命題1）：
   • グループ内の正の結合を強めると、独立したルーティングと比較して、実現されたエキスパート負荷の分散が増大します（局所的なコヒーレンスの代償）。
   • 非負のグループ間反対強度は、同じグループ内結合強度における「フラットな」結合（独立したグループ）と比較して、この分散を減少させます。ただし、必ずしも独立したルーティングのベースラインよりも分散を低減させるわけではありません。
4. ルーティングのみの適応： 著者は、小さなコントローラーが、凍結された特徴量に基づいて依存性のダイヤル（$\rho_g, \alpha_{g,g'}$）を設定する適応手法を定式化しました。このコントローラーはスコア関数エスティメータ（REINFORCE）を用いて訓練され、凍結されたベースモデルのフォワードパスと、コントローラーに限定された勾配のみを必要とします。

結果およびパイロット研究
本論文は、TinyStories で学習された 15.8M パラメータのデコーダー専用 MoE を用いた小規模なパイロット研究を提示しています。評価中、ベースモデルは凍結されています。

• 不変性の検証： 合成テストにより、H-CGA 下での順序付き Top-K リストおよびエキスパート包含確率の経験的頻度が、モンテカルロ誤差の範囲内で独立したルーティングと一致することが確認されました。これは、理論的不変性を検証するものです。
• 統計的効果：
  • グループ内の結合（$\rho$）を増加させると、ウィンドウ内のエキスパート集合のジャカード類似度が大幅に増加し、ウィンドウあたりの異なるエキスパートの使用数が減少しました。
  • グループ間の反対強度（$\alpha$）を増加させると、グループ境界におけるジャカード類似度が大幅に低下し、グループ間の需要をデカップリングできることが確認されました。
• 学習： 3つのパラメータを持つ学習されたスカラー・コントローラーは安定して訓練されましたが、独立したベースラインと比較して検証クロスエントロピーの改善は見られませんでした。学習された結合強度は低く、これは「単一のトークン・オブジェクティブのみを使用する場合、周辺保存によってコントローラーが特定できる信号が制限される」という論文の観察と一致しています。
• 負荷分散： パイロット研究では、集計負荷の変動係数（CV）の変化は最小限であり、この要約指標は、命題1で提案された条件付き分散の境界を直接テストするものではないことが指摘されました。

意義と主張
本論文は、H-CGA が MoE ルーティングにおける新しい自由度、すなわちルーティング・ノイズの結合依存性を明らかにするものであると主張しています。その主な意義は、理論的かつメカニズム的なものです。

• 安全性： クロストークンの依存性を変更しても、いかなるトークンの周辺的なルーティング挙動も変化させないという「安全保証」を提供し、学習された好みが損なわれないことを保証します。
• 制御： 局所的なエキスパート集合のコヒーレンスと、グローバルな負荷分散をトレードオフするための、2つの補完的な「ダイヤル」を提供します。
• 効率性： エキスパートの重みを更新するコストを回避し、フォワード方向のみでベースモデルを評価する、ルーティングのみの適応を可能にします。

限界と謙虚な姿勢
著者は、スコープと限界について明示しています。

• レイヤー局所的な保証： 不変性の結果は、特定のルーティング・レイヤーに入力される隠れ状態とロジットに条件付けられています。結合がダウンストリームの隠れ状態を変更する可能性があるため、これらはマルチレイヤー・ネットワークにおけるエンドツーエンドの不変性を保証するものではありません。
• タスクレベルの利得なし： パイロット研究は、ルーティングのみの適応がダウンストリーム・タスクの性能を向上させることを確立していません。著者は、周辺保存が意味するところとして、単一のトークン分布のみに依存する目的関数ではコントローラーを特定できないことを指摘しています。有用な信号は、結合効果（例：アテンションの相互作用）から生じる必要がありますが、それらは弱い可能性があります。
• 分散 vs オーバーフロー： 本手法は負荷の分散を理論的に特性化していますが、パイロット研究では条件付き負荷分散や容量オーバーフローを直接測定していません。本論文は容量オーバーフローを解決すると主張しているのではなく、反対ダイヤルがそれを対処するためのメカニズムであることを示唆しています。
• オーバーヘッド： 本手法は、すべてのエキスパートに対してガンベル摂動を生成し、調整する必要があり、標準的なルーティングと比較して計算オーバーヘッドが生じる可能性があります。

要約すると、本論文は、凍結された MoE においてクロストークンのルーティング依存性を制御するための数学的に厳密なフレームワークを確立しており、個々のトークンのルーティング則を変更することなく、コヒーレンスと負荷分散を操作できることを証明しています。同時に、タスク性能に対する即時的な実用的利益については慎重な姿勢を保っています。