Mixture of Experts
Mixture of Experts(MoE)は、複数の専門化されたエキスパート・サブネットワークと、各入力トークンごとにそのうちの少数だけを有効化するルーターを使用するニューラルネットワークアーキテクチャです。これにより、各トークンあたりの実際の計算量を少数のエキスパートに限定しながら、モデルの総パラメータ数を大幅に増やすことができます。
MoEは、1つの大規模なニューラルネットワークを複数のエキスパート・サブネットワークに置き換え、ルーターを使ってトークンごとに少数のエキスパートだけを選択・有効化するスパース設計です。
総パラメータ数は大きい一方で、各トークンは選ばれたエキスパートだけで処理されるため、同じ計算コストではるかに高いモデル容量を実現できます。
重要なのは、ルーターのtop-kルーティングと、エキスパート間の負荷分散であり、トークンが1つのエキスパートに集中しないよう、補助損失を併用して学習します。
Mixtral 8x7Bは、1トークンあたり8つのエキスパートのうち2つ మాత్రమేを使用するため、総パラメータ数は46.7Bでありながら、1トークンあたりに सक्रिय化されるのは12.9Bにとどまり、Llama 2 70Bと比べて推論が約6倍高速でした。
大規模言語モデルの時代において、MoEはモデル容量を計算コストから切り離すための中核技術であり、GPT、Mixtral、Switch Transformerのような先端モデルの設計を支えています。
Mixture of Expertsとは
Mixture of Experts(MoE)は、モデルを1つの巨大なニューラルネットワークとして扱うのではなく、複数のexpertサブネットワークと、それらの中から選択するrouter(ゲーティングネットワーク)で構成するアーキテクチャです。入力が来ると、ルーターが各トークンをどのエキスパートに送るかを決定し、選ばれたエキスパートだけが計算に参加します。すべてのパラメータを毎回動かす密なモデルとは異なり、MoEはトークンごとに少数のエキスパートだけを有効化し、このsparse activationがその定義的な特徴です。
この設計の利点は、モデル容量を計算コストから切り離せることです。エキスパートを追加するほど総パラメータ数は増えますが、1つのトークンを処理する際の実際の計算は、選ばれた少数のエキスパートに限定されます。TransformerベースのLLMでは、通常、各層のFeed-Forward(FFN)ブロックを複数のエキスパートFFNに置き換える形でMoEが適用されます。
仕組み:ルーティングとスパース活性化
この仕組みの中心にあるのがルーターです。入力トークンベクトルxに対して、ルーターはゲーティング重み行列Wgを掛け、softmaxを適用して各エキスパートのスコアを生成し、その後、スコア上位k個のエキスパートのみを選択します。これをtop-k routingと呼びます。Switch Transformerはk=1(単一エキスパート)、Mixtralはk=2(2つのエキスパート)を採用しています。選ばれなかったエキスパートはそのトークンに対してフォワードパスを実行しないため、ここで計算量が削減されます。
# Conceptual routing (top-k)
gates = softmax(x @ W_g) # score for each expert
top = topk(gates, k) # select top-k experts (e.g., k=2)
y = sum( gates[i] * expert_i(x) for i in top ) # only selected experts computeスパース・ルーティングには、1つ課題があります。学習が進むにつれて、トークンが一部の人気エキスパートに流れ込み、他のエキスパートがほとんど使われなくなり、偏りが生じることがあります。これを防ぐため、MoEにはload-balancing auxiliary loss専門家が均等に使われるよう促します。Hugging Face の MoE 解説では、「すべての expert に同等の重要性を与える」ために補助損失を追加することが説明されており、ルーター z-loss のような手法にも触れています。これは、ルーターに入力される大きなロジットを抑制して学習の安定性を高めるものです。
Dense モデルとの比較
項目 | Dense モデル | MoE(Sparse)モデル |
|---|---|---|
活性化の範囲 | すべてのトークンで全パラメータを使用 | トークンごとに選択された expert のみを使用 |
計算量とパラメータ数 | 計算量はパラメータ数に比例して増加する | 総パラメータ数は大きいが、トークンごとの計算量は概ね一定に保たれる |
事前学習速度 | ベースライン | 同等の品質により速く到達する(Switch Transformer では約 7 倍) |
メモリ(VRAM) | アクティブなパラメータのみを読み込む | 実際には使われない expert も含め、すべての expert をメモリに保持する必要がある |
ファインチューニング | 比較的安定している | 過学習しやすく、別個のハイパーパラメータが必要 |
重要性と実証
MoE の概念自体は 1991 年の Jacobs らによる「Adaptive Mixture of Local Experts」にさかのぼりますが、現代の大規模ニューラルネットワークへの本格的な適用における転機となったのは、2017 年の Shazeer らによる「Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer」(arXiv:1701.06538)でした。この論文では、数千の expert から構成される sparse-gated MoE 層を提案し、計算効率を大きく損なうことなくモデル容量を 1000 倍以上に拡張できることを示しました。
Google の Switch Transformer(Fedus ら、arXiv:2101.03961)は、その後 top-k ルーティングを最小限の k=1 に簡素化し、ルーティング計算と通信コストを削減しながら品質を維持しました。さらに、最大 1.571 兆パラメータの Switch-C モデルを学習しました。同じ計算予算で、Dense な T5-XXL と比べて事前学習速度が約 7 倍向上したと報告されています。
実世界で公開されている LLM において MoE の影響を明確に示した事例としては、Mistral AI のMixtral 8x7B(arXiv:2401.04088)。各層は8つのエキスパートで構成され、ルーターは各トークンごとに2つを選択するため、総パラメータ数は46.7Bですが、トークンごとにアクティブになるのは12.9Bです。Mistralの公式発表によると、Mixtralは32kトークンのコンテキストに対応し、「6x faster inference」により大半のベンチマークでLlama 2 70Bを上回り、GPT-3.5と同等またはそれ以上の性能を示します。実質的には、12.9Bの密なモデルと同じ速度とコストで、はるかに大規模なモデルの品質を実現しています。
ただし、トレードオフは明確です。推論時には、ほとんど使われない場合でもすべてのエキスパートをメモリに読み込む必要があるため、VRAM負荷が大きくなります。また、ファインチューニング時には、MoEは密なモデルよりも過学習しやすい傾向があります。Hugging FaceのMoE解説では、「事前学習時の困惑度が同じでも、疎なモデルは推論寄りの下流タスクで密なモデルに後れを取ることがある」と指摘しつつ、MoEは密なモデルよりもインストラクションチューニングの恩恵を受けやすいとも述べています。