Mixture-of-Expertsを用いた効率的な大規模言語モデルの台頭

自然言語処理（NLP）の世界では、大規模でより優れた言語モデルを構築する追求が、多くの最近の進歩の原動力となっています。ただし、これらのモデルはサイズが大きくなるにつれて、トレーニングと推論のための計算要件が増大し、利用可能なハードウェアリソースの限界に達してしまいます。

ここで、Mixture-of-Experts（MoE）という手法が登場します。これは、計算負荷を軽減しながら、大規模でより強力な言語モデルのトレーニングを可能にすることを約束しています。以下では、MoEについて、その起源、内部メカニズム、トランスフォーマー 기반の言語モデルにおけるその応用について説明します。

Mixture-of-Expertsの起源

Mixture-of-Experts（MoE）の概念は、1990年代初頭に研究者が条件付き計算のアイデアを探求したときに遡ることができます。ここでは、ニューラルネットワークの部分が入力データに基づいて選択的に活性化されます。この分野の先駆的な作品の1つは、1991年のJacobs et al.の「Adaptive Mixture of Local Experts」という論文でした。この論文では、各々が入力空間の異なる領域に特化したニューラルネットワークのアンサンブルのための教師あり学習フレームワークが提案されました。

MoEの基本的な考え方は、入力データのサブセットを処理する責任を持つ複数の「エキスパート」ネットワークを持つことです。ゲーティングメカニズム（通常は別のニューラルネットワーク自身）が、どのエキスパート（またはエキスパート）が特定の入力を処理するかを決定します。このアプローチにより、モデルは各入力に対して全モデルの容量を使用するのではなく、関連するエキスパートのみを活性化することで、計算リソースをより効率的に割り当てることができます。

これ以降、研究者は条件付き計算のアイデアを探求し、階層的なMoE、条件付き計算の低ランク近似、確率ニューロンとハード閾値活性化関数を介した勾配の推定などの技術を開発しました。

トランスフォーマーにおけるMixture-of-Experts

Mixture of Experts

MoEのアイデアは数十年前からありますが、トランスフォーマー基盤の言語モデルへのその応用は比較的最近のことです。トランスフォーマーは、現在の言語モデルのデファクトスタンダードとなりましたが、複数の層で構成されており、各層には自己注意メカニズムとフィードフォワードニューラルネットワーク（FFN）があります。

トランスフォーマーにMoEを適用するための重要な革新は、密なFFN層を、複数のエキスパートFFNとゲーティングメカニズムで構成される疎なMoE層に置き換えることです。ゲーティングメカニズムは、各入力トークンにどのエキスパート（またはエキスパート）を処理するかを決定し、モデルが入力シーケンスごとにエキスパートのサブセットのみを選択的に活性化できるようにします。

MoEをトランスフォーマーに適用する初期の作品の1つは、2017年のShazeer et al.による「Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer」という論文でした。この作品では、疎なゲーティングメカニズムを使用するMoE層の概念が導入され、エキスパートの選択プロセスに疎性とノイズが追加され、各入力に対してエキスパートのサブセットのみが活性化されることが保証されました。

その後、他の多くの作品がトランスフォーマーへのMoEの応用をさらに進歩させ、トレーニングの不安定性、負荷バランス、効率的な推論などの課題に対処しました。注目すべき例としては、Switch Transformer（Fedus et al., 2021）、ST-MoE（Zoph et al., 2022）、GLaM（Du et al., 2022）があります。

言語モデルのMixture-of-Expertsの利点

言語モデルのMoEを使用する主な利点は、モデルのサイズを拡大しながら、推論時の計算コストを相対的に一定に維持できることです。各入力トークンに対してエキスパートのサブセットのみを選択的に活性化することで、MoEモデルは、密なモデルと同じ表現力を達成しながら、はるかに少ない計算量で済むことができます。

例えば、7億パラメータの密なFFN層を持つ言語モデルを考えてみましょう。もしも、この層を、各々が7億パラメータを持つ8つのエキスパートで構成されるMoE層に置き換えるなら、パラメータの総数は56億に増加します。ただし、推論時に各トークンに対して2つのエキスパートのみを活性化する場合、計算コストは14億パラメータの密なモデルのものと同等になります。なぜなら、2つの7億パラメータの行列乗算が計算されるからです。

この推論時の計算効率は、リソースが限られている展開シナリオ（例：モバイルデバイスまたはエッジコンピューティング環境）で特に貴重です。さらに、トレーニング時の計算要件の削減は、エネルギー節約と炭素排出量の削減につながり、持続可能なAI慣行への注目と一致しています。

課題と考慮事項

MoEモデルは魅力的な利点を提供しますが、その採用と展開もいくつかの課題と考慮事項を伴います:

1. トレーニングの不安定性: MoEモデルは、密なモデルと比較してトレーニングの不安定性に陥りやすいことが知られています。この問題は、エキスパートの活性化の疎性と条件性に起因し、勾配の伝播と収束の課題につながる可能性があります。Zoph et al.（2022）が提案したルーターのz損失のようなテクニックは、これらの不安定性を緩和するために提案されていますが、さらに研究が必要です。
2. ファインチューニングと過剰適合: MoEモデルは、特にダウンストリームタスクのデータセットが比較的小さい場合、ファインチューニング中に過剰適合しやすい傾向があります。この挙動は、MoEモデルの容量と疎性の増加に起因し、トレーニングデータに過度に特化する可能性があります。過剰適合を緩和するために、慎重な正則化とファインチューニング戦略が必要です。
3. メモリ要件: MoEモデルは、推論時の計算コストを削減できますが、密なモデルと比較してメモリ要件が高いことがあります。これは、すべてのエキスパートの重みをメモリに読み込む必要があるため、各入力に対してサブセットのみが活性化されていても、です。メモリ制約は、リソースが限られているデバイスでのMoEモデルのスケーラビリティを制限する可能性があります。
4. 負荷バランス: 計算効率を最適化するには、エキスパート間で負荷をバランスさせることが重要です。つまり、エキスパートの1つが過負荷にならないようにし、他のエキスパートが未使用のままにならないようにします。トレーニング中に補助的な損失を使用し、エキスパートの容量係数を慎重に調整することで、この負荷バランスが達成されます。容量係数は、各エキスパートに割り当て可能なトークンの最大数を決定します。
5. 通信オーバーヘッド: 分散トレーニングと推論シナリオでは、MoEモデルは、エキスパート間で活性化と勾配情報を交換する必要があるため、追加の通信オーバーヘッドを導入する可能性があります。効率的な通信戦略とハードウェアに依存したモデル設計は、このオーバーヘッドを緩和するために不可欠です。

これらの課題があるにもかかわらず、MoEモデルは大規模でより優れた言語モデルを可能にするため、研究者はこれらの問題を解決し、軽減するための重要な努力を払っています。

例：Mixtral 8x7BとGLaM

MoEを言語モデルに実際に適用する例を示すために、2つの注目すべき例、Mixtral 8x7BとGLaMを考えてみましょう。

Mixtral 8x7Bは、Mistral言語モデルのMoEバリアントで、Anthropicによって開発されました。8つのエキスパートで構成されており、各エキスパートは7億パラメータを持つため、合計56億パラメータになります。ただし、推論時に各トークンに対して2つのエキスパートのみが活性化されるため、計算コストは14億パラメータの密なモデルのものと同等になります。7億パラメータの行列乗算を2回計算するためです。

Mixtral 8x7Bは、70億パラメータのLlamaモデルを上回るパフォーマンスを示し、推論時間も大幅に短縮しました。さらに、自然言語の指示に従う能力を高めるために、Mixtral-8x7B-Instruct-v0.1という、指示に特化したバージョンもリリースされています。

もう1つの注目すべき例は、Googleが開発した大規模なMoEモデル、GLaM（Google Language Model）です。GLaMは、デコーダー専用のトランスフォーマーアーキテクチャを採用し、1.6兆トークンの巨大なデータセットでトレーニングされました。モデルは、ファインショットとワンショット評価で優れたパフォーマンスを示し、GPT-3と同等の品質を達成しながら、GPT-3のトレーニングに必要なエネルギーの3分の1のみで済みました。

GLaMの成功は、その効率的なMoEアーキテクチャに帰因できます。これにより、膨大な数のパラメータを持つモデルをトレーニングしながら、計算要件を抑えることができました。モデルはまた、MoEモデルが密なモデルと比較してエネルギー効率と環境に優しい可能性を実証しました。

Grok-1アーキテクチャ

GROK MIXTURE OF EXPERT

Grok-1は、効率とパフォーマンスを最大化するように設計された、トランスフォーマー基盤のMoEモデルです。以下は、その主な仕様です:

1. パラメータ: 314億パラメータを備えるGrok-1は、現在最も大規模なオープンLLMです。ただし、MoEアーキテクチャのおかげで、任意の時点で有効な重みは25%（約86億パラメータ）のみとなり、処理能力が向上します。
2. アーキテクチャ: Grok-1は、8つのエキスパートで構成されるMixture-of-8-Expertsアーキテクチャを採用しており、各トークンは推論時に2つのエキスパートによって処理されます。
3. 層: モデルは64のトランスフォーマー層で構成されており、各層にはマルチヘッド自己注意と密なブロックが含まれます。
4. トークン化: Grok-1は、131,072トークンの語彙サイズを備えるSentencePieceトークナイザーを使用します。
5. 埋め込みと位置エンコーディング: モデルには6,144次元の埋め込みが特徴であり、回転位置エンコーディングを採用しています。これにより、従来の固定位置エンコーディングと比較して、データのよりダイナミックな解釈が可能になります。
6. 注意: Grok-1は、クエリに対して48の注意ヘッド、キーと値に対して8の注意ヘッドを使用し、各ヘッドのサイズは128です。
7. コンテキスト長: モデルは、bfloat16精度を使用して、最大8,192トークンのシーケンスを処理できます。

パフォーマンスと実装の詳細

Grok-1は、MMLUスコアで73%を達成し、LLaMa 2 70BとMixtral 8x7Bを上回るパフォーマンスを示しました。これは、さまざまなテストでその効率と正確性を実証しています。

ただし、Grok-1はその巨大さ故に、相当なGPUリソースを必要とすることに注意する必要があります。現在のオープンソースリリースでは、カスタムカーネルを避けるために、MoE層の非効率的な実装が使用されています。

それでも、モデルはアクティベーション分割と8ビット量子化をサポートしており、パフォーマンスを最適化し、メモリ要件を削減することができます。

xAIは、Grok-1をApache 2.0ライセンスの下でリリースしました。これにより、重みとアーキテクチャが世界中のコミュニティで利用および貢献できるようになりました。

オープンソースリリースには、Grok-1モデルをロードして実行する方法を示すJAXの例コードリポジトリが含まれています。ユーザーは、トレントクライアントまたはHuggingFace Hubを直接使用して、チェックポイント重みをダウンロードできます。これにより、この画期的なモデルへのアクセスが容易になります。

言語モデルのMixture-of-Expertsの将来

大規模でより優れた言語モデルの需要が拡大するにつれて、MoE技術の採用はさらに勢いを増すことが予想されます。現在の研究努力は、トレーニングの不安定性を改善し、ファインチューニング時の過剰適合を緩和し、メモリと通信要件を最適化することに焦点を当てています。

1つの有望な方向性は、階層的なMoEアーキテクチャの探求です。ここでは、各エキスパート自身が複数のサブエキスパートで構成されます。このアプローチにより、より大規模なモデルを可能にするだけでなく、計算効率を維持したまま、大規模モデルと同等の表現力を達成できる可能性があります。

さらに、MoEモデルを効率的に処理するハードウェアとソフトウェアシステムの開発は、活発な研究分野です。MoEモデルの疎な計算パターンを効率的に処理するように特別に設計されたアクセラレータと分散トレーニングフレームワークは、パフォーマンスとスケーラビリティをさらに向上させる可能性があります。

また、MoE技術を言語モデリングの他の進歩、例えば疎な注意メカニズム、効率的なトークン化戦略、またはマルチモーダル表現と統合することで、さらに強力で汎用的な言語モデルを開発できる可能性があります。これらのモデルは、幅広いタスクに取り組む能力を持ちます。

結論

Mixture-of-Experts手法は、大規模でより優れた言語モデルを実現するための強力なツールとして浮上しています。入力データに基づいてエキスパートを選択的に活性化することで、MoEモデルは、密なモデルと比較して計算効率を向上させることを約束しています。トレーニングの不安定性、過剰適合、メモリ要件などの課題が残っていますが、MoEモデルは計算効率、スケーラビリティ、環境への配慮という点で大きな利点を提供します。

自然言語処理の分野が何が可能かを推進し続けるにつれて、MoE技術の採用は次世代の言語モデルの開発に重要な役割を果たす可能性があります。MoEをモデルアーキテクチャの他の進歩、トレーニング技術、ハードウェア最適化と組み合わせることで、さらに強力で汎用的な言語モデルを実現できます。これらのモデルは、人間と自然で無駄のない方法で理解し、コミュニケーションをとることができます。