人工知能

大規模言語モデルにおけるMixture-of-Expertsの台頭

mm
公開日
更新日
Mixture of Experts Grok Mistral

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

ここで、Mixture-of-Experts（MoE）という手法が登場し、計算負荷を軽減しながら、大規模で高性能な言語モデルをトレーニングできる可能性を提供している。以下では、MoEの概念、内部メカニズム、トランスフォーマーベースの言語モデルへの応用について説明する。

Mixture-of-Expertsの起源

Mixture-of-Experts（MoE）の概念は、1990年代初頭に、入力データに基づいてニューラルネットワークの部分を選択的に活性化するというアイデアを探求した研究者们によって生み出された。 この分野における先駆的な作品の1つは、Jacobs et al.による1991年の「Adaptive Mixture of Local Experts」という論文であり、入力空間の異なる領域に特化したニューラルネットワークのアンサンブルを対象とした監督学習フレームワークを提案した。

MoEの基本的な考え方は、入力データのサブセットを処理する専門家ネットワークを複数持つことである。ゲーティングメカニズム（通常はニューラルネットワーク自体）が、どの専門家を入力データに適用するかを決定する。このアプローチにより、モデルは各入力データに対して関連する専門家のみを活性化することで、計算リソースをより効率的に割り当てることができる。

これ以降、研究者们は、階層的なMoE、条件付き計算の低ランク近似、確率的ニューロンとハード閾値活性化関数を通じて勾配を推定する技術などの開発に取り組んできた。

トランスフォーマーベースの言語モデルにおけるMixture-of-Experts

Mixture of Experts

Mixture of Experts

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

トランスフォーマーにMoEを適用する際の重要な革新は、密なFFN層を疎なMoE層に置き換えることである。各MoE層には複数の専門家FFNとゲーティングメカニズムが含まれており、入力トークンごとにどの専門家を活性化するかを決定する。

トランスフォーマーにMoEを適用した初期の研究の1つは、Shazeer et al.による2017年の「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層を持つ言語モデルを考えてみよう。MoE層に置き換える場合、8つの専門家で構成され、各専門家に7億パラメータがあると仮定すると、合計56億パラメータになる。ただし、推論時に各トークンに対して2つの専門家のみを活性化する場合、計算コストは14億パラメータの密なモデルのコストと同等になる。

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

課題と考慮事項

MoEモデルは魅力的な利点を提供するが、その採用と展開は複数の課題と考慮事項を伴う。

  1. トレーニングの不安定性: MoEモデルは、密なモデルと比較してトレーニングの不安定性に陥りやすい。この問題は、専門家の活性化が疎で条件付きであるため、勾配伝播と収束に課題が生じる。Zoph et al.（2022）によって提案されたルーターのz損失（z-loss）などのテクニックが、不安定性を軽減するために提案されているが、さらに研究が必要である。
  2. ファインチューニングと過学習: MoEモデルは、ファインチューニング中に過学習しやすい。この現象は、MoEモデルの容量と疎性の増加によるものであり、トレーニングデータに過度に特化する可能性がある。過学習を軽減するために、慎重な正則化とファインチューニング戦略が必要である。
  3. メモリ要件: MoEモデルは、密なモデルと比較して、推論時にメモリ要件が高くなることがある。これは、すべての専門家の重みをメモリにロードする必要があるためである。メモリ制約は、リソースが限られているデバイスでのMoEモデルのスケーラビリティを制限する可能性がある。
  4. ロードバランシング: 計算効率を最適化するために、専門家間のロードバランシングが重要である。各専門家が過負荷にならないように、また他の専門家が未使用にならないようにする必要がある。ロードバランシングは、トレーニング中に補助損失を使用し、専門家の容量係数を慎重に調整することで実現される。
  5. 通信オーバーヘッド: 分散トレーニングと推論シナリオでは、MoEモデルは、専門家間の活性化と勾配情報の交換に伴う追加の通信オーバーヘッドを導入する可能性がある。効率的な通信戦略とハードウェアに依存したモデル設計が、このオーバーヘッドを軽減するために不可欠である。

これらの課題にもかかわらず、MoEモデルは大規模で高性能な言語モデルを可能にするという潜在的な利点により、研究者们の間で大きな関心を集めている。

例：Mixtral 8x7BとGLaM

MoEを言語モデルに適用する実際の例として、Mixtral 8x7BとGLaMの2つを考えてみよう。

Mixtral 8x7Bは、アンソロジックによって開発されたMistral言語モデルのMoEバリアントである。8つの専門家で構成されており、各専門家には7億パラメータがある。合計で56億パラメータになるが、推論時に各トークンに対して2つの専門家のみを活性化するため、計算コストは14億パラメータの密なモデルのコストと同等になる。

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

もう1つの注目すべき例は、Googleによって開発された大規模なMoEモデルであるGLaMである。GLaMは、デコーダーのみのトランスフォーマーアーキテクチャを採用しており、1.6兆トークンの巨大なデータセットでトレーニングされている。GLaMは、ファインショット評価とワンショット評価で高性能を発揮しており、GPT-3と同等の品質を達成しながら、GPT-3のトレーニングに必要なエネルギーの3分の1しか使用していない。

GLaMの成功は、効率的なMoEアーキテクチャに归することができる。GLaMは、多数のパラメータを持つモデルをトレーニングすることを可能にし、計算コストを抑えることで、エネルギー効率と環境への配慮に優れたモデルとなっている。

Grok-1アーキテクチャ

GROK MIXTURE OF EXPERT

GROK MIXTURE OF EXPERT

Grok-1は、効率性とパフォーマンスを最大化するために設計された、トランスフォーマーベースのMoEモデルである。以下に、Grok-1の主な仕様を示す。

  1. パラメータ数: 314億パラメータを持ち、現在最大のオープンLLMである。ただし、MoEアーキテクチャにより、各時点で有効な重みは約86億パラメータ（全体の25%）のみである。
  2. アーキテクチャ: Grok-1は、8つの専門家で構成されるMixture-of-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をモデルアーキテクチャ、トレーニング技術、ハードウェア最適化の他の進歩と組み合わせることで、さらに強力で多様な言語モデルが可能になり、人間との自然でシームレスなコミュニケーションが可能になる。

mm

私は過去5年間、機械学習とディープラーニングの魅力的世界に没頭してきました。私の情熱と専門知識は、AI/MLに特に焦点を当てた50以上の多様なソフトウェアエンジニアリングプロジェクトに貢献することになりました。私の継続的な好奇心は、自然言語処理という分野にも私を引き付け、さらに探求したいと思っています。