Mixtral of Experts

In plain words

Most language models use every part of their network for every word they process. Mixtral changes that inside each of its 32 layers. The feed-forward part of each layer is split into 8 separate sub-networks, called experts, and a small router decides which 2 of them handle each token. The other 6 stay idle for that token.

The effect is a model that stores a lot of knowledge but does only a fraction of the work per token. The paper reports about 47 billion parameters in total and about 13 billion used for any one token. The rest of the design follows its predecessor, Mistral 7B. The model reads up to 32,000 tokens of context, and the authors report that it matches or beats Llama 2 70B and GPT-3.5 on most of the benchmarks they ran. An instruction-tuned version is also released.

The trade-off is memory. All 47 billion parameters must still be stored, and the authors say sparse layers suit batched workloads better.

TL;DR

Mixtral 8x7B is a sparse mixture-of-experts Transformer with open weights under Apache 2.0 §1. Each layer has 8 SwiGLU experts, and a linear router sends each token to the top 2, mixing their outputs with softmax weights §2.1. It has about 47B total and 13B active parameters and a 32k dense context §2 §3. Author-reported results match or beat Llama 2 70B and GPT-3.5 on most benchmarks, and the code is the official mistral-src repository.

Problem

Dense models spend the same compute on every token, so raising quality means raising the cost of every forward pass. The paper aims for a model with more parameters but without a matching rise in per-token compute, and it reports quality that matches or beats Llama 2 70B §1 §2.1. Sparse mixture-of-experts layers offer that, but the paper names two costs: the routing overhead, and the need to balance work across GPUs when experts are spread over devices §2.1. The target context is a fully dense 32k tokens §2.

Method

Mixtral reuses the Mistral 7B architecture, with two stated changes: it keeps a fully dense 32k context, and every feed-forward block is replaced by a mixture-of-experts layer §2. The main model dimensions are in tab. 1: 32 layers, model width 4096, 32 query heads with 8 key/value heads, and 8 experts with 2 used per token tab. 1.

Sparse mixture-of-experts layer §2.1. The layer output combines the expert outputs, each scaled by a gate value from a gating network [§2.1, display 1]. The paper uses the top-K softmax form: the gate keeps the K largest router logits and applies a softmax over only those §2.1, display 2. For Mixtral, K = 2 §2.1. The expert is a SwiGLU feed-forward block §2.1. The full token output is given in §2.1, display 3.

Steps for one token in one layer. The reference implementation applies this per token, to a whole batch at once, grouped by expert moe.py.

1. The token passes through the attention block of the layer, which is the Mistral 7B style with 8 key/value heads shared by 32 query heads §2 tab. 1 transformer_layers.py.
2. A bias-free linear router maps the token to 8 logits, one per expert §2.1 transformer_layers.py.
3. The top 2 logits are selected, and softmax is taken over those two only, so their weights sum to 1 §2.1, display 2 moe.py.
4. Each selected SwiGLU expert runs on the token. Its output is multiplied by its weight and summed §2.1, display 3 moe.py.
5. The layer's output is added back to the residual stream. The paper does not describe this step in detail, and the normalisation details are in the code only transformer_layers.py.

# Sparse MoE feed-forward for one token  [§2.1, display 2-3]
# Input:  x (dim,); router weight (n_experts, dim); experts (SwiGLU)
# Output: y (dim,)
def moe_feed_forward(x, router_weight, experts, k):
    # Bias-free router logit per expert  [§2.1]
    logits = router_weight @ x                 # (n_experts,)
    # Keep the k largest logits  [§2.1, display 2]
    top_indices = top_k_indices(logits, k)     # (k,)
    # Softmax over kept logits only, so weights sum to 1 [code: moe.py#L26-L27]
    weights = softmax(logits[top_indices])     # (k,)
    y = zeros_like(x)
    # Weighted sum of the selected expert outputs  [§2.1, display 3]
    for weight, index in zip(weights, top_indices):
        y = y + weight * experts[index](x)
    return y

# SwiGLU expert  [§2.1] [code: transformer_layers.py#L96-L106]
# Input:  x (dim,); w1, w3 (hidden_dim, dim); w2 (dim, hidden_dim)
# Output: (dim,)
def swiglu_expert(x, w1, w3, w2):
    # Gate branch passes through SiLU  [code: transformer_layers.py#L105-L106]
    gate = silu(w1 @ x)                        # (hidden_dim,)
    # Linear branch, multiplied elementwise with the gate  [§2.1]
    up = w3 @ x                                # (hidden_dim,)
    return w2 @ (gate * up)                    # (dim,)

Routing and expert placement §2.1. Experts can run on one GPU with sparse kernels (MegaBlocks is cited) or be spread across GPUs with expert parallelism. In the second case, tokens are routed to the GPU that holds their expert and the output is sent back. The paper notes that this can unbalance the work across GPUs.

Context §2 §3.2. The model is trained with a 32k context. The passkey retrieval test in §3.2 reports 100% accuracy regardless of context length or position. Perplexity on a long-document subset falls as context grows §3.2 fig. 4.

Instruction model §4. Supervised fine-tuning on an instruction dataset is followed by Direct Preference Optimization §4. The fine-tuning data and hyperparameters are not stated.

Hyperparameters

Architecture values are from tab. 1. Other values are stated where shown.

Results

Author-reported. The authors re-ran every baseline with their own evaluation pipeline §3. The setting for each row is given in the bullets below.

Evaluation protocol §3. Commonsense reasoning is 0-shot (Hellaswag, Winogrande, PIQA, SIQA, OpenbookQA, ARC-E, ARC-C, CommonsenseQA). World knowledge is 5-shot (NaturalQuestions, TriviaQA). Reading comprehension is 0-shot (BoolQ, QuAC). Math is GSM8K 8-shot with maj@8 and MATH 4-shot with maj@4. Code is HumanEval 0-shot and MBPP 3-shot. Aggregates are MMLU 5-shot, BBH 3-shot and AGI Eval 3-5-shot, English multiple choice only §3.

Base model vs Llama 2 and Mistral 7B tab. 2. Scores are percentages as reported.

Source: tab. 2. The paper's caption says Mixtral matches or beats Llama 2 70B on almost all benchmarks with about 5x fewer active parameters tab. 2. Llama 2 70B is ahead on HellaSwag, WinoGrande and TriviaQA tab. 2.

Against Llama 2 70B and GPT-3.5 tab. 3. These rows use different protocols from tab. 2, as noted in the pitfalls.

Multilingual §3.1 tab. 4. Scores are percentages.

Long context §3.2 fig. 4. Passkey retrieval reaches 100% accuracy in the reported setting, regardless of context length or passkey position. The text does not state the exact lengths tested.

Bias §3.3. BBQ accuracy: LLaMA 2 70B 51.5%, Mixtral 56.0%. BOLD sentiment (average ± standard deviation):

Instruction model §4. MT-Bench score 8.30 (see tab. 3); the authors describe it as the best open-weights model at the time of writing. The authors also report an LMSys Chatbot Arena Elo of 1121 for the instruct model, against 1117 for Claude-2.1 and the best GPT-3.5-Turbo, 1111 for Gemini Pro and 1077 for Llama-2-70b-chat, from a snapshot dated 22 December 2023 §4.

Routing analysis §5 tab. 5. The first-choice expert is picked about 12.5% of the time under uniform routing. Measured first-choice repetition on consecutive tokens is close to random at layer 0 and well above random at layers 15 and 31. Repetition of either first or second choice is about 46% under uniform routing.

Source: tab. 5 §5. The authors find no clear topic specialisation of experts §5 fig. 7.

Paper vs code

Code: the official repository mistralai/mistral-src at code_commit 9eaeb91c17450e09021b6065a1d5cc69876507c8, Apache-2.0 LICENSE. It is linked from the paper's abstract §1. The repository also documents 8x22B checkpoints README.md, which this card does not cover.

• Routing and gating. Consistent. The paper's top-K softmax §2.1, display 2 matches the code's top-k then softmax over the k logits moe.py. The router has no bias §2.1 and is a bias-free linear layer transformer_layers.py.
• Expert function. Consistent. SwiGLU §2.1 matches the code's SiLU-gated feed-forward transformer_layers.py.
• Where MoE applies. Consistent. The paper replaces every feed-forward block §2. The code builds an MoE layer in each transformer block when the model config has moe set transformer_layers.py transformer.py.
• Context length. The paper states a 32k dense context §2 tab. 1. The code's RoPE table is precomputed for 128,000 positions transformer.py. This is a capacity of the code, not a claim in the paper.
• Sliding window. The paper says the context is dense §2. The code supports an optional sliding window args.py, and the value for the released model is not in this repo. It cannot be verified here.
• RoPE base and norm epsilon. Not stated in the paper. The code uses defaults of 1e6 transformer.py and 1e-6 transformer_layers.py when the model config omits them. The released model's config values are not in this repo.
• Evaluation pipeline. The paper uses its own evaluation harness §3. This repo has no benchmark harness (searched for MMLU, HellaSwag and GSM8K). The results in tab. 2 and tab. 3 cannot be reproduced from this repo.
• Training. The repo is for inference. The paper states no training objective for the MoE router, and the repo has none, so it is not stated.

No discrepancies found in the MoE math, expert function or placement.

Pitfalls

• "47B tokens" in the text. The comparison text says the model has "47B tokens compared to 70B" §3. It means 47B parameters. Implement the 47B figure as parameters, not tokens.
• MT-Bench reference. The instruct section points to tab. 2 for the MT-Bench score §4. Table 2 has no MT-Bench column. The 8.30 score is in tab. 3.
• Protocols differ between tables. GSM8K is 8-shot with maj@8 in §3 and tab. 2 (74.4%), but 5-shot in tab. 3 (58.4%). The MBPP number in tab. 2 uses the hand-verified subset §3. Do not compare numbers across tab. 2 and tab. 3 as if they share a protocol.
• TriviaQA without context. The authors did not provide Wikipedia context for TriviaQA, which differs from the Llama 2 paper's setup §3. Their Llama numbers are re-runs, not copied from Llama 2.
• Memory is 47B, compute is 13B. The paper states that memory cost follows the 47B sparse count, which is still below Llama 2 70B §3. It also says the routing and extra memory loads add overhead, and that sparse layers suit batched workloads §3. No latency or throughput measurements are given in the text.
• Renormalised top-2 weights. The gate weights sum to 1 over the two selected experts only. The weights of the other six experts are not part of the output, so the softmax is not over all eight. Code computes the softmax in float32 and casts back to the input dtype moe.py.
• Load balancing is open. The paper names load balancing as a challenge under expert parallelism §2.1. No auxiliary loss, capacity limit or token dropping is described in the paper, and the repo's inference path has none moe.py. Do not assume one.
• Reference loop is not the paper's kernel. The reference MoE layer loops over the experts and scatters results back moe.py. The paper points to sparse kernels (MegaBlocks) and to vLLM integration §1 §2.1. The repo's loop is a different implementation from those.
• Context ceiling. The paper trains at 32k §2. The code allows 128,000 positions in its RoPE table transformer.py. The paper's long-context evaluation does not state the lengths it tested §3.2.
• Routing analysis is descriptive. The routing measurements use subsets of The Pile validation set §5. The paper gives no statistical tests for them, so treat the figures as descriptive.
• Instruct results are snapshots. The MT-Bench and arena numbers depend on the judge and on a dated leaderboard snapshot §4 fig. 6.
• Figure and table numbers. The LaTeX source has no compiled PDF in raw/src. Figure and table numbers here follow caption order in the source. Check them against the PDF before relying on them.
• Equations are unnumbered. The paper's displayed equations have no numbers. They are cited here as "display 1", "display 2" and "display 3" in §2.1, counted in source order.

Lineage

• Builds on: Mistral 7B, whose modifications Mixtral reuses §2. The preference-tuning step uses DPO §4.
• Related sparse-expert work named in the paper: the sparsely-gated MoE layer (arXiv:1701.06538) §2.1; GShard (arXiv:2006.16668), which replaces every other block and uses a different second-expert gate §2.1; MegaBlocks (arXiv:2211.15841) for sparse kernels §2.1. No cards exist for these, so nothing from them is used here.
• Compared against: LLaMA 2 70B (arXiv:2307.09288) tab. 2 tab. 3; GPT-3.5, Claude-2.1 and Gemini Pro §4 (no arXiv ids in the paper).
• Superseded by: not known from wiki content.