Mixture of Experts (MoE)とは？ 現代のLLMを支えるアーキテクチャ

Mixture of experts MoE AIは、複数の専門化されたサブネットワーク（エキスパートと呼ばれる）に計算を分割するニューラルネットワークアーキテクチャです。軽量のゲーティングネットワークが、各入力トークンを処理するエキスパートを選択します。任意の時点でアクティブになるのはごく一部のみです。この選択的ルーティングにより、MoEはすべてのトークンにすべてのパラメータを使用する従来の密なモデルと区別されます。

この設計が注目を集めているのは、総パラメータ数が数百億に達しても、推論コストをはるかに小さなモデルと同等に抑えられるためです。Mixtral 8x7BやGPT-4のいくつかのバージョンがこのパターンを採用しています。仕組みを理解することで、モデル規模と効率に関する主張を評価する助けになります。

Key Takeaways

• Mixture of experts MoE AIは、単一の大規模フィードフォワードネットワークを、多くの小さなエキスパートネットワークと、適切なサブセットを選択するルーターに置き換えます。

• スパース活性化により、各トークンにつき固定数のエキスパートのみが処理するため、総パラメータ数が増えてもトークンあたりの計算量はほぼ一定に保たれます。

• Mixtral 8x7Bなどのモデルは、MoE版が推論時に同程度のアクティブパラメータを使用しながら、ベンチマークで密なベースラインと同等かそれ以上を達成できることを示しています。

• エキスパート間の負荷分散とルーターの安定性に追加の手法が必要となるため、訓練は密な訓練より複雑なままです。

• 同じアーキテクチャはオープンなモデルとプロプライエタリなシステムの両方に現れており、この手法が研究から本番の大規模言語モデルへ移行したことを示しています。

Mixture of Experts Definition

Mixture of expertsモデルは、エキスパートサブネットワークの集合とゲーティング関数を含みます。ゲーティング関数はエキスパートと同じ入力を受け取り、エキスパートプール上の確率分布を出力します。最高スコアのトップkエキスパートのみが入力を受け取り、出力に寄与します。その他のエキスパートはそのトークンに対して非アクティブのままです。

コアの特性には、モジュール性、スパース性、条件付き計算が含まれます。モジュール性は、各エキスパートが異なるデータパターンやタスクタイプに特化できることに現れます。スパース性は、トークンあたりのアクティブパラメータの割合が低いことに現れます。条件付き計算は、アクティブなセットがトークンごとに変化することに現れます。

これらの特性により、MoEモデルは密なトランスフォーマーと区別されます。密なモデルではすべての重みがすべてのフォワードパスに参加します。MoEモデルでは、任意のトークンに対して大部分の重みがアイドル状態になります。その結果、トークンあたりの計算量をほぼ一定に保ちつつ、より高い総容量を実現します。

How Mixture of Experts Works

このアーキテクチャは、ルーター、エキスパートプール、出力結合ステップの3つの層に分解されます。

Router network - selecting experts

ルーターは小さな線形層または浅いMLPで、入力の隠れ状態をエキスパートスコアのベクトルにマッピングします。ソフトマックスがスコアを確率に変換します。システムはこれらの確率に従ってトップkのエキスパートを選択します。エキスパート総数が8以上であっても、kは通常1または2です。

Expert networks - specialized computation

各エキスパートは通常、標準的なトランスフォーマーブロック内の位置ごとのFFNと同じ構造のフィードフォワードネットワークです。エキスパートは単一の大規模FFNより小さいため、訓練時のメモリ予算を超えずに多数を維持できます。訓練中に異なるエキスパートが異なるデータ分布からの勾配を受け取ることで、特化が生まれます。

Output combination - weighted summation

選択されたエキスパートの出力は、ルーターの確率で乗算され合計されます。これによりMoE層からの最終的な寄与が生成されます。アテンションを含むトランスフォーマーブロックの残りの部分は変更されません。

スパース活性化により、アクティブなパラメータ数を低く抑えます。7BパラメータのFFNに相当するサイズのエキスパートを8つ持つモデルでも、k=1のときアクティブになる計算量は約7Bモデル程度です。残りのエキスパートはメモリに残りますが、そのトークンの行列乗算には参加しません。

Mixture of Experts Explained Through Examples

Mixtral 8x7Bは各層に8つのエキスパートを含み、トークンごとに2つを活性化します。したがって推論時のアクティブパラメータ数は12-13Bの密なモデルに相当しますが、総パラメータ数は46Bを超えます。標準ベンチマークでは、生成時のメモリ帯域幅を抑えつつ、いくつかのカテゴリでLlama 2 70Bを上回ります。

GPT-4はより多くのエキスパートと高い総パラメータ数を持つMoE設計を採用していると広く報告されています。公開情報は限定的ですが、高い総容量と制御された推論コストを組み合わせるパターンは同じ原理に一致します。GoogleのSwitch TransformerやGLaMなどの他の本番システムも、さらに大規模で同一の構造に従っています。

これらのモデルに共通するパターンは、ルーターが異なるトークンタイプを異なるエキスパートに送ることを学習することです。コードトークンはあるエキスパート群へ、自然言語トークンは別の場所へルーティングされる場合があります。この暗黙の特化により、訓練時のサンプル効率が向上します。

Real-World Applications

オープンウェイトのチャットモデルでは、この設計により管理可能なコストでより長いコンテキストをサポートします。ユーザーはMixtral 8x7Bをローカルまたは控えめなクラウドGPU上で実行できます。アクティブメモリ使用量が密な13Bモデルに近いためです。内部アシスタントを展開する企業は、提供ハードウェアの比例した増加なしに高い出力品質を得られます。

スケーリング則を研究する研究グループは、密な訓練が非現実的になる総パラメータ数を超えた後も性能が向上し続けるかどうかをテストするためにMoEを使用します。Switch Transformerに関する論文では、補助的な負荷分散損失と組み合わせることで、1兆パラメータを超えるMoEモデルが安定して訓練できることを示しました。

Common Questions About mixture of experts MoE AI

Q: mixture of experts MoE AIは、同じ総サイズの密なモデルと比べてどのように推論コストを削減しますか？

A: 選択されたエキスパートのみが行列乗算を実行します。ルーターのコストは小さく、メモリ帯域幅は総重みではなくアクティブな重みによって決まります。その結果、生成速度はアクティブなエキスパートのサイズに追従します。

Q: mixture of experts MoE AIは特別なハードウェアを必要としますか？

A: 標準的なGPUで十分です。効率的な推論はゼロ寄与のエキスパートをスキップするカーネルから恩恵を受けますが、アーキテクチャ自体はモデル重みがロードされれば既存のトランスフォーマーランタイム上で動作します。

Q: 訓練は密なトランスフォーマーとどう異なりますか？

A: ルーターはエキスパートとともに訓練する必要があります。補助損失によりエキスパートの利用率を均衡に保ちます。これらの項がなければ、一部のエキスパートがほとんどトークンを受け取らず、有用な表現を学習できません。

Q: mixture of experts MoE AIは密なモデルと同じようにファインチューニングできますか？

A: はい。同じ教師ありファインチューニングと選好最適化のパイプラインが適用できます。各例でルーターとアクティブなエキスパートのみが勾配を受け取るため、実際にはバックワードパス時のメモリ使用量を削減できます。

Q: ルーターがすべてのトークンを同じエキスパートに送るとどうなりますか？

A: そのエキスパートが不均衡な勾配を受け取り、他のエキスパートは未訓練のままになります。単一エキスパートへの崩壊を防ぐため、事前訓練中に負荷分散損失とルータージッタが追加されます。