Mixture of Experts (MoE) の解説

Mixture of Experts (MoE) は、トランスフォーマーモデルが密なモデルに比べて事前学習を大幅に高速化し、総パラメータ数に対して推論速度も速くなるようにします。これは、密なフィードフォワードネットワーク (FFN) 層を、入力トークンごとに少数の「エキスパート」だけを活性化するスパース MoE 層に置き換えることで実現されます。

Mixture of Experts のコアアーキテクチャ

MoE モデルは標準的な FFN 層を、スパース MoE 層とルーティング機構の組み合わせに置き換えます。このアーキテクチャは主に 2 つのコンポーネントから構成されます：

  • Sparse MoE Layers: 単一の密な層の代わりに、これらの層は複数の「エキスパート」を含み、各エキスパートは通常ニューラルネットワーク（多くは FFN）です。エキスパート自体が MoE になるような階層構造で配置することも可能です。
  • Gate Network (Router): 学習されたルーターが、どのトークンをどのエキスパートに送るかを決定します。ルーターはネットワーク全体と同時に事前学習され、トークンを 1 つまたは複数のエキスパートに送るように設定できます。

MoE は計算効率を提供しますが、特有のトレードオフも伴います。すべてのエキスパートをメモリにロードする必要があるため、VRAM が大量に必要です（単一のフォワードパスではごく少数のエキスパートしか活性化されません）。例えば、Mixtral 8x7B は 47B パラメータのモデルを保持できるだけの VRAM が必要ですが、トークンあたり 2 つのエキスパートを使用した場合の推論計算量（FLOPs）は 12B モデルに近くなります。

スパース性とルーティングメカニズム

スパース性は、条件付き計算（例ごとにシステムの一部だけを活性化）を用いることで、計算量を比例的に増やさずにモデルのサイズをスケールさせることを可能にします。

ゲーティング関数

ルーティングは、どのエキスパート $E$ を活性化するかを決定するゲーティングネットワーク $G$ によって処理されます。従来の設定ではソフトマックス関数が使用されます。Noisy Top-k Gating のような高度な手法では、調整可能なノイズを導入し、上位 $k$ の値を選択して、学習中の探索と負荷分散を改善します。

負荷分散とエキスパート容量

少数のエキスパートが過度に利用され、他が無視される「ウィナー・テイク・オール」シナリオを防ぐために、MoE は以下を採用します：

  • Auxiliary Loss: ルーターがトークンをすべてのエキスパートに均等に分配するよう促す追加の損失関数。
  • Expert Capacity: 単一のエキスパートが処理できるトークン数の上限。エキスパートが容量に達した場合、トークンはドロップされるか、残差接続を介して次の層に送られます。

MoE 実装の進化

GShard と Switch Transformers

Google の GShard は、トップ2 ゲーティングとランダムルーティング（2 番目のエキスパートが重み比例で選択される）を使用して、トランスフォーマーを 6000 億パラメータ以上にスケールさせました。

Switch Transformers は、単一エキスパート戦略を導入することでさらにシンプル化しました。各トークンを 1 つのエキスパートにだけルーティングすることで、ルーター計算を削減し、エキスパートのバッチサイズを半減し、通信コストを低減しつつ品質を維持しました。このアーキテクチャは T5-XXL に対して 4 倍の事前学習速度向上を達成しました。

安定性と専門化

学習の安定性は、Router z-loss（ST-MoE で導入）などの手法で管理されます。これは、ゲーティングネットワークに入る大きなロジットにペナルティを課し、丸め誤差を減らすものです。

エキスパートの専門化に関する研究では、エンコーダーのエキスパートはしばしば浅い概念（例：句読点や固有名詞）に特化し、デコーダーのエキスパートはあまり専門化しないことが示されています。多言語設定では、負荷分散要件のためエキスパートは単一言語に特化することはほとんどありません。

ファインチューニングとインストラクションチューニング

スパースモデルは密なモデルに比べて過学習しやすいため、効果的なファインチューニング戦略には以下が含まれます：

  • Higher Regularization: スパース層内でより高いドロップアウト率を使用する。
  • Selective Freezing: MoE 層を凍結し、非 MoE 層を更新することで、品質を保ちつつ学習を高速化できる。
  • Hyperparameter Adjustment: スパースモデルはしばしば小さなバッチサイズと高い学習率から恩恵を受けます。

最近の研究では、MoE は密なモデルに比べて インストラクションチューニング から大幅に恩恵を受けることが示されています。MoE と同等の T5 を比較した場合、インストラクションチューニングによる MoE の性能向上（Flan-MoE vs MoE）は、密なモデルの向上（Flan T5 vs T5）よりも大きかったです。

計算効率と並列化

並列化戦略

MoE の膨大なパラメータ数に対処するため、さまざまな並列化手法が使用されます：

  • Expert Parallelism: エキスパートを異なるワーカーに分散させます。MoE 層では、トークンは特定のエキスパートが存在するワーカーにルーティングされ、非 MoE 層ではデータ並列と同様に動作します。

ハードウェア最適化

  • Capacity Factor: 品質と通信コストのトレードオフ。容量係数が高いほど品質は向上しますが、アクティベーションのメモリとデバイス間通信が増加します。
  • Block-Sparse Operations: MegaBlocks のようなプロジェクトは、従来のバッチ行列乗算をブロックスパース演算に置き換え、トークンの割り当てが不均衡でもトークンをドロップせずに処理し、顕著な速度向上を実現します。

比較: スパース MoE と 密モデル

Feature Sparse MoE Dense Model
Pretraining Compute より効率的 / 高速 非効率的
Inference Speed より速い（総パラメータ数に対して） 遅い（総パラメータ数に対して）
VRAM Requirement 高い（すべてのエキスパートをロード必要） 低い（総パラメータ数に対して）
Best Use Case 高スループット、マシン多数 低スループット、VRAM 制限

オープンソース MoE エコシステム

いくつかのフレームワークとモデルが MoE をオープンコミュニティにもたらしました：

  • Frameworks: MegaBlocks、Fairseq、OpenMoE。
  • Models: Google の Switch Transformers（最大 1.6T パラメータ）、Meta の NLLB MoE（54B）、Mistral の Mixtral 8x7B（Llama 2 70B を上回り、推論が高速）。

Sources