推荐一篇非常优质的技术科普文章,Maarten Grootendorst的《A Visual Guide to Mixture of Experts (MoE)》用超过50张可视化图,把MoE技术的原理解释得非常易懂。
虽然文章较老(去年10月发的),但还是非常适合对MoE一知半解的人从头理解。
⬇️AI导读:
## MoE:一场关于“专家”与“路由”的精妙协作
MoE,即“专家混合模型”,本质上是一种提升LLM质量的技术,它通过整合多个子模型(也就是“专家”)来实现。构成MoE的核心,无非是两部分:
**专家(Experts)**: 这些“专家”并非我们想象中那种专精于“心理学”或“生物学”的领域大拿。它们通常是独立的FFNN(前馈神经网络)层,其“专长”在于处理特定上下文中的特定token。简单说,它们更擅长处理词汇层面的语法信息,而非某个宏大领域的知识。
**路由器(Router)或门控网络(gate network)**: 这是MoE的“大脑”,负责决定哪些token应该被发送给哪些专家处理。它会根据输入生成一个概率分布,然后选择最匹配的专家。
想象一下,一个标准的Transformer解码器架构中,FFNN层是处理上下文信息的关键。传统的FFNN是“密集模型”,意味着所有参数都会被激活。而MoE则将其拆解成多个“专家”,形成“稀疏模型”——每次只激活其中一部分专家。这样一来,模型在训练时,每个专家学习不同的信息;在推理时,只调用最相关的专家。
一个token在MoE模型中,会像走迷宫一样,穿过多个解码器块,每个块里都有不同的专家组合。这意味着不同的token会走上不同的“路径”,最终生成文本。这种动态选择,让模型在保持性能的同时,大大提升了效率。
## 路由的艺术与负载均衡的挑战
路由机制是MoE的灵魂,它决定了模型如何高效地利用这些专家。但问题也随之而来:如果路由器总是偏爱某些专家,而冷落另一些,就会导致专家利用率不均,甚至有些专家根本得不到充分训练。这在训练和推理阶段都会造成问题。
为了解决这个问题,**负载均衡**成了关键。目标是让所有专家在训练和推理时都能获得同等的重要性。
* **KeepTopK策略**: 这是最直接的负载均衡方法之一。它通过引入可训练的(高斯)噪声,并设定一个阈值K,只保留概率最高的K个专家,将其他专家的权重设为负无穷,从而确保它们在SoftMax后概率为0。这样,每次只会激活少数几个专家。
* **Token Choice**: KeepTopK策略的延伸,它允许将每个token路由到一个(top-1路由)或多个(top-k路由)选定的专家。这种方式的好处是,可以对不同专家的贡献进行加权和整合。
* **辅助损失(Auxiliary Loss)**: 也叫“负载均衡损失”,它被添加到网络的常规损失函数中。通过计算每个专家的“重要性分数”和“变异系数(CV)”,辅助损失会强制专家之间保持同等的重要性。CV值越高,说明专家间重要性差异越大,模型就会努力降低CV,实现均衡。
除了专家选择的均衡,**专家容量(Expert Capacity)**也是一个重要考量。如果某个专家接收的token过多,超出了它的处理能力,多余的token就会被发送给下一个专家,甚至直接进入下一层(这被称为“token溢出”)。限制专家容量,能有效防止过度训练和资源浪费。
## 从Switch Transformer到Soft-MoE:架构的演进
MoE的早期应用存在训练不稳定的问题,而**Switch Transformer**的出现,简化了MoE的架构和训练流程,显著提升了稳定性。它将传统的FFNN层替换为“Switching Layer”,这是一个稀疏MoE层,每个token只选择一个专家(Top-1路由)。Switch Transformer还引入了**容量因子(Capacity Factor)**,直接影响专家容量,避免了计算资源的浪费或性能下降。其简化的辅助损失函数,通过平衡token分派比例与路由器概率,进一步实现了均匀路由。
MoE并非语言模型的专属。**Vision-MoE (V-MoE)**将MoE引入了视觉模型(如ViT)。ViT将图像分割成小块(patches),这些patch被当作token处理。V-MoE用稀疏MoE替换了ViT编码器中的密集FFNN,通过增加专家数量,实现了视觉模型的规模化扩展。为了应对图像patch数量庞大和低容量带来的token溢出,V-MoE引入了**批次优先级路由(Batch Priority Routing)**,优先处理更重要的patch,确保关键信息不丢失。
更进一步,**Soft-MoE**则尝试从离散的token分配转向“软性”分配。它通过将输入(patch嵌入)与一个可学习矩阵Φ相乘,生成“路由器信息”,然后通过SoftMax更新patch嵌入,使其成为所有patch嵌入的加权平均。这样,所有patch都被“混合”后发送给每个专家,输出再与路由器矩阵相乘,实现了“软性”的token处理,而非硬性分配。
## Mixtral 8x7B:稀疏参数与活跃参数的权衡
MoE模型最吸引人的地方,莫过于其独特的计算需求。以**Mixtral 8x7B**为例,它拥有8个专家,每个专家大小为5.6B参数。这意味着,你需要加载**8 x 5.6B (46.7B)**的“稀疏参数”(即模型总参数)到内存中。但在推理时,它只会激活其中的2个专家,所以实际使用的“活跃参数”只有**2 x 5.6B (12.8B)**。
这揭示了一个关键的权衡:MoE模型需要更多的显存来加载所有专家,但推理速度却更快,因为它只激活了一小部分参数。这对于追求高性能和高效率的大模型部署来说,无疑是一条充满潜力的路径。