技術要約：改善された疎なMixture of ExpertsのためのSelective Sinkhorn Routing

問題提起

疎なMixture-of-Experts (SMoE) モデルは、トークンごとに一部のエキスパートのみを活性化させることで、計算効率を維持しながらモデルの容量をスケールアップできる手法を提供します。しかし、Softmaxベースのルーティングに依存する従来のSMoEモデルは、「ルーティング崩壊（routing collapse）」、すなわち少数のエキスパートが過剰に利用される一方で、他のエキスパートがほとんど利用されないという問題に直面することがよくあります。

既存の緩和戦略は、通常、以下の手法に依存しています：

1. 補助的な目的関数: 不均衡なエキスパート利用にペナルティを与えるロードバランシング損失やz-loss。これらは、目的関数の不一致、学習の不安定性、および複雑さの増大を招く可能性があります。
2. ノイズ付きゲーティング: 探索を促進するためにノイズを注入する方法であり、これは学習可能なパラメータやハイパーパラメータを追加します。
3. Sinkhornベースのルーティング: 最適輸送（OT）のためにSinkhornアルゴリズムを使用する従来の手法は、多くの場合、輸送マップを重みの割り当てではなく、エキスパートの選択（top-k）のためだけに用います。さらに、ゲーティング行列がSoftmax手法と同じように勾配ベースの学習を通じて直接最適化されないため、ルーティングの柔軟性を低下させたり、反復的なSinkhorn更新による大幅な学習オーバーヘッドが発生したりすることがあります。

手法

本論文では、軽量なメカニズムであるSelective Sinkhorn Routing (SSR) を提案します。これは、エントロピー正則化付きの最適輸送問題として、トークンからエキスパートへの割り当てを再定式化したものです。

コアとなる定式化

著者らは、$m$ 個のトークンを $n$ 個のエキスパートに割り当てる問題を、エントロピー正則化を伴う最大コスト最適輸送問題としてモデル化しています。目標は、特定の制約を満たしつつ、適合性を最大化する輸送計画 $\hat{\Pi}$ を見つけることです：

• 正値性: $\Pi > 0$。
• 行和: 各トークンはその全質量をエキスパートにルーティングする ($\Pi \mathbf{1}_n = \mathbf{1}_m$)。
• 列和: 各エキスパートは等しい期待総負荷を受け取る ($\Pi^\top \mathbf{1}_m = \frac{m}{n} \mathbf{1}_n$)。

コスト行列 $C$ はゲーティングスコア $S$ から導出されます。著者らは2つのバリアントを提案しています：

• 線形コスト: $C = S$。
• Softmaxコスト: $C_{i,:} = \text{softmax}(S_{i,:})$。これはSinkhorn反復中の数値オーバーフローを防ぐために使用されます。

輸送計画を単なる選択（top-k）にのみ使用する従来のOTベースの手法とは異なり、SSRは $\hat{\Pi}$ の値をルーティング重みの計算に直接使用します。各トークンについて、$\hat{\Pi}$ の対応する行の最大エントリーに基づいてトップ-$k$のエキスパートが選択され、重みはこれらのエントリーから正規化されます。

選択的適用とノイズ注入

課題として、Sinkhornルーティングは $W_g$ がOT目的関数から分離されているため、$W_g$ を更新しないことが特定されました。これに対処するため、SSRはハイブリッド学習戦略を採用しています：

• 選択的適用: 学習中、各MoEブロックは、小さな確率 $p$（例：0.001）でSinkhornルーティングを使用し、それ以外の場合は標準的なSoftmaxゲーティングを使用します。これにより、Softmaxゲーティングが $W_g$ を効果的に更新できる一方で、Sinkhornルーティングがバランスの取れたエキスパート利用シグナルを提供することを可能にします。
• ノイズ注入: エキスパートの崩壊を防ぎ、探索を促進するために、コスト行列にガウスノイズが追加されます ($\tilde{C} = C + \alpha_{noise} \cdot \epsilon$)。理論的分析（命題4.3）によれば、これによりすべてのエキスパートが非ゼロの選択確率を持つことが保証されます。

推論時の挙動

論文では、単一の入力に対するバランスの取れた制約（等しい列和など）を強制することは、入力の具体的な適合度に関わらず一様なルーティングを強いることになるため、逆効果であると主張しています。したがって、推論時にはSinkhornルーティングとノイズ注入の両方が無効化され、一貫した入力依存の予測を確保するために決定論的なSoftmaxルーティングへと回帰します。

主な貢献

1. 新しいルーティングフレームワーク: 補助的なバランシング損失なしに、バランスの取れたエキスパート利用を促進する、エントロピー正則化付き最適輸送と確率的ノイズ注入を統合した手法。
2. 理論的洞察: Sinkhornベースのルーティングとノイズ注入が、探索とバランシングを促進することで学習を助ける一方で、歪んだ割り当てを避けるために推論時には無効にすべきであることを示す証明。
3. 効率性と性能: SSRが、最先端のベースラインと比較して、学習効率、精度、および堅牢性を向上させることを示す広範な評価。

実験結果

著者らは、言語モデリング（WikiText-103, Enwik-8）および画像分類（ImageNet-1K, ImageNet-A, ImageNet-O, ImageNet-R）においてSSRを評価しました。

• 言語モデリング:
  • WikiText-103において、SSRの各バリアントは、Vanilla SMoEおよびロードバランシング損失、z-loss、またはノイズを使用するベースラインを一貫して上回りました。
  • SSR-L w/ noise は、Vanilla SMoE（35.550）に対して1.183減少したテスト・パープレキシティ（PPL）34.367を達成し、次点のベースラインを大幅に上回りました。
  • 極めて重要なことに、SSRの学習オーバーヘッドは最小限（0.33% – 0.65%）であり、フルSinkhornベースのSMoEが**72.47%**のオーバーヘッドを要したのと対照的です。
  • Momentum設定下において、SSRはVanilla SMoEの特徴的なノルム進化を維持しつつ性能を向上させましたが、他のバランシング手法は性能を低下させました。
• バイトレベルモデリング (Enwik-8):
  • SSR-S w/ noise は、最高のテストBits-per-Character (BPC) である1.128を達成し、Vanilla SMoEを0.010 BPC改善しました。
  • フルSinkhornベースのSMoEよりも2.28倍高速でした。
• ビジョンタスク:
  • ImageNet-1Kにおいて、SSR-L w/ noise はTop-1精度77.420を達成し、Vanilla SMoEに対して**+2.368**の利得を得ました。
  • また、ImageNet-OおよびImageNet-Rにおいても強力な堅牢性を示しました。
• アブレーション研究:
  • Sinkhornルーティングを適用する確率 $p$ は極めて重要であり、非常に小さな値（例：$10^{-3}$ または $10^{-4}$）でも、ルーターを過度に制約することなく、バランシング信号を提供するには十分でした。
  • ノイズ注入は、あらゆる設定において一貫して性能を向上させました。
  • 推論時にバランシングメカニズムを無効にすることが最良の結果をもたらし、理論的な主張を裏付けました。

意義と主張

本論文は、Selective Sinkhorn RoutingがSMoE設計における実用的かつ効果的なソリューションであることを主張しています。複雑な補助損失を、効率的で断続的なSinkhornベースの最適化に置き換えることで、SSRは以下を実現します：

• 学習の安定性の向上: 目的関数の不一致なしに、組み込みのエキスパートバランシングを実現。
• 高い精度: テキストおよびビジョンのベンチマーク全体で実証。
• 堅牢性: 入力の破損や敵対的な例に対する性能の強化。
• 効率性: 補助損失を用いる手法およびフルSinkhornルーティングの両方と比較して、最小限の計算オーバーヘッド。

著者らは、この手法が軽量であり、追加の学習可能パラメータを必要とせず（ノイズ・バリアントの場合）、エキスパート選択の柔軟性を維持しながら、学習フェーズにおけるバランスの取れた利用を保証することを強調しています。