スパースMoEルーティングは、Mixture of expertsニューラルネットワークにおいて、各入力トークンを処理するエキスパートサブネットワークを決定するメカニズムである。スパース混合エキスパート層では、ルーターまたはゲーティング関数が入力を評価し、より大きなプールから少数のエキスパート（通常は1つまたは2つ）のみを選択する。これは、すべてのエキスパートがすべての出力に寄与するデンス混合エキスパートとは対照的である。ルーティングの決定はトークンごとに行われ、入力シーケンスの異なる部分が異なるエキスパートによって処理されることを可能にする。この設計により、モデルの総パラメータ数は増加するが、トークンあたりの計算コストはほぼ一定に保たれる。なぜなら、フォワードパスとバックワードパスの間、選択されたエキスパートのみが活性化されるからである。
この概念は1990年代の混合エキスパートに関する初期の研究から生まれたが、スパースルーティングが実用的に重要になったのは2010年代から2020年代にかけてであり、デンスなTransformer (architecture)アーキテクチャの限界を超えて大規模言語モデルをスケールさせる方法として注目された。トークンを専門化されたエキスパートにルーティングすることで、モデルはすべてのパラメータをすべての入力に対して活性化することなく、データ内の多様なパターンを捕捉できる。スパースMoEルーティングは現在、Google DeepMind、OpenAI、その他の研究機関によって開発されたいくつかの著名な本番モデルの中核的なコンポーネントとなっている。
歴史的起源
混合エキスパートの基礎的なアイデアは1990年代初頭に遡る。ロバート・A・ジェイコブス、マイケル・I・ジョーダン、スティーブン・J・ノーラン、ジェフリー・E・ヒントンを含む研究者らは、1991年に適応的混合ローカルエキスパートアーキテクチャを導入し、ゲーティングネットワークが複数のエキスパートネットワークの出力に重み付けを行った。ほぼ同時期に、ジョン・ハンプシャーとアレックス・ワイベルによるメタ・パイネットワークは、音声音素分類に同様の重み付けスキームを適用し、6人の日本語話者からのデータを用いて6つの時間遅延ニューラルネットワークを訓練した。これらの初期システムはデンスな重み付けを使用しており、すべてのエキスパートが各出力に寄与したが、エキスパート関数、ゲーティング関数、重み付き組み合わせという中核的なコンポーネントを確立した。
エキスパートのサブセットのみを活性化するスパースルーティングは、1990年代後半から2000年代初頭に探求されたが、深層学習の台頭とともに注目を集めた。2017年、Google DeepMindの研究者らは、スパースゲーティング混合エキスパート層を発表し、各入力に対してトップkのエキスパートを選択する訓練可能なゲーティングネットワークを導入した。この研究は、スパース活性化が計算量の比例的な増加なしにモデル容量を劇的に増加させ得ることを実証し、後の大規模アプリケーションへの道を開いた。
ルーティングメカニズム
スパースMoE層では、ルーティングメカニズムはゲーティング関数で構成され、通常は線形層とそれに続くエキスパート次元上のソフトマックスからなる。入力トークン表現\(x\)に対して、ルーターは\(n\)個のエキスパートそれぞれのスコアを計算し、多くの場合\(w(x)_i = \text{softmax}(W_g x)_i\)として表される。ここで\(W_g\)は学習可能な重み行列である。ルーターはその後、スコアが最も高いトップkのエキスパートを選択する。ここでkはハイパーパラメータであり、通常は1または2に設定される。層の出力は、選択されたエキスパートの出力の重み付き和であり、重みは選択されたエキスパート間で正規化される。
このトークンごとの選択により、モデルは専門化できる。異なるエキスパートは、異なる言語、ドメイン、構文パターンなど、異なるタイプの入力を処理することを学習できる。ルーティングの決定は推論中は決定的であるが、訓練中はルーターをエキスパートと共同で訓練する必要がある。トップk選択は非微分可能であるため、研究者らはストレートスルー推定器やノイズ付きトップkゲーティングなどの技術を使用する。これは、訓練中にロジットに学習可能なノイズを追加して、探索と負荷分散を促進する。
負荷分散と補助損失
スパースMoEルーティングにおける主要な課題は、エキスパートがほぼ均等に使用されることを保証することである。介入がなければ、ルーターは少数のエキスパートがほとんどのトークンを受け取り、他は十分に活用されない状態に収束する可能性があり、実効容量が減少する。これに対処するため、モデルはエキスパート間の不均衡なトークン分布にペナルティを課す補助的な負荷分散損失を組み込む。Google (AI)の研究者らによって2021年に導入されたSwitch Transformerで一般的なアプローチは、各エキスパートにルーティングされたトークンの割合に、そのエキスパートの平均ルーター確率を掛けたものに比例する損失項を追加するものである。これにより、ルーターがトークンをより均一に分散するよう促される。
もう一つの技術はエキスパート容量であり、これは特定のバッチ内で各エキスパートが処理できるトークン数を制限する。エキスパートが容量に達した場合、超過トークンはドロップされるか、残差接続にルーティングされる。これにより、単一のエキスパートがボトルネックになるのを防ぎ、予測可能な計算を保証する。DeepSeek-V3（2024年）で使用されている補助損失なしルーティングなど、より最近の方法は、明示的な補助損失なしに負荷を均衡させるために動的バイアス調整を使用する。
スケーリングと効率性
スパースMoEルーティングにより、デンスモデルが許容する範囲をはるかに超えてモデルパラメータをスケールできる。例えば、Switch Transformerは最大1.6兆パラメータのモデルを導入したが、各トークンはそのごく一部のみを活性化した。同様に、2023年にMistral AIがリリースしたMixtral 8x7Bは、層ごとに8つのエキスパートを使用し、トークンごとに2つが活性化され、トークンあたりの計算量を抑えながら、より大きなデンスモデルに匹敵する性能を達成した。GoogleのGShard（2020年）はスパースMoEを機械翻訳に適用し、後のGLaM（2021年）やPaLM（2022年）などのモデルも効率性のためにMoE層を組み込んだ。
効率性の利点は、トークンあたりの計算コストが総パラメータ数ではなく活性化されたエキスパート数に比例するという事実から生じる。これにより、モデルはより多くのパラメータを持つことができ、精度を向上させることができる一方、推論と訓練のコストを管理可能に保つことができる。しかし、スパースMoEモデルはすべてのエキスパートパラメータを格納するためにより多くのメモリを必要とし、トークンを異なるデバイスに存在する可能性のあるエキスパートにルーティングする必要があるため、分散訓練において通信オーバーヘッドを導入する。
大規模言語モデルへの応用
スパースMoEルーティングは、大規模言語モデルにおける標準的な技術となっている。多くのオープンウェイトおよびプロプライエタリモデルが、品質とコストのバランスを取るためにMoE層を使用している。例えば、Mistral AIのMixtral 8x7BとMixtral 8x22B、AlibabaのQwen1.5-MoE、DeepSeekのDeepSeek-V2およびV3はすべてスパースルーティングを採用している。プロプライエタリな分野では、OpenAIのGPT-4が混合エキスパートアーキテクチャを使用していると広く報じられているが、同社は詳細を確認していない。AnthropicのClaudeモデルやGoogleのGeminiモデルも、公開声明や特許に基づいてMoE層を組み込んでいる可能性が高い。
これらのモデルは、ルーティングを使用してドメイン、言語、推論タスクにわたってエキスパートを専門化する。例えば、一部のエキスパートは数学的推論を処理し、他はコード生成や多言語テキストに焦点を当てる場合がある。ルーターは入力コンテキストに基づいてトークンを適切なエキスパートに導くことを学習し、すべてのエキスパートがすべてのトークンを処理する必要なしに全体的な性能を向上させる。
訓練の課題
スパースMoEモデルの訓練は、負荷分散以外にも固有の課題を提示する。ルーターは不安定性に悩まされる可能性があり、パラメータの小さな変化がルーティング決定の大きな変化につながり、訓練の発散を引き起こす。ノイズ付きトップkゲーティングやルーターへのドロップアウトなどの技術は、訓練を安定させるのに役立つ。さらに、エキスパートが早期に専門化しすぎて、モデルが一般化に失敗する一種のモード崩壊につながる可能性がある。研究者らは、エキスパートドロップアウトやルーティング正則化などの方法を提案して、これを緩和している。
もう一つの課題は、分散訓練における通信オーバーヘッドである。典型的な設定では、各エキスパートは異なるデバイスに配置され、トークンは元のデバイスから選択されたエキスパートをホストするデバイスに送信されなければならない。この全対全通信は、特に大規模なバッチサイズの場合、ボトルネックになる可能性がある。MegatronやDeepSpeedライブラリなどの効率的な実装は、最適化された通信スケジュールを使用してこのオーバーヘッドを削減する。
最近の開発
最近の研究は、ルーティング品質の向上とMoEモデルのコスト削減に焦点を当てている。Mixture-of-Depths（2024年）アプローチは、スパース性のアイデアを深さ次元に拡張し、トークンを異なるエキスパートにルーティングするだけでなく、層を完全にスキップすることも可能にする。もう一つの方向性は、きめ細かいエキスパート分割であり、エキスパートをより小さく、より多くすることで、より細かい専門化を可能にする。例えば、DeepSeek-V3は層ごとに256のエキスパートを使用し、そのうち8つだけが活性化され、訓練コストを削減しながら最先端の性能を達成している。
ハードウェアとソフトウェアの共同設計も進んでいる。NVIDIAやAMDなどの企業は、MoE推論用にアクセラレータを最適化しており、vLLMやTensorRT-LLMなどのフレームワークは効率的なMoEサービングをサポートしている。エッジ展開への関心の高まりにより、エキスパートプルーニングや量子化など、限られたメモリのデバイスに適合させるためのMoEモデル圧縮に関する研究が行われている。
将来の方向性
スパースMoEルーティングは依然として活発な研究分野である。未解決の疑問には、エキスパートの最適数とルーティングの粒度を自動的に決定する方法、ルーティングをより解釈可能にする方法、MoEをModel Pruningや量子化などの他の効率化技術と組み合わせる方法などがある。モデルがスケールし続けるにつれて、スパースルーティングは能力と計算コストのバランスを取る上で中心的な役割を果たし、法外なリソース要求なしにより強力なAIシステムを可能にするだろう。