Mixtral of Experts

Reading guide

Problem and motivation

Dense autoregressive language models incur computational costs directly proportional to their total parameter count for every processed token. As models scale up to 70B parameters and beyond, serving them at low latency and high throughput becomes computationally expensive. The authors aim to scale model capacity while keeping per-token inference cost comparable to much smaller models, creating an open-weights model capable of matching or outperforming Llama 2 70B and proprietary baselines like GPT-3.5.

Key ideas

• Sparse Mixture of Experts (SMoE): Replace dense feedforward network (FFN) blocks across all 32 transformer layers with 8 distinct expert networks, routing each token to only the top-2 experts.
• Decoupled Parameter Counts: The architecture holds 47B total parameters across its layers, but only executes 13B active parameters per token forward pass, matching the computational footprint of a moderate-sized dense model.
• Full 32k Context Window: Native support for 32,768 token sequence lengths with 100% retrieval accuracy on passkey benchmarks.
• Syntactic Rather than Topic-Based Routing: Empirical inspection of token routing reveals expert selection is driven primarily by syntactic function and temporal locality rather than high-level topic domains.

Method (with the key equations)

Mixtral modifies the Mistral 7B architecture ($32$ layers, hidden dimension $4096$, $32$ attention heads, $8$ key-value heads) by substituting every FFN sub-block with an MoE module consisting of $n=8$ experts ($E_0,\dots ,E_{n-1}$), each implemented as a SwiGLU network with intermediate dimension $14336$.

For an input token representation $x$, a router network computes a sparse gating distribution across experts using the top-$K$ logits ($K=2$): $$G(x):=\text{Softmax}(\text{TopK}(x\cdot W_g))$$ where the Top-K operator assigns $-\infty$ to all coordinates outside the largest $K$ values of logits $\ell =x\cdot W_g\in {\mathbb{R}}^n$: $$(\text{TopK}(\ell ){)}_i:=\{\begin{array}{ll}{\ell }_i& \text{if }{\ell }_i\text{ is in the top-}K\text{ coordinates of }\ell ,\\ -\infty & \text{otherwise.}\end{array}$$

The layer's overall output $y$ is computed as the gating-weighted sum over expert outputs: $$y=\sum _{i=0}^{n-1}G(x{)}_i\cdot {\text{SwiGLU}}_i(x)$$ Efficient execution is achieved through kernels like MegaBlocks, which represent MoE computation as sparse matrix multiplications, alongside Expert Parallelism across multiple GPUs.

Results (with the most important numbers)

• General & Reasoning Benchmarks: On MMLU (5-shot), Mixtral 8x7B achieves 70.6% compared to 69.9% for Llama 2 70B and 70.0% for GPT-3.5. On GSM8K (8-shot, maj@8), it reaches 74.4% versus 69.6% for Llama 2 70B. On the MATH benchmark (4-shot, maj@4), Mixtral achieves 28.4% versus Llama 2 70B's 13.8%.
• Code Generation: Mixtral outperforms Llama 2 70B substantially on HumanEval (40.2% vs 29.3%) and MBPP (60.7% vs 49.8%).
• Multilingual Tasks: Upsampling multilingual data enables Mixtral to exceed Llama 2 70B across French, German, Spanish, and Italian on ARC-Challenge, HellaSwag, and MMLU (e.g., French MMLU: 70.9% vs 64.3%).
• Instruction Fine-Tuning & Arena Elo: Mixtral 8x7B – Instruct (trained with SFT and Direct Preference Optimization) scores 8.30 on MT-Bench and achieved an LMSys Chatbot Arena Elo of 1121, outperforming Claude-2.1 (1117), GPT-3.5-Turbo (1117), and Llama-2-70b-chat (1077).
• Context Retrieval: Achieves 100% retrieval accuracy on the synthetic passkey retrieval task across the entire 32k context window.

Limitations and open questions

• Memory Footprint vs. Compute: While compute cost scales with the 13B active parameters, serving requires enough VRAM to hold all 47B parameters, demanding multi-GPU setups or offloading techniques.
• Expert Parallelism Bottlenecks: The paper discovers high temporal locality—consecutive tokens share the same expert assignments much more frequently than random chance (up to 28.4% first-choice repetition at layer 15 vs 12.5% random). This can cause load imbalance and over-subscription of specific GPUs under Expert Parallelism.
• Lack of Topic Specialization: Router selection does not naturally partition by conceptual domains (e.g., biology vs. philosophy), leaving open questions on how to structure gating if domain-specific expert specialization is explicitly desired.