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 讲解）

三句话总结

🎯 在不产生过高推理延迟和计算成本的情况下，扩展 Transformer 的容量和性能。
🔧 采用稀疏混合专家（SMoE）仅解码器架构，用 8 个 SwiGLU 专家替代标准前馈块，通过可学习的路由器为每个 token 选择 2 个专家。
📈 拥有 47B 总参数，每个 token 仅激活 13B，在 MMLU（70.6% vs 69.9%）和 GSM8K（74.4% vs 69.6%）上超越 Llama 2 70B。

摘要

我们推出了 Mixtral 8x7B，这是一种稀疏混合专家（SMoE）语言模型。Mixtral 拥有与 Mistral 7B 相同的架构，不同之处在于每一层由 8 个前馈模块（即专家）组成。对于每个 token，在每一层中，路由网络会选择两个专家来处理当前状态并合并它们的输出。尽管每个 token 仅由两个专家处理，但在每个时间步选定的专家可以不同。因此，每个 token 可以访问 47B 参数，但在推理过程中仅使用 13B 活跃参数。Mixtral 在 32k token 的上下文长度下进行训练，并在所有评估基准测试中优于或匹敌 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 模型。基础模型和指令模型均在 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.

论文解读

问题与动机

稠密自回归语言模型处理每个 token 所产生的计算成本，与其总参数量直接成正比。随着模型扩展到 70B 参数及以上，以低延迟和高吞吐量进行服务变得极其昂贵。作者旨在扩展模型容量的同时，将每个 token 的推理成本保持在与更小模型相当的水平，打造一个性能媲美或超越 Llama 2 70B 以及 GPT-3.5 等专有基准的开源权重模型。

核心理念

稀疏混合专家（SMoE）： 将全部 32 个 Transformer 层的稠密前馈网络（FFN）模块替换为 8 个独立的专家网络，每个 token 仅路由至 Top-2 专家。
解耦参数量： 该架构跨层包含 47B 总参数，但每个 token 的前向传播仅执行 13B 活跃参数，相当于中等规模稠密模型的计算开销。
完整 32k 上下文窗口： 原生支持 32,768 token 的序列长度，在 passkey 基准测试中实现了 100% 的检索准确率。
偏向句法而非主题的路由： 对 token 路由的实证分析表明，专家的选择主要受句法语义功能和时间局部性驱动，而非高层级的主题领域。

方法（包含核心公式）

Mixtral 对 Mistral 7B 架构（ $32$ 层，隐藏维度 $4096$ ， $32$ 个注意力头， $8$ 个键值头）进行了修改，将每个 FFN 子模块替换为一个包含 $n = 8$ 个专家（ $E_{0}, \dots, E_{n - 1}$ ）的 MoE 模块，每个专家均实现为中间维度为 $14336$ 的 SwiGLU 网络。

对于输入的 token 表示 $x$ ，路由网络使用 top- $K$ logits（ $K = 2$ ）计算专家之间的稀疏门控分布： $G (x) := Softmax (TopK (x \cdot W_{g}))$ 其中 Top-K 算子将 logits $ℓ = x \cdot W_{g} \in R^{n}$ 中最大 $K$ 个值之外的所有坐标赋值为 $- \infty$ ： $(TopK (ℓ))_{i} := {ℓ_{i} - \infty if ℓ_{i} is in the top- K coordinates of ℓ, otherwise.$

该层的总输出 $y$ 计算为专家输出的门控加权和： $y = i = 0 \sum n - 1 G (x)_{i} \cdot SwiGLU_{i} (x)$ 高效执行是通过诸如 MegaBlocks 等算子实现的，该算子将 MoE 计算表示为稀疏矩阵乘法，并结合了跨多个 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，并取得了 1121 的 LMSys Chatbot Arena Elo 分数，超越了 Claude-2.1（1117）、GPT-3.5-Turbo（1117）和 Llama-2-70b-chat（1077）。
长上下文检索： 在整个 32k 上下文窗口的人工 passkey 检索任务中，均达到 100% 的检索准确率。

局限性与未决问题

显存占用与计算量的矛盾： 尽管计算成本与 13B 活跃参数相当，但服务部署需要足够的显存来容纳全部 47B 参数，这就需要多 GPU 配置或参数卸载（offloading）技术。
Expert Parallelism 的瓶颈： 论文发现了高度的时间局部性——连续的 token 分配给相同专家的频率远超随机概率（在第 15 层，首选专家重复率高达 28.4%，而随机期望仅为 12.5%）。这在 Expert Parallelism 下可能会导致负载不均衡以及特定 GPU 负载过重。
缺乏主题特异性： 路由器的选择并不会自然地按概念领域划分（如生物学与哲学），如果明确需要特定领域的专家分工，如何构建门控机制仍是一个未决问题。

关键术语

Sparse Mixture of Experts (SMoE): 一种神经网络架构，其中条件计算将每个输入的 token 路由到前馈子网络（专家）的稀疏子集，而不是激活所有参数。
Active Parameter Count: 在单个 token 的推理过程中实际加载并参与计算的参数数量。在 Mixtral 中，相对于其 47B 的总稀疏参数，该数值为 13B。
SwiGLU: 一种门控线性单元激活变体，结合了 Swish 和逐元素乘法，用作每个专家前馈块的内部架构。
TopK Gating: 一种路由机制，对 token 表示应用线性投影，保留前 K 个最大值，同时将其余值掩码为负无穷大，并通过 softmax 进行归一化。
Expert Parallelism (EP): 一种分布式切分策略，将不同的专家映射到不同的 GPU 上，并在设备之间将 token 激活值路由到其分配的专家。
Direct Preference Optimization (DPO): 一种对齐算法，无需训练显式的强化学习奖励模型，直接在偏好对数据上优化语言模型参数。
Passkey Retrieval: 一种综合评估测试，用于测量模型在长达 32k token 的长上下文提示中，定位并返回随机放置的短密钥（passkey）的能力。
Positional Locality: 观察到的一种趋势，即序列中连续的 token 被重复分配给同一个专家的概率远高于均匀随机概率。

常见问题

为什么 Mixtral 只有 47B 参数却被称为 8x7B？

虽然每层有 8 个专家且架构源自 Mistral 7B，但自注意力权重和嵌入层由所有专家共享。只有前馈子模块被复制了 8 次，因此总参数量为 47B，而不是 8 × 7B = 56B。

在推理期间每个 token 会激活多少个专家？

对于每个 token，路由器在每一层可用的 8 个专家中精确选择 2 个专家（ $K = 2$ ），因此每次前向传播的活跃参数为 13B。

不同的专家是否专攻不同的学术或领域主题？

没有发现明显的领域专业化。在 The Pile 数据集上的分析表明，在 ArXiv、PubMed 和 PhilPapers 等主题中，专家的分配分布几乎完全相同；相反，路由决策与句法、token 标识和位置局部性密切相关。

Mixtral 8x7B 模型遵循什么软件许可证发布？

Mixtral 8x7B 基础模型和 Mixtral 8x7B – Instruct 均在宽松的 Apache 2.0 许可证下发布，可用于学术研究和商业用途。

本页解读由 AI 根据论文生成，可能有错误；细节请以原文为准。

用 XYZNOT 阅读（全文翻译与 AI 讲解）