Mixture of Experts（MoE）におけるルーティングは、現代の機械学習アーキテクチャにおける中核的なメカニズムであり、各入力トークンを専門化されたエキスパートネットワークのサブセットに動的に割り当てる。すべての入力を処理するために全パラメータを活性化する代わりに、ルーター（ゲーティングネットワーク）はトークンを最も関連性の高いエキスパートに導くことを学習し、モデル容量と計算効率のバランスを取る。このアプローチにより、大規模言語モデルは、ネットワークの一部のみがトークンごとに活性化されるため、推論コストを管理可能に保ちながら、数兆パラメータにスケールできる。
この概念は、1990年代初頭の適応的局所エキスパート混合に関する研究に由来するが、2010年代に深層学習におけるスパースゲート型MoE層の導入により注目を集めた。Google DeepMindやOpenAIのモデルで使用されるような現代の実装は、エキスパート上の確率分布を出力する学習済みルーティング関数に依存し、バランスの取れた使用を確保するために補助損失を組み込むことが多い。ルーティングは、最先端の生成AIシステムにおける重要な設計選択となり、トレーニング効率と最終的なモデル品質の両方に影響を与えている。
歴史的発展
入力の特殊化されたコンポーネントへのルーティングのアイデアは、深層学習以前に遡り、アンサンブル法やモジュール型ニューラルネットワークにルーツを持つ。1991年、マイケル・ジョーダンとロバート・ジェイコブスは、ゲーティングネットワークが複数のフィードフォワードネットワークの出力を重み付けする階層的エキスパート混合に関する基礎的な研究を発表した。この初期の枠組みは、入力空間を分割して異なるエキスパートが異なる領域を処理することを学習するという中核原理を確立した。
2017年、Google DeepMind（当時のGoogle Brain）の研究者らは、ノーム・シャジールが主導する論文でスパースゲート型MoE層を導入した。この設計は、ソフトマックスゲーティング関数を使用して各トークンのトップkエキスパートを選択し、kは通常1または2に設定された。著者らは、MoE層が数百のエキスパートにスケールでき、言語モデリングと機械翻訳ベンチマークで最先端の結果を達成しつつ、サンプルあたりの計算コストを削減できることを実証した。この研究は、後のアーキテクチャ、例えばトークンごとに1つのエキスパートのみを選択するようにルーティングを簡素化したSwitch Transformer（2021年）や、大規模並列トレーニングのためのGShardフレームワークに直接影響を与えた。
ルーティングメカニズム
ルーティング関数はいくつかのタイプに分類できる。最も一般的なのはトークン選択ルーティングであり、各トークンが学習済みスコアに基づいて独立にトップkエキスパートを選択する。これは、トークンの隠れ状態の線形射影と、それに続くエキスパートインデックス上のソフトマックスとして実装される。選択されたエキスパートはトークンを処理し、その出力はルーティング確率で重み付けされて合計される。
代替案はエキスパート選択ルーティングであり、各エキスパートがバッチからトップkトークンを選択し、すべてのエキスパートが最小限の負荷を受け取ることを保証する。このアプローチはSwitch Transformerで普及し、後にMixtralなどのモデルで洗練され、負荷不均衡の問題に対処するが、トークンとエキスパートの割り当ての慎重な処理が必要である。
他の変種には、学習済みパラメータなしでトークンを割り当てる決定論的ハッシュ関数を使用するハッシュベースのルーティングや、2レベルのゲーティングシステムが最初にエキスパートのグループを選択し、次にそのグループ内の特定のエキスパートを選択する階層的ルーティングがある。各方法は、柔軟性、計算オーバーヘッド、トレーニング安定性の間でトレードオフを行う。
負荷分散と補助損失
MoEルーティングにおける主要な課題は負荷不均衡である。ルーターが常に少数の人気エキスパートを選択し、他を未トレーニングのままにすることを学習する可能性がある。これを軽減するため、ほとんどの実装は、エキスパート間のトークンの不均一な分布を罰する補助負荷分散損失を追加する。2017年のシャジール論文で導入された標準的な定式化は、各エキスパートにルーティングされたトークンの割合を計算し、クロスエントロピー項を介して均一性を促進する。
Switch Transformerの負荷分散損失などのより最近のアプローチは、ルーターの平均確率とエキスパートごとの実際のトークン数の間のドット積を乗算するより単純な係数を使用する。GShardなどの一部のシステムは、各エキスパートが処理できるトークン数を制限する容量係数を採用し、ルーターに作業を分散させることを強制する。これらの技術は、安定したトレーニングと、一部のエキスパートが死重となるエキスパート崩壊の防止に不可欠である。
アーキテクチャ統合
ルーティングは通常、トランスフォーマーブロック内で適用され、フィードフォワードネットワーク（FFN）をMoE層に置き換える。標準的なトランスフォーマーでは、各トークンはマルチヘッドアテンションサブレイヤーとそれに続く位置ごとのFFNを通過する。MoEトランスフォーマーでは、FFNはそれぞれ独自のパラメータを持つエキスパートFFNのセットと、どのエキスパートを活性化するかを選択するルーターに置き換えられる。
この統合により、モデルはトークンごとのアクティブパラメータ数を一定に保ちながら、総パラメータ数を大きくできる。例えば、各10億パラメータの64エキスパートを持つモデルは、総計640億パラメータを持つが、k=2の場合、トークンごとに20億のみを活性化する。この特性は、Anthropicや他の研究所のMixture of Expertsモデルに見られるように、数兆パラメータモデルへのスケーリングに重要である。
ルーター自体は小さなニューラルネットワークであり、多くの場合、単一の線形層とそれに続くソフトマックスで構成され、エキスパートと共同でトレーニングされる独自のパラメータを持つ。一部のアーキテクチャは、異なる層やアテンションヘッドに別々のルーターを使用し、最近の研究では、トークンタイプや位置に基づいて適応する学習済みルーティングポリシーを探求している。
トレーニングダイナミクス
ルーティングを使用したMoEモデルのトレーニングは、独自の課題を導入する。ルーターの離散的な決定（トップkエキスパートの選択）は微分不可能であり、勾配は選択されたエキスパートの出力をルーティング確率で重み付けしたものにのみ流れる。これにより、移動ターゲット問題が生じる。エキスパートが改善するにつれて、ルーターの好みが変化し、不安定性を引き起こす可能性がある。
これに対処するため、研究者は、トレーニング中にルーティングロジットにガウスノイズを追加して探索を促進するノイズ付きトップkゲーティングなどの技術を使用する。別のアプローチは、ルーティング決定にストレートスルー推定器を使用し、選択をハード割り当てとして扱うが、確率重みを通じて勾配を渡す。さらに、一部の方法は、トレーニング中にソフトマックスの温度をアニーリングして、ルーティング決定を徐々にシャープにする。
負荷分散損失は通常、主要なタスク損失を支配しないように小さな係数（例：0.01）で重み付けされる。実際には、MoEモデルのトレーニングには慎重なハイパーパラメータ調整が必要であり、多くのシステムは負荷分散とルーター信頼性の両方に補助損失を採用している。
大規模言語モデルへの応用
MoEにおけるルーティングは、大規模言語モデルにおける標準的な技術となっている。注目すべき例は以下の通り：
- GoogleのSwitch Transformer（2021年）は、スパースMoEアーキテクチャで1.6兆パラメータにスケールした。
- GShard（2020年）は、数千のTPUコアにわたるMoEモデルの効率的なトレーニングを実証した。
- Mistral AIのMixtral 8x7B（2023年）は、8エキスパートとトップ2ルーティングを使用し、より大きな高密度モデルに匹敵するパフォーマンスを達成した。
- DeepSeek-V3（2024年）は、256エキスパートとトップ8ルーティングを備えた細粒度MoEを採用している。
- OpenAIやAnthropicのいくつかのモデルはMoE層を使用していると報告されているが、正確な詳細は多くの場合プロプライエタリである。
これらのモデルは、ルーティングがコスト効率の高いスケーリングを可能にすることを実証しており、総パラメータが増加してもトークンあたりの計算は一定のままである。これにより、MoEは、Amazon Web Services、Microsoft Azure、Google Cloudなどのクラウドプラットフォームを含む本番環境で大規模モデルを提供するための好ましい選択となっている。
効率性とハードウェアの考慮事項
ルーティングは、分散トレーニングと推論に通信オーバーヘッドを導入する。トークンは適切なエキスパートに送信される必要があり、そのエキスパートは異なるデバイス上にある可能性がある。このオールツーオール通信パターンは、特に多くのエキスパートがある場合にボトルネックになる可能性がある。NVIDIA（明示的にはリストされていないが、暗示されている）やAMDなどのハードウェアベンダーは、MoE操作に最適化されたカーネルを開発しており、GroqやSambaNovaの特殊なAIアクセラレータは、スパース活性化パターンを効率的に処理するように設計されている。
通信を削減するため、一部のシステムはエキスパート並列処理を使用し、エキスパートをデバイス間で複製するか、データ並列処理とエキスパート並列処理を組み合わせたハイブリッドアプローチを使用する。ルーティングの粒度（トークンレベル対ブロックレベル）の選択も効率に影響する。トークンレベルのルーティングはより細かい制御を提供するが、より多くの通信を伴い、ブロックレベルのルーティングはトークンをグループ化してオーバーヘッドを削減する。
課題と将来の方向性
その成功にもかかわらず、MoEにおけるルーティングはいくつかの未解決の問題に直面している。1つは、エキスパートの特殊化と一般化の間のトレードオフである。過度に特殊化されたエキスパートは、新しいタスクにうまく転移しない可能性がある。もう1つは、ルーターが新しいデータ分布に適応する必要があるため、MoEモデルのファインチューニングが難しいことである。トークンごとにアクティブなエキスパートの数が変化する適応的ルーティングや、現在のトークン以外のコンテキストを考慮する学習済みルーティングポリシーに関する研究が進行中である。
将来の方向性には、モデルプルーニングや量子化などの他の効率技術とルーティングを組み合わせることや、マルチモーダル入力を処理できるルーターの開発が含まれる。モデルがスケールし続けるにつれて、ルーティングは容量と計算のバランスを取るための重要なメカニズムであり続け、言語を超えてビジョンや強化学習への潜在的な応用がある。
関連項目
参考文献
- Shazeer, N., et al. (2017). Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer.
- Fedus, W., et al. (2021). Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity.
- Lepikhin, D., et al. (2020). GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding.
- Jiang, A. Q., et al. (2023). Mixtral of Experts.
- DeepSeek-AI (2024). DeepSeek-V3 Technical Report.