技術要約：疎なMixture-of-Expertsにおける不連続性の幾何学的および確率論的解析

問題提起

疎なMixture-of-Experts（SMoE）アーキテクチャは、Top-$k$エキスパート選択による条件付きルーティングを可能にすることで、計算コストを制御しつつ大規模なパラメータ数を実現し、最先端の言語および視覚モデルをスケールアップさせるための中心的な技術となっています。しかし、このメカニズムは入力・出力写像に固有の不連続性を導入します。具体的には、Top-$k$ゲーティングメカニズムは、入力を固定された$k$個のエキスパートが活性化する多面体領域へと分割します。これらの領域の境界において、ゲーティングスコアがタイ（同値）となる場合、入力への微小な摂動が、完全に異なるエキスパートセットへの「ハードスイッチ」を引き起こし、結果として制御不能な出力変動をもたらします。

先行研究では、これらの不連続性が認識されており、微分可能なルーティングの代替案（例：Soft MoE, SMEAR, ReMoE）が提案されてきました。しかし、これらの手法は、自己回帰生成に必要な因果構造を損なうか、あるいは根本的に異なるゲーティングメカニズムを用いてゼロから再学習を行う必要があります。標準的なTop-$k$ SMoEにおける不連続性の幾何学的構造と確率的挙動に関する体系的な理論的解析は、これまで欠落していました。

手法

本論文は、幾何学的測度論と確率過程という2つの主要な観点からこの問題にアプローチしています。

1. 不連続性の幾何学的解析

著者らは、スイッチングイベントにおけるタイ（同値）が発生したエキスパートの数によって、不連続性を**次数（order）**として分類しています。

• 次数-1（Order-1）: 活性化している$k$番目のエキスパートと、非活性化している$(k+1)$番目のエキスパートとの間のペアワイズなタイ。
• 次数-$n$（Order-$n$）: しきい値において$n+1$個のエキスパートが同時にタイとなる状態。

測度論的なスライシング議論を用いることで、本論文は有界な入力空間内における、これらの不連続集合の$\epsilon$-厚層近傍（$\epsilon$-thickened neighborhoods）の漸近的な体積推定値を導出しています。解析は、ユークリッド的な$\epsilon$-厚層化と、より計算的に扱いやすい（ロジット差に基づく）$\ell_\infty, \epsilon$-厚層化の両方をカバーしています。核心となる幾何学的知見は、すべての不連続集合はルベーグ測度ゼロであるが、その$\epsilon$-厚層は無視できない体積を持つということです。決定的なことは、入力空間の半径$R$が増大するか、あるいは$\epsilon$が減少するにつれて、高次の不連続近傍の体積は低次のものに対して漸近的に消失することです。

2. 確率的解析

本論文は、入力空間におけるランダムな摂動を、伊藤拡散過程（Itô diffusion process）としてモデル化しています。著者らは、2つの主要な確率的特性を分析しています。

• 到達時間（Hitting Times）: あるTop-$k$セル内から始まるランダムな拡散パスは、有限時間内に不連続境界に到達することがほとんど確実に（almost surely）起こることを証明しています。さらに、最初の接触（first hit）は、ほとんど確実に次数-1の不連続において発生します。高次の不連続（余次元 $\ge 2$）は、次元 $d \ge 2$ におけるブラウン運動が固定点をほとんど確実に通過しないのと同様に、拡散過程にとって（最初の接触確率という観点からは）事実上「不可視」です。
• 滞在時間（Occupation Times）: 著者らは、ランダムなパスが異なる次数の不連続の$\epsilon$-近傍に費やす期待時間の境界を導出しています。これらの境界は、拡散過程が次数-1の境界付近で時間の大部分を過ごし、高次の不連続に対する滞在時間は急速に減衰することを裏付けています。

3. 提案手法：$\ell_\infty, \epsilon$ ローカル平滑化

入力が次数-1の不連続の近くに存在する可能性が高いという知見に基づき、既存のSMoEモデルの基礎となるTop-$k$ルーティング・ロジックを変更することなく適用できる、単純な平滑化メカニズムを提案しています。

• メカニズム: この手法は、非Top-$k$のロジットが$k$番目のロジットの範囲内にある（不連続集合の$\ell_\infty, \epsilon$-厚層内にある）入力を特定します。
• 平滑化: これらの特定の入力に対して、境界付近のエキスパートのロジットに「ログ・スマートステップ（log-smoothstep）」関数を適用します。これにより、境界付近のエキスパートをソフトに組み込み、ハードな切り替えではなく連続的な遷移を実現します。
• 適応的$\epsilon$: 訓練中に$\epsilon$を適応的に調整するための境界損失（boundary loss）を導入し、平均的な活性化エキスパート数が目標予算（例：$k + 0.5$）内に収まるようにします。

主な貢献

1. 漸近的測度境界: 本論文は、$\epsilon$-厚層化された次数-1の不連続集合が入力空間の体積を支配し、高次の不連続集合は消失していく相対的な体積しか持たないことを確立しました。これは、ユークリッド距離および$\ell_\infty$計量の両方において成立します。
2. 確率的な到達保証: ランダムな摂動は有限時間内に不連続に遭遇すること、そして最初の遭遇はほとんど確実に次数-1の不連続であることを証明しました。
3. 滞在時間の境界: 小さな$\epsilon$の領域において、高次の不連続付近での滞在時間は無視できるほど小さいことを示す明示的な境界を導出しました。
4. 実用的な平滑化メカニズム: Top-$k$ SMoEの疎性と因果構造を維持しつつ、不連続付近で連続性を強制する、直接的かつ計算効率の高い平滑化手法を提案しました。

実験結果

提案されたSmoothSMoEは、言語および視覚タスクにおいて評価されました。

• 連続性の検証: CIFAR-10およびWikiText-103を用いた可視化により、バニラSMoEが境界で鋭いジャンプを示す一方で、SmoothSMoEは、入力の摂動がゼロに近づくにつれて出力の差が消失する連続的な出力写像を生成することが示されました。
• 言語モデリング: WikiText-103およびEnWiki-8において、SmoothSMoEはパープレキシティ（PPL）およびBits-per-Character（BPC）の両方において、バニラSMoEおよびReMoE（完全微分可能な代替案）の両方を上回りました。特に、攻撃を受けたデータセットに対する堅牢性が向上しました。
• GLUEベンチマーク: 自然言語理解タスク（CoLA, MRSR, MNLIなど）において、SmoothSMOEは比較対象のベースラインの中で最高水準の性能を達成し、次点のメソッドに対して平均約0.5%のスコア向上を実現しました。
• 視覚タスク: DomainBedベンチマークにおいて、平滑化されたバリアントであるSmoothGMoEは、ベースラインのGMoEに対し、特にTerraIncognitaデータセットにおいて一貫した精度向上を示しました。
• 大規模モデル: 5.6Bパラメータの視覚言語モデルを用いた実験により、本手法が効果的にスケールすることが確認されました。アニールド・バリアント（推論時に平滑化を除去したもの）は、ハードなTop-$k$推論を維持しながら、フル平滑化モデルと同等の性能を達成しました。

意義と主張

本論文は、その主要な貢献は、SMoEの不連続性を厳密に特徴付ける理論的解析であり、それが低次の不連続性が不安定性の主要な原因であることを明らかにしている点にあると主張しています。この知見は、グローバルな微分可能ルーティングの計算的・構造的なコストを回避しつつ、標的を絞った平滑化戦略を動機付けています。

著者らは、自らの手法が以下のことを達成すると断言しています。

• 連続性の強制: SMoE写像の不連続付近における連続性を数学的に保証します。
• 疎性の維持: Soft MoEやトークン・マージングのアプローチとは異なり、自己回帰生成に不可欠な疎で因果的なルーティング構造を維持します。
• オーバーヘッドの最小化: 理論的境界および実験結果は、高次の不連続（多くの追加エキスパートを活性化させる必要があるもの）は統計的に無視できるため、計算オーバーヘッドが小さいことを示唆しています。
• 性能の向上: 平滑化は正則化として機能し、基礎となるTop-$k$アーキテクチャを変更したりゼロから再学習したりすることなく、訓練ダイナミクスを安定させ、多様なタスクにおける経験的な性能を向上させます。

本論文は、敵対的または構造的な摂動は想定されているランダム拡散モデルから逸脱する可能性があるものの、提示された理論的枠組みは、Top-$k$ルーティングに内在する不安定性を理解し、軽減するための強固な基盤を提供するものであると結論付けています。