技術的サマリー：MoE 専門化における幾何学的非対称性

問題定義

Mixture-of-Experts（MoE）アーキテクチャは、固定された計算コストで大規模なパラメータ数を可能にすることで、Transformer モデルをスケーリングするための標準的なメカニズムとなっているが、専門化されたエキスパートの内部幾何学的構造は未だ十分に理解されていない。既存の分析は、主にルーティング統計、負荷分散、またはシステムレベルのパフォーマンスに焦点を当てている。エキスパートが関数空間（実行する変換）と表現空間（ルーティングされたトークンが隠れ状態内でどのように組織化されているか）のいずれにおいてどのように異なるかを理解するという点に、重要なギャップが存在する。

現在の直観では、MoE 層は入力分布を互いに排他的なクラスタに分割し、各エキスパートが局所的な近似器として機能すると考えられている。しかし、実証的証拠は、エキスパートがしばしば重複するトークン領域で動作することを示しており、以下の問いを提起する：事前学習済み MoE Transformer の内部において、ルーティングはどのような幾何学的構造を誘発するか？ 具体的には、エキスパートは異なる関数を実装しているのか、またそれらの表現は互いにどのように関連しているのか？

手法

著者は、関数空間と表現空間の両方における MoE 層を分析するための統合されたヤコビアン-PCA-グラスマン枠組みを導入する。この枠組みは、事前学習済みモデル（Mistral、Qwen）と、ゼロから訓練された制御された 3 層 Transformer モデルに適用される。

1. 関数的幾何学：エキスパート局所ヤコビアン

個々のエキスパートの機能をルーティングダイナミクスから独立して特徴づけるため、著者はエキスパート局所ヤコビアン（$J_e(x) = \partial f_e(x) / \partial x$）を分析する。

• 指標: 各エキスパートについてデータセット全体で平均ヤコビアン（$\bar{J}_e$）を計算し、これをベクトル化する。
• 整合性: これらのベクトル化されたヤコビアンの間の正規化されたフロベニウス内積（コサイン類似度）を通じて、ペアワイズな関数的整合性を測定する。
• 解釈: 低いコサイン類似度は関数的非相関を示し、これはエキスパートが同じ関数の冗長な変種ではなく、局所的に異なる変換を実装していることを意味する。

2. 表現幾何学：ルーティング済み PCA とグラスマン距離

ルーティングされた表現の幾何学を特徴づけるため、著者はルーティング重みで重み付けされた隠れ状態（$h_e(x) = g_e(x)h(x)$）を収集する。

• PCA: ルーティングされた隠れ状態に対して標準的な主成分分析（PCA）を行い、主方向と分散スペクトルを特定する。
• 部分空間距離: 異なるエキスパートのトップ-$n$ PCA 部分空間間のグラスマン測地距離を用いて、エキスパート間の幾何学的分離を定量化する。
• 解釈: この指標は、完全な部分空間間の多方向的な分離を捉える。距離がゼロに近い場合は冗長な機能的役割を意味し、理論的最大値（$\frac{\pi}{2}\sqrt{n}$）に近づく場合は直交性を意味する。

3. 制御されたルーティング実験

ルーティングメカニズムの効果を訓練ダイナミクスから分離するため、著者は固定されたエキスパートパラメータを持つ制御された 3 層モデルにおいて 2 つの領域を比較する。

• Top-k ルーティング: 疎で離散的な選択。
• フルソフトルーティング: 密で連続的な混合。

主要な結果

1. 構造的非対称性：関数的非相関対表現的重複

事前学習済み MoE Transformer（Mistral-8x7B および Qwen1.5-MoE-A2.7B）全体において、著者は一貫した構造的非対称性を観察する。

• 強い関数的非相関: クロスエキスパートヤコビアン整合性は一貫してゼロに近い（平均 $\approx$ 0.000–0.068）。これは、エキスパートが冗長ではなく、異なる局所変換を実装していることを示している。
• 部分的な表現的重複: 関数的非相関にもかかわらず、ルーティングされた表現は明確だが部分的に重複する部分空間を占有する。グラスマン距離は有意義（分離を示す）であるが、完全な直交性に対する理論的最大値よりも著しく低い。
• 結論: 関数的非相関と表現的重複は共存する。エキスパートは、入力空間を互いに排他的なクラスタに分割するのではなく、共有多様体上の重複する領域において局所的に非相関な演算子として機能する。

2. 幾何学的制御ノブとしてのルーティング疎性

制御された実験により、ルーティングの鋭さがこの幾何学を直接形成することが明らかになった。

• Top-k ルーティング: より鋭い関数的分離（低いヤコビアン整合性）と、より大きなグラスマン距離（より大きな部分空間の発散）を誘発する。
• フルソフトルーティング: 高いヤコビアン整合性（冗長なマッピングへの収束を示す）と、収束したグラスマン距離（高い部分空間の重複）を伴う、絡み合ったエキスパート構造をもたらす。
• 含意: ルーティング疎性は、専門化を強制する主要な要因である。

3. 密層との比較

MoE エキスパートは、密なフィードフォワード層と比較して、分散スペクトルにおいてより急峻な固有値の減衰を示す。最初の数個の主成分は、MoE エキスパートにおいて分散の著しく高い割合を捉えており、これはルーティングがトークンを低次元でより焦点の絞られた部分空間に集中させていることを示唆している。

4. モデル間の一貫性

観察された幾何学的パターン（低いヤコビアン整合性 + 高いが最大値未満のグラスマン分離）は、異なるモデル（Mistral、Qwen）、層（初期、中間、後期）、およびデータセット全体で一貫しており、これは特定の訓練実行のアーティファクトではなく、MoE アーキテクチャの頑健な性質であることを示唆している。

意義と主張

本論文は、現代の Transformer アーキテクチャにおける条件付き計算を研究するための一般的な診断枠組みを提供すると主張している。主な貢献は以下の通りである。

1. 統合された診断ツール: 単純なルーティング統計や活性化パターンを超えて、エキスパートの専門化を分析するための体系的な方法を提供し、関数空間と表現空間を共同で特徴づける。
2. MoE の幾何学的解釈: 結果は、MoE 層が共有表現多様体のソフトな分割を実装するという見方を支持する。エキスパートは機能的に明確であるが、重複する部分多様体上で動作する。
3. メカニズム的洞察: この研究は、ルーティングの鋭さ（疎性）を専門化の度合いに対する直接的な制御メカニズムとして特定し、アーキテクチャ設計の選択を測定可能な幾何学的構造に結びつける。
4. モデル設計への示唆: 著者は、表現の類似性（例えば、活性化の重複）のみに基づくエキスパートの統合や剪定戦略は誤っている可能性があると示唆する。エキスパートは重複する表現を持ちながら異なる機能を実行し得るため、真の冗長性を正しく特定するには機能的指標（ヤコビアン整合性など）が必要となる可能性がある。

本論文は、その分析が局所的（一次ヤコビアン）かつ線形的（PCA）であり、大域的な非線形挙動や訓練中の幾何学の進化を捉えていないことを認めつつ、控えめな範囲を維持している。しかし、それは MoE モデルが計算をどのように組織化するかについての基礎的な幾何学的視点を確立している。