[Paper Review] Mixtral of Experts

Albert Q. Jiang, Alexandre Sablayrolles+24 more|arXiv (Cornell University)|Jan 8, 2024

Topic Modeling115 citations

TL;DR

Mixtral 8x7B is a sparse mixture-of-experts language model that uses 8 experts per layer with 2 chosen per token, achieving 47B sparse and 13B active parameters, outperforming Llama 2 70B and GPT-3.5 on many benchmarks, with an instruction-tuned variant that surpasses several chat models.

ABSTRACT

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.

Motivation & Objective

• Motivate the development of open, high-performing open-weight SMoE language models.
• Propose Mixtral 8x7B as an open, multilingual, long-context decoder-only SMoE with efficient inference.
• Demonstrate that per-token computation is kept small (K=2, active params = 13B) while overall parameter count is large (47B).
• Show that Mixtral achieves state-of-the-art or competitive results across math, code, multilingual tasks, and instruction-following scenarios.

Proposed method

• Describe Mixtral 8x7B architecture: 32 layers, 8 experts per MoE layer, two experts selected per token by a router.
• Explain gating: Softmax(TopK(x · W_g)) with K=2 to select experts per token.
• Replace FFN blocks with MoE blocks; use SwiGLU as the expert function.
• Context length extended to 32768 tokens; open-source weights under Apache 2.0.
• Detail training with multilingual data and context window, plus an instruction-tuned variant Mixtral 8x7B – Instruct.
• Outline deployment considerations: Megablocks kernels, vLLM integration, and EP load-balancing considerations.

Figure 1: Mixture of Experts Layer. Each input vector is assigned to 2 of the 8 experts by a router. The layer’s output is the weighted sum of the outputs of the two selected experts. In Mixtral, an expert is a standard feedforward block as in a vanilla transformer architecture.

Experimental results

Research questions

• RQ1 Can Mixtral 8x7B achieve competitive or superior performance to larger dense models (e.g., Llama 2 70B) across standard benchmarks?
• RQ2 What are the efficiency and scalability benefits of sparse Mixture-of-Experts when using a fixed per-token active parameter budget?
• RQ3 Does instruction-tuning (Mixtral Instruct) yield superior human-evaluation performance relative to open-weight rivals?
• RQ4 How does Mixtral perform on long-context tasks and multilingual benchmarks?
• RQ5 What bias and safety characteristics does Mixtral exhibit relative to comparable open models?

Key findings

• Mixtral 8x7B outperforms or matches Llama 2 70B on most benchmarks, especially in mathematics and code generation.
• Active parameters per token are 13B, with a total sparse parameter count of 47B.
• Mixtral–Instruct surpasses GPT-3.5 Turbo, Claude-2.1, Gemini Pro, and Llama 2 70B–chat on human benchmarks.
• Significant multilingual gains over Llama 2 70B in French, German, Spanish, and Italian.
• Long-context tests show 100% passkey retrieval accuracy across context lengths; perplexity improves with longer context.
• Mixtral exhibits reduced bias and more positive sentiment in BBQ/BOLD benchmarks compared to Llama 2 70B.

Figure 2: Performance of Mixtral and different Llama models on a wide range of benchmarks . All models were re-evaluated on all metrics with our evaluation pipeline for accurate comparison. Mixtral outperforms or matches Llama 2 70B on all benchmarks. In particular, it is vastly superior in mathemat