Mixtral of Experts

Albert Q. Jiang; Alexandre Sablayrolles; Antoine Roux; Arthur Mensch; Blanche Savary; Chris Bamford; Devendra Singh Chaplot; Diego de las Casas; Emma Bou Hanna; Florian Bressand; Gianna Lengyel; Guillaume Bour; Guillaume Lample; Lélio Renard Lavaud; Lucile Saulnier; Marie-Anne Lachaux; Pierre Stock; Sandeep Subramanian; Sophia Yang; Szymon Antoniak; Teven Le Scao; Théophile Gervet; Thibaut Lavril; Thomas Wang; Timothée Lacroix; William El Sayed

論文 · arXiv 2401.04088

Mixtral of Experts

Albert Q. Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Emma Bou Hanna, Florian Bressand, Gianna Lengyel, Guillaume Bour et al.

公開日 2024-01-08cs.LGcs.CL arXiv PDF

XYZNOT で読む（全文翻訳と AI 解説）

3行まとめ

🎯 過度な推論レイテンシや計算コストを発生させることなく、Transformerの容量と性能をスケーリングすること。
🔧 標準のフィードフォワードブロックを8つのSwiGLUエキスパートに置き換え、学習されたルーターを介してトークンごとに2つのエキスパートを選択するSparse Mixture of Experts（SMoE）デコーダーオンリーアーキテクチャ。
📈 トークンあたりわずか13Bのアクティブパラメータで合計47Bのパラメータにアクセスし、MMLU（70.6% vs 69.9%）およびGSM8K（74.4% vs 69.6%）でLlama 2 70Bを上回る性能を達成。

概要

本論文では、Sparse Mixture of Experts（SMoE）言語モデルであるMixtral 8x7Bを提案します。MixtralはMistral 7Bと同じアーキテクチャを持ちますが、各層が8つのフィードフォワードブロック（すなわちエキスパート）で構成されている点が異なります。各層において、ルーターネットワークがトークンごとに2つのエキスパートを選択して現在の状態を処理し、それらの出力を結合します。各トークンは2つのエキスパートしか参照しませんが、選択されるエキスパートはタイムステップごとに異なる場合があります。その結果、各トークンは47Bのパラメータにアクセスできますが、推論時に使用されるアクティブパラメータは13Bに留まります。Mixtralは32kトークンのコンテキストサイズで学習され、評価されたすべてのベンチマークにおいてLlama 2 70BおよびGPT-3.5と同等以上の性能を達成しました。特に、数学、コード生成、および多言語ベンチマークにおいてMixtralはLlama 2 70Bを大幅に上回っています。さらに、指示追従向けにファインチューニングされたモデルであるMixtral 8x7B - Instructも提供しており、人間による評価ベンチマークでGPT-3.5 Turbo、Claude-2.1、Gemini Pro、Llama 2 70B - chatを凌駕しています。ベースモデルとInstructモデルは両方ともApache 2.0ライセンスの下で公開されています。

英語の原文概要

We introduce Mixtral 8x7B, a Sparse Mixture of Experts (SMoE) language model. Mixtral has the same architecture as Mistral 7B, with the difference that each layer is composed of 8 feedforward blocks (i.e. experts). For every token, at each layer, a router network selects two experts to process the current state and combine their outputs. Even though each token only sees two experts, the selected experts can be different at each timestep. As a result, each token has access to 47B parameters, but only uses 13B active parameters during inference. Mixtral was trained with a context size of 32k tokens and it outperforms or matches Llama 2 70B and GPT-3.5 across all evaluated benchmarks. In particular, Mixtral vastly outperforms Llama 2 70B on mathematics, code generation, and multilingual benchmarks. We also provide a model fine-tuned to follow instructions, Mixtral 8x7B - Instruct, that surpasses GPT-3.5 Turbo, Claude-2.1, Gemini Pro, and Llama 2 70B - chat model on human benchmarks. Both the base and instruct models are released under the Apache 2.0 license.

論文解説

課題と動機

Dense（高密度）な自己回帰型言語モデルは、処理されるトークンごとに総パラメータ数に正比例する計算コストが発生します。モデルが70Bパラメータ以上にスケールアップすると、低レイテンシかつ高スループットで運用する（サービングする）ための計算コストが極めて高くなります。著者らは、トークンあたりの推論コストをはるかに小規模なモデルと同等に抑えつつモデル容量をスケーリングし、Llama 2 70BやGPT-3.5などの商用ベースラインに匹敵またはそれを凌駕するオープンウェイトモデルの構築を目指しました。

主要なアイデア

Sparse Mixture of Experts (SMoE): 全32層のTransformerレイヤーにおいて、Denseなフィードフォワードネットワーク（FFN）ブロックを8つの独立したエキスパートネットワークに置き換え、各トークンを上位2つのエキスパートのみにルーティングします。
パラメータ数の分離: アーキテクチャ全体としては47Bの総パラメータを保持しますが、トークンの順伝播ごとに実行されるアクティブパラメータは13Bのみであり、中規模なDenseモデルと同等の計算負荷に抑えられます。
ネイティブ32kコンテキストウィンドウ: 32,768トークンのシーケンス長をネイティブサポートし、Passkeyベンチマークで100%の検索精度を達成しています。
トピックではなく構文に基づくルーティング: トークンルーティングの実証的検証により、エキスパートの選択は高レベルなトピック領域ではなく、主に構文的機能や局所的な時間的連続性（temporal locality）によって駆動されていることが明らかになりました。

手法（主要な数式を含む）

MixtralはMistral 7Bのアーキテクチャ（ $32$ 層、隠れ層次元 $4096$ 、 $32$ アテンションヘッド、 $8$ Key-Valueヘッド）を変更し、すべてのFFNサブブロックを $n = 8$ のエキスパート（ $E_{0}, \dots, E_{n - 1}$ ）からなるMoEモジュールに置き換えています。各エキスパートは中間次元 $14336$ のSwiGLUネットワークとして実装されています。

入力トークン表現 $x$ に対し、ルーターネットワークは上位 $K$ 個のロジット（ $K = 2$ ）を用いて、エキスパートに対する疎なゲーティング分布を計算します: $G (x) := Softmax (TopK (x \cdot W_{g}))$ ここで、Top-K演算子はロジット $ℓ = x \cdot W_{g} \in R^{n}$ の最大値上位 $K$ 個以外のすべての座標に $- \infty$ を割り当てます: $(TopK (ℓ))_{i} := {ℓ_{i} - \infty if ℓ_{i} が ℓ の上位 K 個の座標に含まれる場合, それ以外 .$

レイヤー全体の出力 $y$ は、エキスパート出力のゲーティング重み付け総和として計算されます: $y = i = 0 \sum n - 1 G (x)_{i} \cdot SwiGLU_{i} (x)$ MoEの計算をスパース行列乗算として表現するMegaBlocksなどのカーネルと、複数GPUにまたがるExpert Parallelismを併用することで、効率的な実行が実現されています。

結果（主要な数値）

一般知識および推論ベンチマーク: MMLU（5-shot）において、Mixtral 8x7Bは70.6%を達成し、Llama 2 70Bの69.9%やGPT-3.5の70.0%を上回りました。GSM8K（8-shot, maj@8）では74.4%に達し、Llama 2 70Bの69.6%を凌駕しています。MATHベンチマーク（4-shot, maj@4）では、Mixtralは28.4%を記録し、Llama 2 70Bの13.8%を大きく引き離しました。
コード生成: MixtralはHumanEval（40.2% vs 29.3%）およびMBPP（60.7% vs 49.8%）において、Llama 2 70Bを大幅に上回りました。
多言語タスク: 多言語データのアップサンプリングにより、MixtralはARC-Challenge、HellaSwag、MMLUにおいて、フランス語、ドイツ語、スペイン語、イタリア語全般でLlama 2 70Bを凌駕しました（例：フランス語MMLUで70.9% vs 64.3%）。
インストラクションファインチューニングとArena Elo: Mixtral 8x7B – Instruct（SFTおよびDirect Preference Optimizationで学習）はMT-Benchで8.30を記録し、LMSys Chatbot Arena Eloでは1121を達成して、Claude-2.1（1117）、GPT-3.5-Turbo（1117）、Llama-2-70b-chat（1077）を上回りました。
コンテキスト検索: 32kのコンテキストウィンドウ全体にわたる人工的なPasskey検索タスクにおいて、100%の検索精度を達成しました。

限界と未解決の課題

メモリフットプリントと計算量のトレードオフ: 計算コストは13Bのアクティブパラメータに比例するものの、サービング時には47Bすべてのパラメータを保持するのに十分なVRAMが必要となり、マルチGPU構成やオフロード手法が必須となります。
Expert Parallelismのボトルネック: 論文では高い時間的局所性が発見されており、連続するトークンがランダムな確率を大幅に上回る頻度で同じエキスパートに割り当てられます（第15層ではランダム時の12.5%に対し、最大28.4%の第1選択の重複）。これにより、Expert Parallelism下で特定のGPUに負荷が偏り、過剰サブスクリプションが生じる可能性があります。
トピック特化の欠如: ルーターの選択は概念的なドメイン（例：生物学 vs 哲学）ごとに自然に分割されるわけではないため、ドメイン固有のエキスパートの専門分化を明示的に行いたい場合に、どのようにゲーティングを構造化すべきかという課題が残されています。

重要用語

Sparse Mixture of Experts (SMoE): すべてのパラメータを活性化するのではなく、条件付き計算によって各入力トークンをフィードフォワードサブネットワーク（エキスパート）の疎なサブセットにルーティングするニューラルネットワークアーキテクチャ。
Active Parameter Count: 単一トークンの推論時に実際にロードされて計算に使用されるパラメータ数。Mixtralではスパースパラメータの総数47Bに対して13B。
SwiGLU: Swishと要素ごとの積を組み合わせたGated Linear Unit活性化関数のバリエーションであり、各エキスパートのフィードフォワードブロックの内部アーキテクチャとして使用される。
TopK Gating: トークン表現に線形射影を適用し、上位K個の最大値を保持して残りを負の無限大にマスクした上で、Softmaxによって正規化するルーティングメカニズム。
Expert Parallelism (EP): 異なるエキスパートを別々のGPUにマッピングし、トークンのアクティベーションを割り当てられたエキスパートが存在するデバイスへとルーティングする分散並列化戦略。
Direct Preference Optimization (DPO): 明示的な強化学習の報酬モデルを学習させることなく、選好ペアデータを用いて言語モデルのパラメータを直接最適化するアライメント手法。
Passkey Retrieval: 最大32kトークンまでの長いコンテキストプロンプト内のランダムな位置に配置された短いPasskeyを特定し、正しく返すモデルの能力を測定する人工評価テスト。
Positional Locality: シーケンス内で連続するトークンが、一様ランダムな確率を大幅に上回って同じエキスパートに繰り返し割り当てられる観察された傾向。

よくある質問

パラメータ数が56Bではなく47Bであるのに、なぜMixtralは8x7Bと呼ばれるのですか？

層ごとに8つのエキスパートが存在し、アーキテクチャがMistral 7Bに基づいているものの、Self-Attentionの重みと埋め込み層はすべてのエキスパート間で共有されています。フィードフォワードサブブロックのみが8つ複製されているため、8 × 7B = 56Bではなく、総パラメータ数は47Bになります。

推論中、トークンごとにいくつのエキスパートが活性化されますか？

各層において、ルーターは利用可能な8つのエキスパートの中から各トークンに対して正確に2つのエキスパート（ $K = 2$ ）を選択するため、順伝播ごとのアクティブパラメータは13Bになります。

異なるエキスパートは異なる学術分野やドピックごとに専門分化していますか？

明らかなドメイン専門分化は見られませんでした。The Pileを用いた分析では、ArXiv、PubMed、PhilPapersなどの主題間において、エキスパートの割り当て分布は事実上同一であることが示されました。ルーティングの決定はドメインではなく、構文、トークンの同一性、および位置の局所性に強く相関しています。

Mixtral 8x7Bモデルはどのようなソフトウェアライセンスの下でリリースされていますか？

Mixtral 8x7BベースモデルとMixtral 8x7B – Instructはどちらも、学術および商用利用が可能な寛容なApache 2.0ライセンスの下で公開されています。

このページの解説は AI が論文から作成したもので、誤りを含む場合があります。詳細は原文をご確認ください。

XYZNOT で読む（全文翻訳と AI 解説）