スパース混合エキスパート(Sparse Mixture of Experts、Sparse MoE)は、複数の専門化されたニューラルネットワーク(エキスパートと呼ばれる)を、各入力に対して少数のエキスパートのみを活性化するゲーティング機構と組み合わせた機械学習アーキテクチャである。このアプローチは、すべての入力に対してすべてのパラメータを使用する稠密モデルとは対照的である。入力を関連するエキスパートに選択的にルーティングすることで、Sparse MoEは計算コストを比例的に増加させることなくモデル容量を増やし、大規模言語モデルやその他の深層学習システムのスケーリングにおける重要な技術となっている。
混合エキスパート(MoE)の概念は、1990年代初頭にアンサンブル学習の一形態として始まり、複数の学習器が問題空間を均質な領域に分割するものであった。ロバート・ジェイコブス、マイケル・ジョーダン、ジェフリー・ヒントンによる初期の研究では、局所エキスパートの適応的混合が導入され、各エキスパートが入力空間の異なる部分に特化した。その後、ハンプシャーとワイベルのメタPiネットワークがMoEを音声認識に適用し、6つの時間遅延ニューラルネットワークを訓練して異なる日本語話者の音素を分類した。これらの基盤システムは稠密重み付けを使用し、すべてのエキスパートがすべての出力に寄与したが、ゲーティング関数はその寄与を重み付けすることを学習した。
Sparse MoEは、稠密MoEの計算限界に対する実用的な解決策として登場した。稠密MoEでは、ゲーティング関数がすべてのエキスパート出力の加重和を計算するため、各入力に対してすべてのエキスパートを評価する必要がある。このオーバーヘッドを削減するために、トークンごとに少数のエキスパートのみを選択するスパース活性化が導入された。重要な革新は、入力に基づいて上位k個のエキスパート(通常k=1またはk=2)を選択する訓練可能なルーターであり、モデルがほとんどのエキスパートを完全にスキップできるようにする。このスパースルーティングは、2017年の論文「Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer」で、Googleのノーム・シェイザーらによって普及し、スパースMoEが数千のエキスパートにスケーリングしながら合理的な推論コストを維持できることを実証した。
アーキテクチャとルーティング
Sparse MoE層は、エキスパートネットワークのセット、ゲーティングまたはルーティング関数、負荷分散機構の3つの主要コンポーネントで構成される。各エキスパートは通常、フィードフォワードニューラルネットワークであるが、任意の微分可能な関数でもよい。ルーターは入力トークン表現を受け取り、学習された線形射影上のソフトマックスを介してエキスパート上の確率分布を生成する。最高確率を持つ上位k個のエキスパートのみが活性化され、その出力は正規化されたルーター重みを使用して結合される。
ルーティング関数は、性能と効率の両方にとって重要である。トークンを適切なエキスパートに割り当てることを学習しつつ、エキスパートがほぼ均等に使用されるようにして、過少利用を避ける必要がある。一般的な問題は「ルーター崩壊」であり、ルーターが常に同じ少数のエキスパートを選択し、特化の目的を無効にする。この問題に対処するため、現代の実装では、不均一なエキスパート使用を罰する補助負荷分散損失を使用する。例えば、2021年にGoogleが導入したSwitch Transformerは、トークンごとに1つのエキスパートのみを選択する簡略化されたルーティング戦略を使用し、エキスパート間の均一なトークン分布を促進する負荷分散損失と組み合わせている。
訓練ダイナミクス
Sparse MoEモデルの訓練は、稠密モデルと比較して特有の課題を提示する。ルーターとエキスパートは共同で訓練する必要があるが、エキスパートの離散的な選択により、ルーティング決定は微分不可能である。ほとんどの実装では、訓練中にソフトマックス近似を使用し、ルーター重みは連続的であり、上位k選択は推論中にのみ適用される。これにより、勾配がルーターを通じて流れることが可能になるが、訓練と推論の動作の間に不一致が生じる。
もう一つの課題は、訓練中の安定性である。Sparse MoEモデルは、特にルーターが過度に自信を持つ場合や、エキスパートが非常に異なる数のトークンを受け取る場合に、訓練不安定性を示すことがある。ランダムに選択されたエキスパートを訓練中にドロップするエキスパートドロップアウトや、大きなルーターロジットを罰するルーターz損失などの技術が、訓練を安定化させるのに役立つ。2023年にMistral AIがリリースしたMixtralモデルは、Sparse MoEが層ごとに8つのエキスパートを持ち、トークンごとに2つだけを選択しながら、稠密モデルに匹敵する総パラメータ数を維持しつつ、最先端の性能を達成できることを実証した。
大規模言語モデルへの応用
Sparse MoEは、計算を比例的に増加させずにパラメータをスケーリングできるため、大規模言語モデル(LLM)の標準技術となっている。Switch Transformerは、1.6兆パラメータのスパースMoEモデルが効率的に訓練でき、同様の計算予算の稠密モデルよりも優れた性能を達成することを示した。このアプローチは、いくつかの主要なAI研究機関で採用されている。GoogleのGShardアーキテクチャは、スパースMoEを機械翻訳に適用し、ソース言語に基づいてトークンをエキスパート間でルーティングした。最近では、Mixtral 8x7BやDeepSeek-MoEなどのモデルが、スパースMoEがオープンソースLLMで強力な性能を提供できることを実証している。
Sparse MoEの計算効率は、推論において特に価値があり、各トークンに対してエキスパートの一部のみをメモリにロードする必要がある。これにより、数千億のパラメータを持つモデルが、コンシューマーGPUなどの限られたメモリのハードウェアで実行できる。しかし、すべてのエキスパート重みを格納するメモリフットプリントは依然として大きく、異なるエキスパートを異なるデバイスに配置するエキスパート並列処理などの技術が、モデルを複数のアクセラレータに分散させるためによく使用される。
稠密モデルとの比較
Sparse MoEモデルは、パラメータと計算のトレードオフにおいて稠密モデルと根本的に異なる。Nパラメータの稠密モデルは、各入力に対してすべてのNパラメータを使用するため、計算コストはパラメータ数に比例して増加する。N総パラメータを持つが、トークンごとにk個のアクティブエキスパートのみを持つスパースMoEモデルは、トークンごとに約k/Nのパラメータを使用し、同じ計算予算の稠密モデルよりもはるかに多くのパラメータを持つことができる。これにより、スパースMoEモデルは、推論レイテンシを増加させることなく、より多くの知識を獲得し、より多様なタスクを処理できる。
しかし、スパースMoEモデルが普遍的に優れているわけではない。エキスパート数、kの値、負荷分散戦略の慎重な調整が必要である。ルーターは小さなオーバーヘッドを追加し、離散的なルーティングは予測が難しい特化につながる可能性がある。稠密モデルは訓練と展開がより簡単であり、多くのタスクでは、稠密モデルとスパースモデルの性能差は小さい。稠密アーキテクチャとスパースアーキテクチャの選択は、計算、メモリ、目標性能の特定の制約に依存する。
ハードウェアと実装の考慮事項
Sparse MoEを効率的に実装するには、専門的なハードウェアサポートとソフトウェアフレームワークが必要である。ルーティング操作は不規則なメモリアクセスパターンを導入し、異なるトークンが異なるエキスパートに送られる可能性がある。これにより、GPU上で負荷不均衡が発生し、一部のデバイスが多くのトークンを処理する一方で、他のデバイスがアイドル状態になることがある。過剰なトークンを破棄するトークンドロップや、エキスパートをデバイス間で分散するエキスパート並列処理などの技術が、これらの問題を軽減するのに役立つ。主要なクラウドプロバイダーであるGoogle Cloud、Amazon Web Services、Microsoft Azureは、大規模MoE訓練と推論に最適化されたインフラストラクチャを提供している。
ハードウェアベンダーもスパースMoE向けの最適化を始めている。NVIDIA GPUはスパーステンソル操作をサポートし、AMDとIntelはスパース計算を処理する機能を備えたアクセラレータを開発している。Google Cloudは、MoE層に関わる行列乗算に適したTensor Processing Unit(TPU)を提供している。PyTorchやTensorFlowなどのソフトウェアフレームワークは、MoE層のネイティブサポートを追加し、負荷分散ユーティリティや分散訓練プリミティブを含んでいる。
将来の方向性
Sparse MoEに関する研究は進化を続けており、いくつかの有望な方向性がある。1つの分野は、より良い特化と負荷分散を達成するためのルーティングアルゴリズムの改善である。もう1つは、エキスパート間でパラメータを共有したり、低ランク近似を使用したりして、多くのエキスパートを格納するメモリフットプリントを削減する方法の開発である。また、言語モデルを超えて、コンピュータビジョンや強化学習にスパースMoEを適用することへの関心もあり、この技術が多様な入力を処理するモデルのスケーリングに役立つ可能性がある。
Sparse MoEと、Transformer (architecture)の変種やResidual Network (ResNet)設計などの他のアーキテクチャ革新との統合は、活発な研究分野である。モデルが成長し続けるにつれて、スパース活性化は大規模AIシステムを実現可能にする上でますます重要な役割を果たす可能性が高い。エキスパート特化と一般化のバランス、およびルーティングと訓練ダイナミクスの間の相互作用は、次世代のLarge language modelアーキテクチャを形成する未解決の研究課題である。
結論
Sparse MoEは、機械学習アーキテクチャにおける重要な進歩を表し、計算効率を維持しながらモデルを前例のない規模にスケーリングすることを可能にする。アンサンブル学習の起源から最先端のLLMでの現在の役割まで、この技術は多用途で効果的であることが証明されている。入力ごとにエキスパートのサブセットのみを活性化することで、Sparse MoEは容量と計算の間の有利なトレードオフを達成し、現代のAIシステムの基盤となっている。ハードウェアとソフトウェアが進化し続けるにつれて、Sparse MoEは、ますます高性能で効率的なモデルを構築するための重要なツールであり続けると期待されている。