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

Paper · 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.

Veröffentlicht 2024-01-08cs.LGcs.CL arXiv PDF

Mit Übersetzung und KI lesen

Zusammenfassung in drei Zeilen

🎯 Skalierung von Transformer-Kapazität und -Leistung ohne untragbare Inferenzlatenz und Rechenkosten.
🔧 Eine reine Sparse-Mixture-of-Experts-(SMoE)-Decoder-Architektur, die Standard-Feedforward-Blöcke durch 8 SwiGLU-Experten ersetzt und über einen gelernten Router 2 Experten pro Token auswählt.
📈 Greift auf insgesamt 47B Parameter zu, nutzt aber nur 13B aktive Parameter pro Token und übertrifft Llama 2 70B auf MMLU (70,6 % vs. 69,9 %) sowie GSM8K (74,4 % vs. 69,6 %).

Abstract

Wir stellen Mixtral 8x7B vor, ein Sparse Mixture of Experts (SMoE) Sprachmodell. Mixtral besitzt dieselbe Architektur wie Mistral 7B, mit dem Unterschied, dass jede Schicht aus 8 Feedforward-Blöcken (d. h. Experten) besteht. Für jedes Token wählt ein Router-Netzwerk auf jeder Schicht zwei Experten aus, um den aktuellen Zustand zu verarbeiten und deren Ausgaben zu kombinieren. Obwohl jedes Token nur zwei Experten durchläuft, können sich die ausgewählten Experten zu jedem Zeitschritt unterscheiden. Infolgedessen hat jedes Token Zugriff auf 47B Parameter, nutzt jedoch während der Inferenz nur 13B aktive Parameter. Mixtral wurde mit einer Kontextlänge von 32k Token trainiert und übertrifft oder erreicht Llama 2 70B sowie GPT-3.5 über alle evaluierten Benchmarks hinweg. Insbesondere übertrifft Mixtral Llama 2 70B in Mathematik, Codegenerierung und mehrsprachigen Benchmarks deutlich. Zudem stellen wir mit Mixtral 8x7B - Instruct ein instruktionsoptimiertes Modell bereit, das GPT-3.5 Turbo, Claude-2.1, Gemini Pro und das Llama 2 70B - Chat-Modell in menschlichen Benchmarks übertrifft. Sowohl das Basis- als auch das Instruct-Modell werden unter der Apache-2.0-Lizenz veröffentlicht.

Original-Abstract (Englisch)

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.

Leseleitfaden

Problem und Motivation

Dichte autoregressive Sprachmodelle verursachen Rechenkosten, die für jedes verarbeitete Token direkt proportional zu ihrer Gesamtparameteranzahl sind. Bei einer Skalierung auf 70B Parameter und mehr wird das Serving mit niedriger Latenz und hohem Durchsatz rechenintensiv und kostspielig. Die Autoren zielen darauf ab, die Modellkapazität zu vergrößern und gleichzeitig die Inferenzkosten pro Token auf dem Niveau deutlich kleinerer Modelle zu halten. So soll ein Open-Weights-Modell geschaffen werden, das mit Llama 2 70B und proprietären Baselines wie GPT-3.5 mithalten oder diese übertreffen kann.

Zentrale Ideen

Sparse Mixture of Experts (SMoE): Ersetzung dichter Feedforward-Netzwerk-Blöcke (FFN) über alle 32 Transformer-Schichten hinweg durch 8 distinkte Expertennetzwerke, wobei jedes Token nur an die Top-2-Experten geroutet wird.
Entkoppelte Parameterzahlen: Die Architektur verfügt über insgesamt 47B Parameter in ihren Schichten, aktiviert jedoch nur 13B aktive Parameter pro Vorwärtsdurchlauf eines Tokens, was dem Rechenaufwand eines mittelgroßen dichten Modells entspricht.
Vollständiges 32k-Kontextfenster: Native Unterstützung für Sequenzlängen von 32.768 Token mit 100 % Retrieval-Genauigkeit bei Passkey-Benchmarks.
Syntaktisches statt themenbasiertes Routing: Empirische Analysen des Token-Routings zeigen, dass die Expertenauswahl primär durch syntaktische Funktionen und zeitliche Lokalität gesteuert wird und nicht durch übergeordnete Themendomänen.

Methode (mit den zentralen Gleichungen)

Mixtral modifiziert die Architektur von Mistral 7B ( $32$ Schichten, Hidden Dimension $4096$ , $32$ Attention Heads, $8$ Key-Value Heads), indem jeder FFN-Unterblock durch ein MoE-Modul ersetzt wird. Dieses besteht aus $n = 8$ Experten ( $E_{0}, \dots, E_{n - 1}$ ), die jeweils als SwiGLU-Netzwerk mit einer Zwischendimension von $14336$ implementiert sind.

Für eine Eingabe-Token-Repräsentation $x$ berechnet ein Router-Netzwerk eine dünnbesetzte Gating-Verteilung über die Experten mittels der Top- $K$ -Logits ( $K = 2$ ): $G (x) := Softmax (TopK (x \cdot W_{g}))$ wobei der Top-K-Operator allen Koordinaten außerhalb der größten $K$ Werte der Logits $ℓ = x \cdot W_{g} \in R^{n}$ den Wert $- \infty$ zuweist: $(TopK (ℓ))_{i} := {ℓ_{i} - \infty wenn ℓ_{i} unter den Top- K -Koordinaten von ℓ liegt, andernfalls.$

Die Gesamtausgabe der Schicht $y$ wird als Gating-gewichtete Summe über die Ausgaben der Experten berechnet: $y = i = 0 \sum n - 1 G (x)_{i} \cdot SwiGLU_{i} (x)$ Eine effiziente Ausführung wird durch Kernel wie MegaBlocks erreicht, die MoE-Berechnungen als dünnbesetzte Matrixmultiplikationen abbilden, ergänzt durch Expert Parallelism über mehrere GPUs hinweg.

Ergebnisse (mit den wichtigsten Zahlen)

Allgemeine und Reasoning-Benchmarks: Bei MMLU (5-shot) erzielt Mixtral 8x7B 70,6 % im Vergleich zu 69,9 % bei Llama 2 70B und 70,0 % bei GPT-3.5. Bei GSM8K (8-shot, maj@8) erreicht es 74,4 % gegenüber 69,6 % bei Llama 2 70B. Im MATH-Benchmark (4-shot, maj@4) erzielt Mixtral 28,4 % gegenüber 13,8 % bei Llama 2 70B.
Codegenerierung: Mixtral übertrifft Llama 2 70B bei HumanEval (40,2 % vs. 29,3 %) und MBPP (60,7 % vs. 49,8 %) deutlich.
Mehrsprachige Aufgaben: Durch das Upsampling mehrsprachiger Daten übertrifft Mixtral Llama 2 70B in Französisch, Deutsch, Spanisch und Italienisch bei ARC-Challenge, HellaSwag und MMLU (z. B. französisches MMLU: 70,9 % vs. 64,3 %).
Instruction Fine-Tuning & Arena Elo: Mixtral 8x7B – Instruct (trainiert mit SFT und Direct Preference Optimization) erreicht einen Score von 8,30 auf MT-Bench und einen LMSys Chatbot Arena Elo-Wert von 1121, womit es Claude-2.1 (1117), GPT-3.5-Turbo (1117) und Llama-2-70b-chat (1077) übertrifft.
Kontext-Retrieval: Erreicht 100 % Retrieval-Genauigkeit bei der synthetischen Passkey-Retrieval-Aufgabe über das gesamte 32k-Kontextfenster hinweg.

Limitationen und offene Fragen

Speicherbedarf vs. Rechenaufwand: Während die Rechenkosten mit den 13B aktiven Parametern skalieren, erfordert das Serving genügend VRAM, um alle 47B Parameter im Speicher zu halten, was Multi-GPU-Setups oder Offloading-Verfahren notwendig macht.
Engpässe bei Expert Parallelism: Die Studie stellt eine hohe zeitliche Lokalität fest – aufeinanderfolgende Token teilen sich dieselben Expertenzuweisungen deutlich häufiger als durch reinen Zufall erwartet (bis zu 28,4 % Wiederholung der Erstwahl auf Schicht 15 gegenüber 12,5 % bei Gleichverteilung). Dies kann bei Expert Parallelism zu Ungleichgewichten bei der Lastverteilung und einer Überlastung einzelner GPUs führen.
Fehlende thematische Spezialisierung: Die Auswahl durch den Router segmentiert nicht erkennbar nach konzeptionellen Fachgebieten (z. B. Biologie vs. Philosophie). Dies wirft offene Fragen auf, wie das Gating strukturiert werden sollte, wenn eine domänenspezifische Expertenspezialisierung explizit gewünscht ist.

Schlüsselbegriffe

Sparse Mixture of Experts (SMoE): Eine neuronale Netzwerkarchitektur, bei der konditionale Berechnung jedes Eingabe-Token an eine dünnbesetzte Teilmenge von Feedforward-Unternetzwerken (Experten) weiterleitet, anstatt alle Parameter zu aktivieren.
Active Parameter Count (Aktive Parameteranzahl): Die Anzahl der Parameter, die während der Inferenz für ein einzelnes Token tatsächlich geladen und berechnet werden; bei Mixtral sind dies 13B im Vergleich zu seinen 47B gesamten Sparse-Parametern.
SwiGLU: Eine Aktivierungsvariante der Gated Linear Unit, die Swish und elementweise Multiplikation kombiniert und als interne Architektur für jeden Experten-Feedforward-Block dient.
TopK Gating: Ein Routing-Mechanismus, der eine lineare Projektion auf die Token-Repräsentation anwendet, die Top-K höchsten Werte beibehält, alle übrigen auf minus unendlich setzt und das Ergebnis via Softmax normalisiert.
Expert Parallelism (EP): Eine verteilte Partitionierungsstrategie, die unterschiedliche Experten auf getrennte GPUs verteilt und Token-Aktivierungen über Geräte hinweg zu ihren zugewiesenen Experten routet.
Direct Preference Optimization (DPO): Ein Alignment-Algorithmus, der die Parameter von Sprachmodellen direkt auf Präferenzpaaren optimiert, ohne ein explizites Reinforcement-Learning-Belohnungsmodell zu trainieren.
Passkey Retrieval: Ein synthetischer Evaluierungstest, der die Fähigkeit eines Modells misst, einen kurzen Passkey zu finden und zurückzugeben, der zufällig in langen Kontext-Prompts von bis zu 32k Token platziert wurde.
Positional Locality (Positionale Lokalität): Die beobachtete Tendenz, dass aufeinanderfolgende Token in einer Sequenz weitaus häufiger demselben Experten zugewiesen werden, als es bei einer gleichmäßigen Zufallsverteilung der Fall wäre.

Fragen

Warum heißt Mixtral 8x7B, wenn es nur 47B statt 56B Parameter hat?

Obwohl es 8 Experten pro Schicht gibt und die Architektur von Mistral 7B abgeleitet ist, werden die Self-Attention-Gewichte und Embedding-Schichten von allen Experten gemeinsam genutzt. Lediglich die Feedforward-Unterblöcke sind 8-fach repliziert, was zu 47B Gesamtparametern anstelle von 8 × 7B = 56B führt.

Wie viele Experten werden bei der Inferenz für jedes Token aktiviert?

Der Router wählt auf jeder Schicht für jedes Token genau 2 Experten ( $K = 2$ ) aus den 8 verfügbaren Experten aus, was zu 13B aktiven Parametern pro Vorwärtsdurchlauf führt.

Spezialisieren sich unterschiedliche Experten auf verschiedene akademische oder fachliche Themenbereiche?

Es wurde keine offensichtliche Domänenspezialisierung festgestellt. Analysen auf The Pile zeigten nahezu identische Verteilungen der Expertenzuweisungen über Fächer wie ArXiv, PubMed und PhilPapers hinweg; Routing-Entscheidungen korrelieren stattdessen stark mit Syntax, Token-Identität und zeitlicher Lokalität.

Unter welcher Softwarelizenz werden die Mixtral 8x7B-Modelle veröffentlicht?

Sowohl das Basismodell Mixtral 8x7B als auch Mixtral 8x7B – Instruct werden unter der permissiven Apache-2.0-Lizenz für akademische und kommerzielle Zwecke veröffentlicht.

Die Übersicht auf dieser Seite wurde von einer KI aus dem Paper erstellt und kann Fehler enthalten. Details stehen im Paper.

Mit Übersetzung und KI lesen