杀疯了的开源专家模型 Mixtral 8x7B 论文公开啦!
作者 | 付奶茶、王二狗
上个月法国初创公司 Mistral AI 开源的一个8x7B MoE模型Mixtral 8x7B引爆了AI社区。
一是因为它的性能击败了LLama2和GPT-3.5,二是因为,Mixtral是一个专家混合模型(Mixture-of-Experts model,MoE),使用了传闻中 OpenAI 用于GPT-4的架构,但是参数相比GPT-4要小很多,堪称是缩小版“GPT-4”。
而就在这几天,Mixtral 8x7B发布了模型背后的论文,下面就带大家一起来深入了解 Mixtral 8x7B的独特之处。
论文题目:
Mixtral of Experts
论文链接:
https://arxiv.org/abs/2401.04088
Github代码:
https://github.com/mistralai/mistral-src
官方主页:
https://mistral.ai/news/mixtral-of-experts/
模型架构
Mixtral架构的显著特征是每层由8个前馈块组成,这些“专家层”取代了传统Transformer模型中的前馈网络(FFN)块。这些前馈块在每一层由路由网络为每个token选择,并且在推断时,每个token仅与其中两个专家层互动。这种独特的设计使得每个token能够接触到47B个参数,实际上却只使用了13B个活跃参数。
对于给定的输入x,MoE模块的输出是由专家网络输出的加权和决定,这些权重由门控网络的输出给出。即,对于n个专家网络{, , ..., },专家层的输出可以通过以下公式计算得出:
这里,代表门控网络对第i个专家层的n维输出的权重,而是第i个专家网络的输出。如果门控向量是稀疏的,那么可以避免计算那些门控值为零的专家层的输出。为了实现,作者采用了一种简单而高效的方法,即对线性层的Top-K逻辑值应用softmax函数:
K值——即每个token使用的专家层数量——是一个调节每个token处理计算量的关键超参数。通过在保持K值不变的情况下增加n值,可以在保持计算成本大致不变的同时增加模型的参数数量。这导致了模型的总参数数量(通常称为稀疏参数数量)随n的增加而增长,而处理单个token的参数数量(称为活跃参数数量)则随K的增加而增长,直到达到n的数量。
MoE层可以在单个GPU上通过高性能的专门内核高效运行。例如,Megablocks技术将MoE层的前馈网络(FFN)操作转化为大型稀疏矩阵乘法,从而显著提高执行速度,并自然地处理了分配给不同专家的不同数量token的情况。此外,MoE层还可以通过标准的模型并行技术以及一种称为专家层并行(EP)的特殊分区策略分布到多个GPU上。在MoE层执行期间,分配给特定专家处理的token会被路由到相应的GPU上进行处理,然后专家层的输出会被返回到原始token的位置。Mixtral采用与专家函数相同的SwiGLU架构,并设置K为2。这意味着每个token被路由到两个具有不同权重集的SwiGLU子块。因此,对于输入token x,其输出y的计算如下:
实验
Mixtral模型在其训练中采用了32k token的上下文长度,并在多项评估基准测试中展现出了卓越性能,其表现不仅匹敌,甚至超越了Llama 2 70B和GPT-3.5。特别值得一提的是,在数学、代码生成和多语言等领域的基准测试中,Mixtral的表现远超过Llama 2 70B。
此外,研究团队还推出了一款经过特定指令优化的Mixtral模型——Mixtral 8x7B – Instruct。在人类评估基准测试中,这一版本超过了GPT-3.5 Turbo, Claude-2.1, Gemini Pro和Llama 2 70B – chat。
该模型在多语言数据上进行了预训练,尤其在涵盖法语、德语、西班牙语和意大利语的多语言基准测试中,明显优于Llama 2 70B。
此外,论文还对模型在偏见基准测试中的表现进行了评估。结果显示,在BBQ基准测试中,Mixtral展示了较少的偏见,而在BOLD基准测试中,则显示了更加平衡的情感特性。
论文还深入探讨了模型的路由机制。研究人员对路由器选择的专家层进行了小规模分析,关注于训练过程中是否有专家层针对特定领域(如数学、生物学、哲学等)发展出专门化。通过在The Pile验证数据集的不同子集上测量被选中专家的分布,结果显示在模型的第0层、第15层和第31层(即模型的首尾层)并没有出现明显的基于主题的专家分配模式。例如,分配给ArXiv论文(LaTeX编写)、生物学(PubMed摘要)和哲学(PhilPapers)文档的专家分布在所有层级中都非常相似。
尽管如此,研究人员在分析DM数学数据时,注意到专家层分布的细微差异。这种差异可能是由数据集的合成性质以及其在自然语言谱中覆盖范围有限导致的,特别是在模型的第一层和最后一层,这些层的隐藏状态与输入和输出嵌入紧密相关。这表明路由器在某种程度上表现出结构化的句法行为。在不同领域(如Python代码、数学和英语)的文本样本中,每个token都用不同颜色的背景高亮显示,以表示其被选中的专家层。
例如,在Python和英语中,像‘self’和‘Question’这样的单词通常通过同一个专家进行路由,即便它们涉及多个token。同样,在代码中,如缩进符号这样的token总是被分配给同一专家层,尤其是在与模型输入和输出更为相关的第一层和最后一层。
研究人员还观察到,连续的token经常被分配给同一专家层。这种现象在The Pile数据集中表现出一定程度的位置局部性。数据表明,在每个领域和层级中,连续token被分配给同一专家层的比例明显高于随机分配的情况,尤其是在模型的高层。这一发现对于理解如何优化模型以提高训练和推断速度具有重要意义。
总结
该论文介绍了首个在开源模型中实现最先进性能的混合专家网络 Mixtral 8x7B。它的指令执行版本Mixtral 8x7B Instruct,在人类评估基准测试中超过了Claude-2.1、Gemini Pro和GPT-3.5 Turbo。
Mixtral的一个显著优势在于其每个时间步骤仅需利用两个专家层,因此每个token仅使用13B活跃参数,而此前每个token需要70B参数的顶级模型Llama 2 70B相比,其性能已大幅超越。研究团队已经将他们经过训练和微调的模型在Apache 2.0许可下公开发布,欢迎大家快快试用~
微信扫码关注该文公众号作者