Mixture of Expert（MoE）とは？構造、トレーニング方法、4つのバリエーションなどの説明

Mixture of Expert（MoE）とは、複数のエキスパートモデルを統合し、全体のパフォーマンスを向上させるためのアプローチです。この手法では、各エキスパートモデルがデータセットをもとにトレーニングされ、それぞれのエキスパートが特定のデータに特化しています。そして、これらのエキスパートを組み合わせることで、モデルの出力を生成します。通常は、重みを調整して単一の最終出力を作成します。この技術は1991年から存在していましたが、Mistral が Mistral-8x7B モデルを発表したことで、さらに注目を集めるようになりました。本記事では、Mixture of Expert (MoE) とは、構造、トレーニング方法、バリエーション、メリット・デメリット、そしてMoEを使用するモデルについて詳しく説明します。

1. Mixture of Expert（MoE）の構造

MoE アーキテクチャの基盤は、エキスパートとゲーティングメカニズムにあります。

エキスパート: 各エキスパートは、特定のタスクやデータセットにおいて適切に機能するよう、特定の種類のデータでトレーニングされた個別のニューラルネットワークです。この専門化により、エキスパートは複雑なデータを効率的に処理できるようになります。
ゲーティングメカニズム: このメカニズムは、各入力データに最も適したエキスパートを選択するために機能します。ゲーティングメカニズムはルーターの役割を果たし、最終的な出力に貢献する関連性の高いエキスパートのみを選択することで、MoE の重要な要素となっています。

2. MoEのトレーニング方法

MoE をトレーニングするには、エキスパートの重みとゲーティングメカニズムの両方を最適化する必要があります。目標は、特定の入力に対して最適な結果を得るためにエキスパートをトレーニングし、ルーターが常に最適なエキスパートに入力を配信することです。

一般的に使用されてきた従来のアプローチは、期待値を最大化するためのEMアルゴリズムです。まず、ルーターが現在の予測に基づいてエキスパートに確率を割り当てます。次に、エキスパートとルーターの重みを更新し、モデル全体のパフォーマンスを最大化します。
現在では、強化学習法がEM の代替手法として用いられることが増えており、計算の複雑さを軽減しつつトレーニングプロセスを効率化し、優れた成果を上げています（代表的な例としてDeepseekがあります）。

3. MoEのバリエーション

MoEのバリエーションは、トレーニング方法やアーキテクチャの変更に基づいて分類されます。

スパース MoE: 特定の入力に対してエキスパートの一部のみを使用します。例えば、Mistral 8x7B では、トークンごとに8人のエキスパートのうち2人が活用されます。

ハード MoE: 1つの入力に対し、予測を行う特定のエキスパートを1人だけ選択します。

ソフト MoE: 各入力に確率を割り当て、その確率を基にエキスパートを活用します。この確率は最終的な出力の予測に使用されます。例えば、4人のエキスパートと1つの入力がある場合、ルーターは各エキスパートに [0.1, 0.2, 0.3, 0.4] の確率を割り当て、最終出力に影響を与える割合を決定します。

階層型 MoE: エキスパートを複数のレベルに分けて入力を処理します。例えば、第1レベルではエキスパートが一般的なタスクに対応し、第2レベルではコーディングに特化したエキスパートが活躍します。その中には、さらに特定のプログラミング言語に特化したエキスパートも存在します。

4. MoEのメリットとデメリット

メリット:

計算負荷の軽減: Dense モデルとは異なり、MoE モデルでは推論時の計算量が削減されるため、GPU の計算速度が制約を受けることがありません。ただし、すべてのエキスパートを GPU にロードする必要があるため、RAM の要件は変わりません。
柔軟性と拡張性: MoE は多種多様なエキスパートモデルやゲーティングメカニズムを統合できるため、大規模なデータセットや複雑なタスクの処理に対応できる柔軟性とスケーラビリティを備えています。

デメリット:

トレーニングの複雑さ: MoE モデルをトレーニングするには、エキスパートとゲーティングメカニズムの両方を同時に最適化する必要があるため、プロセスが非常に複雑で手間がかかります。
メモリ使用量の多さ: MoE モデルは推論時に一部のエキスパートパラメータのみを活用しますが、それでもすべてのエキスパートをデバイスに読み込む必要があるため、メモリ消費が増大します。
負荷分散の課題: エキスパートが複数のデバイスに分散される場合、計算の効率が低下することがあります。

5. MoEを使用するモデル

5.1 Mixtral-8x7B

Mixtral-8x7B は通常のデコードアーキテクチャを採用していますが、フィードフォワードネットワーク層をゲーティングに置き換え、それに続く8人のエキスパートで構成されています。この構造が32回繰り返されることで、Mixtral-8x7B モデルが形成されます。
Mixtral-8x7B では、トークンごとのゲーティングによって、8人のエキスパートのうち2人のエキスパートへルーティングされます。次に、これら2人のエキスパートのゲーティング確率に基づき出力が合計されます。この後、ソフトマックス層とロジット層を通過して最終出力が得られます。
トレーニング中、トークンはほぼ均等にエキスパートに割り当てられるため、分散がスムーズに行われます。

Mixtral-8x7B は合計46.7Bのパラメーターを持ちますが、推論時には13億個のパラメーターのみが使用されます。その結果、同様なパラメーターサイズを持つ他のDenseモデルと比較して推論が高速化されます。

5.2 DeepSeek

DeepSeek MoE では、動作を改良するために以下のような変更が加えられています：

エキスパートの分割: エキスパートをmNのグループに分割し、それぞれのグループからmKをアクティブ化します。具体的には、フィードフォワードネットワークの隠れ層のサイズを元のサイズの$\frac{1}{m}$に縮小し、各エキスパートは小さなm人の専門家に分割され、合計mN人のエキスパートを生成します。この設計によりアクティブ化されるエキスパートの数はm倍に増加し、効率的な処理が可能となります。

共有エキスパート: 従来の方法では、各エキスパートが同じ共通知識を学習していたため、パラメーターの冗長性が生じていました。共有エキスパートは、異なるコンテキストから共通知識を取り込み、統合する役割を担います。これにより、共通知識における冗長性を削減し、他のエキスパートが個別の知識学習に専念できるようになります。

また、DeepSeek はエキスパートのルーティングを最適化し、計算ボトルネックや負荷の不均衡を解消します。これにより、特定のエキスパートが完全にトレーニングされないといった問題を回避し、モデル全体のパフォーマンスを向上させます。