技術的サマリー：2 段階最適化を備えた階層的混合エキスパート（Hi-MoE）

1. 問題定義

スパース混合エキスパート（MoE）モデルは、トークンを少数のエキスパートネットワークにルーティングすることで容量を拡張します。しかし、これらのモデルの訓練には、負荷分散とエキスパートの専門化の間の根本的なトレードオフが存在します。

• ルーティングの崩壊（Routing Collapse）： 過激な多様性メカニズムは、しばしば「ルーティングの崩壊」を引き起こします。これは、少数のエキスパートがトークン選択を支配し、他のエキスパートがアイドル状態になる現象です。これにより、容量の効率的な利用が阻害され、デバイス間で負荷の偏り（ストレイガー）が生じます。
• グループ化のジレンマ： 最近のアプローチは、エキスパートをグループ化し、各グループから固定数の選択を強制すること（例：$M$個のグループのそれぞれから $K/M$ 個のエキスパートを選択）によって不均衡を緩和します。これはデバイスレベルの計算の均一性を保証しますが、ルーティングの自由度を制限します。この制約はしばしば冗長性をもたらします。異なるグループのエキスパートが、類似したトークンの混合で最適化されるため、類似した機能を学習し、エキスパートレベルの崩壊をグループ間の「多様性の崩壊」に置き換えてしまうのです。

既存の損失ベースの負荷分散手法（例：GShard）は、タスク勾配と衝突して収束を遅らせる傾向があり、損失不要な手法は冗長性を防ぐのに失敗する可能性があります。ハードウェアに整合した負荷分散を強制しつつ、意味のある相補的な専門化を促進するフレームワークが必要です。

2. 手法：Hi-MoE

著者は、ルーティング制御をグループ間バランスとグループ内専門化という 2 つの結合されたレベルに分解する階層的フレームワークであるHi-MoEを提案します。

アーキテクチャ

• グループ化された構造： エキスパートは $M$ 個の互いに排他的なグループに分割されます（多くの場合、GPU シャードに整合します）。
• ルーティング制約： 各トークンに対して、ルーターは各グループ $m$ から固定数のエキスパート（$K_m$）を選択します。
• 2 段階最適化： フレームワークは、階層の各レベルに対して特定の正則化項を導入します。

主要コンポーネント

1. グループ間正則化（$R_{inter}$）：

   • 目的： グループ間での公平なトラフィック分布を強制し、デバイス利用のバランスを保証すること。
   • メカニズム： グループ間の選択後のルーティング重みの $\ell_2$ 集中度を最小化します。特定のグループにルーティング質量が集中することを罰することで、グループ負荷の変動係数（CV）を最小化する代理指標として機能します。
   • 理論的根拠： 定理 4.1 は、この項を最小化することがグループレベルの不均衡に対する保証付きの上限を提供することを証明しています。

2. グループ内正則化（$R_{intra}$）：

   • 目的： グループ内での相補的な専門化を促進し、エキスパートが冗長な機能を学習するのを防ぐこと。
   • メカニズム： 選択前のルーティング分布の $\ell_2$ 集中度を最大化します（実質的にオーバーラップ防止項）。これにより、ルーターは多くのエキスパートを低信頼度で混合するのではなく、グループ内の特定のエキスパートに対して決定的でスパースなコミットメントを行うように促されます。
   • 理論的根拠： 定理 4.2 は、ルーティングのオーバーラップを減少させることがエキスパート間の勾配結合を減少させ、それぞれが異なる最適化軌跡をたどることを可能にすることを示しています。定理 4.3 は、バランスの取れた周辺分布の下で、ルーティング集中度を増加させることが、入力と選択されたエキスパート間の相互情報を最大化することを示しており、これは専門化の代理指標となります。

3. 安定化メカニズム：

   • グループ内の多様性シグナルが退化した領域にドリフトするのを防ぐため、著者はルーターのロジットの指数移動平均（EMA）を用いたバイアス補正ルーティング計算を採用します。これにより、個別の負荷分散損失を追加することなく、歴史的に好まれたエキスパートの持続的な過選択を抑制します。

目的関数

全体の訓練目的関数は、タスク損失、標準的な負荷分散損失、および提案された階層的項を組み合わせます。
$$\mathcal{L} = \mathcal{L}_{task} + \mathcal{L}_{load} + R_{intra} + R_{inter}$$
著者はこれを制約付き最適化問題として定式化し、係数 $\lambda_{intra}$ と $\lambda_{inter}$ をそれぞれ多様性とバランスの制約に対するラグランジュ乗数として扱います。

3. 主要な貢献

• 緊張関係の特定： 厳格なグループ制約を通じてデバイスレベルのバランスを強制することは利用を安定化させるが、グループ間の冗長性を増加させる可能性があることを、本論文は特定しました。
• 階層的フレームワーク： Hi-MoE は、これらの懸念を分離するモジュール式でドロップイン可能な目的関数を導入します。
  • ハードウェアに整合したバランスのためのグループ間目的関数。
  • 機能的な多様性と専門化のためのグループ内目的関数。
• 理論的保証： 正則化項をグループ不均衡の境界、勾配の結合解除、およびルーティング情報の増加（専門化）に関連付ける形式的な証明を提供しています。
• 実証的検証： ビジョンおよび言語ドメイン全体で一貫した改善を示し、Hi-MoE がパープレキシティ/精度と負荷分散（CV）の間のパレートフロンティアを拡張することを証明しました。

4. 実験結果

著者は、Hi-MoE を 3 つの異なる事前学習設定で評価しました。

A. コンピュータビジョン（Tiny ImageNet 上の Swin-MoE）

• 設定： 64 個のエキスパート（16 個のグループ 4 つ）を備えた 10 億パラメータモデル。
• 結果： Hi-MoE は、Vanilla GShard、Loss-free 分散、および MoGE と比較して、最高の Top-1 精度（59.13%）および最良の負荷分散（CV 0.29）を達成しました。
• 観察： アテンションマップの可視化により、Hi-MoE のグループは明確な視覚的専門化（例、あるグループは身体部位に、別のグループは背景表面に焦点を当てる）を発展させたのに対し、ベースラインではより多くの冗長性が示されたことがわかりました。

B. 言語モデリング（OpenWebText 上の nanoGPT）

• 設定： 8 個のエキスパート（2 個のグループ 4 つ）を備えた 10.6 億パラメータモデル。
• 結果： Hi-MoE は、最低のパープレキシティ（2.947）および最良のバランス（CV 0.23）を達成しました。
• アブレーション： $R_{inter}$ を除去するとバランスが劣化し、$R_{intra}$ を除去すると品質が劣化しました。バイアス補正は安定化剤として機能しました。
• 効率性： 階層的ルーターはレイヤーあたり 1.3ms から 1.6ms の最小限のオーバーヘッドを追加する一方で、フラットなルーティングと比較して同期オーバーヘッドを約 3 倍削減しました。

C. 大規模事前学習（OLMoE-7B vs. Hi-MoE-7B）

• 設定： Dolma 混合データセットの 585.4 億トークンで訓練された 70 億パラメータモデル。
• パープレキシティ： Hi-MoE-7B は、標準的な OLMoE-7B ベースラインと比較して、11 の検証ドメイン全体で平均してパープレキシティを**5.6%**削減しました。
• 下流タスク： Hi-MoE は、MMLU スライス（STEM で +22.2%、人文科学で +11.9% など）および常識推論タスクで顕著な向上を示しました。
• 負荷分散： Hi-MoE はエキスパートのバランス（CV）を40%（0.153 から 0.257）改善し、大規模訓練におけるリソース効率のボトルネックを直接解決しました。

5. 意義と主張

本論文は、Hi-MoE がスパース MoE 訓練に内在する「バランスと多様性」の緊張関係を解決すると主張しています。ルーティングの階層性（グループ対エキスパート）を明示的にモデル化することにより、この手法はモデルが以下のことを可能にします。

1. 効率的なスケーリング： モデルの条件付き容量を犠牲にすることなく、均一な GPU 利用を維持する。
2. 専門化された表現の学習： グループ化されたエキスパートが機能的に冗長になる「多様性の崩壊」を防ぐ。
3. 品質の向上： ビジョンおよび言語タスク全体で、パープレキシティおよび下流タスクのパフォーマンスにおいて一貫した向上を達成する。

著者は、Hi-MoE を既存の安定化メカニズムの代替としてではなく、エキスパートグループ化を単なるシステム制約ではなく、原則的な最適化メカニズムとする補完的なモジュール式強化として位置づけています。この研究は、階層的な調整が、兆単パラメータモデルにおける信頼性の高い品質向上を解き放つための viable な道であることを示唆しています。