Mixture of Experts（MoE）

Mixture of ExpertsMoE1. 使用背景大模型这几年有一个很明显的趋势参数越大能力往往越强。但问题也很直接。传统稠密模型里参数一旦变大训练和推理的计算量也会跟着一起涨。也就是说模型虽然更强了但代价也越来越高。所以后面就有一个很自然的想法能不能把模型总参数做得很大但每次前向时只激活其中一小部分参数MoE就是在这个背景下出现的。它最核心的思想可以概括成一句话把一个大层拆成很多个专家experts每个token只路由给其中少数几个专家处理。这样一来模型的总参数量可以很大但单个token实际用到的激活参数量和计算量不一定同步变大。2. 理论基础1先看普通Transformer里的FFN在标准Transformer block里除了attention之外另一个核心部分就是前馈网络FFN。一个典型的FFN可以写成FFN(x)W2 σ(W1x) \text{FFN}(x)W_2\,\sigma(W_1x)FFN(x)W2​σ(W1​x)这里xxx是当前token的隐藏状态W1,W2W_1,W_2W1​,W2​是可学习参数σ(⋅)\sigma(\cdot)σ(⋅)表示激活函数。在稠密模型里每个token都会经过同一套FFN参数。也就是说不管输入token是什么所有token共享同一个专家。2MoE到底改了什么MoE最核心的改动就是把原来“只有一个FFN”的位置换成“有很多个专家FFN”的位置。假设有NNN个专家那么第iii个专家可以写成Ei(x) E_i(x)Ei​(x)一般每个EiE_iEi​本质上还是一个FFN只不过各自参数不同。与此同时再加一个router / gate决定当前token应该送给哪些专家。所以MoE层最本质的结构就是两部分一组专家网络一个路由器。3. 从稠密FFN到MoE层1最一般的形式如果不考虑稀疏性MoE层可以写成所有专家加权求和MoE(x)∑i1Ngi(x) Ei(x) \text{MoE}(x)\sum_{i1}^{N} g_i(x)\,E_i(x)MoE(x)i1∑N​gi​(x)Ei​(x)其中gi(x)g_i(x)gi​(x)表示router给第iii个专家分配的权重。但如果每次都把所有专家都算一遍那计算量还是很大这样就失去MoE的意义了。2稀疏激活所以MoE真正关键的地方在于不是所有专家都激活而是只选top-k个专家。设router输出分数后只保留分数最高的kkk个专家那么MoE层可以写成MoE(x)∑i∈TopK(x)g~i(x) Ei(x) \text{MoE}(x)\sum_{i\in \text{TopK}(x)} \tilde g_i(x)\,E_i(x)MoE(x)i∈TopK(x)∑​g~​i​(x)Ei​(x)这里g~i(x)\tilde g_i(x)g~​i​(x)表示在被选中的专家上重新归一化后的权重。这就是MoE最重要的点总参数很多但每个token只走少数几个专家。3直观理解可以把它理解成普通FFN像是“一个老师教所有学生”MoE更像是“有很多老师但每个学生只分给其中几个最合适的老师”。所以MoE本质上是一种条件计算conditional computation。4. Router是怎么工作的1最常见的router形式对输入token表示xxxrouter一般先做一个线性映射再接softmaxssoftmax(Wrx) s\text{softmax}(W_r x)ssoftmax(Wr​x)其中s∈RNs\in \mathbb{R}^Ns∈RN表示当前token对每个专家的打分。然后根据sss选择top-k个专家。所以router做的事情很简单判断当前token更适合交给哪些专家处理。2top-1和top-2MoE里最常见的两种路由方式就是top-1 routing每个token只选1个专家top-2 routing每个token选2个专家。top-1更便宜、更简单top-2通常表达能力更强一些但通信和实现也更复杂。所以后面很多MoE工作一个核心差别其实就在这里到底每个token激活几个专家。3为什么router很关键因为MoE强不强不只是看专家数量更看路由是否合理。如果router总把大量token送去同一个专家那别的专家就学不到东西如果router分得太平均所有专家又可能变得差不多专门化不明显。所以router其实决定了MoE能不能真正形成“专家分工”。5. MoE最大的难点负载均衡1为什么会不均衡在理想情况下我们希望不同专家都能被充分利用。但现实里router很容易偏向少数专家。结果就是有些专家特别忙容量不够有些专家几乎没token进去学不到东西整个训练会不稳定甚至效率很差。这就是MoE里最经典的问题之一load imbalance负载不均衡。2capacity问题在实际实现里每个专家通常都会设置一个容量上限也就是每一步最多接收多少token。如果某个专家分到的token太多超出的token就可能被丢弃或者被迫做额外处理。所以router不只是“选谁”还牵涉到“会不会把某个专家挤爆”。3辅助负载均衡损失为了缓解这个问题很多MoE方法会额外加一个auxiliary loss让router尽量不要过度偏向少数专家。直观上理解这个loss就是在鼓励两件事专家被选中的频率不要太失衡router的概率分布不要长期塌到少数专家上。所以MoE里常说的aux loss本质上就是给router加“别太偏科”的约束。6. 代表性工作1Sparsely-Gated MoE最早把大规模稀疏专家层做起来2017年的Sparsely-Gated MoE可以看作现代MoE的起点。它最核心的贡献就是把“很多专家 稀疏门控 条件计算”这套东西真正做成了可训练的大规模层。这篇工作给人的一个很强信号就是参数容量可以涨得非常大但计算代价不一定按同样比例上涨。也就是说MoE第一次很明确地把“参数量”和“单次计算量”拆开了。2GShard把MoE放进Transformer大规模训练GShard的重要性在于它不是只提出一个MoE层而是把MoE和大规模Transformer训练真正接起来了。它用的是top-2 routing并结合自动分片和并行训练把MoE推到了更大规模。所以GShard这一步很关键因为它说明了MoE不只是一个层级想法而是真的可以用来训练超大模型。3Switch Transformer把router大幅简化Switch Transformer可以说是MoE发展里非常关键的一步因为它做了一个非常干脆的改动把top-2 routing简化成top-1 routing。也就是说每个token只去一个专家。这个改动听起来好像只是“少选一个专家”但影响很大计算更简单通信成本更低实现更直接训练稳定性更容易控制。所以Switch的一个核心贡献就是把MoE从“能做”推到“更容易做”。4ST-MoE更关注训练稳定性和可迁移性后面ST-MoE这类工作开始更系统地讨论MoE到底怎样才能训得更稳、迁移得更好。这说明MoE发展到后面问题已经不只是“能不能堆很大”而是训练稳不稳fine-tuning稳不稳sparse模型能不能真正拿到transfer收益。所以ST-MoE更像是对MoE工程化和实用化的一次系统整理。5Mixtral开源LLM里最有代表性的MoE之一Mixtral出来以后MoE在开源LLM里的存在感一下子变强了很多。它最典型的特征就是有8个专家每个token选2个专家只激活一部分参数但整体能力很强。Mixtral让很多人第一次直观感受到MoE不是“研究里很大但不好用”的东西而是可以在实际大语言模型里做出非常强性能的架构。6DeepSeekMoE更强调专家专门化DeepSeekMoE这条线的关注点很明确不仅要稀疏还要让专家更专门化。它一个很有代表性的想法是把专家分得更细同时引入shared experts去吸收通用知识把剩下的routed experts留给更细粒度、更专门的分工。这个思路其实很自然。因为在很多MoE里一个常见问题就是专家虽然分开了但功能上还是很重叠。DeepSeekMoE就是在往前再走一步不仅让专家存在还要让专家真正各干各的。7. MoE为什么有效1参数量和计算量被拆开了这是MoE最根本的价值。在稠密模型里参数量大前向计算通常也会跟着变大但在MoE里虽然总参数很多单个token只激活少数专家所以实际计算不一定同步膨胀。这就给了模型一个很有吸引力的性质可以用较低的激活计算换取更高的总参数容量。2不同token可能需要不同知识从直觉上看语言里不同token、不同上下文本来就可能需要不同类型的知识和变换。比如代码token、数学token、普通叙述token它们需要的模式未必一样。那么让所有token都用同一套FFN其实未必最合理。MoE则允许模型学出某种“分工处理”的机制。3更像条件计算我觉得MoE最自然的一种理解就是不是所有计算都值得对每个token重复一遍。有些计算只对某类输入更有用那就让对应专家去做。这本质上是一种资源分配问题。8. MoE的问题出在哪1训练不稳定MoE虽然强但一开始并不好训。router、aux loss、capacity、通信模式这些东西都会影响训练稳定性。所以MoE不是简单“把FFN复制很多份”就完事了它背后有一整套训练技巧。2通信成本高在分布式训练里MoE经常会遇到一个很现实的问题token被路由到不同专家而不同专家可能分布在不同设备上。这就会带来大量all-to-all通信。所以很多时候MoE的瓶颈不只是算力而是路由之后的数据搬运。3专家未必真的专门化这是一个很有意思的问题。理论上大家都希望不同专家学不同知识但如果router设计不好或者负载均衡约束过强可能会出现很多专家学得差不多专门化不明显专家利用率不高。所以后来很多工作都在想办法增强expert specialization。4推理并不总是“白赚”MoE常说“总参数很大但激活参数少”这没错。但在真实系统里推理效率还会受很多因素影响expert并行方式路由带来的访存和通信开销batch size和token分布稀疏实现是否足够高效。所以MoE虽然在理论上很省但工程上并不是自动就会特别省。9. 从更高一层看MoE到底是什么1它不是简单“加参数”MoE当然是在增加参数但它更重要的一点不是“参数更多”本身而是参数是有条件地被激活的。也就是说MoE真正带来的变化不是静态规模而是动态计算路径。2它本质上是一种条件计算架构给定不同输入模型可以走不同专家路径。所以MoE更像是在Transformer里嵌入了一种“输入驱动的模块选择机制”。从这个角度看MoE不是单纯的大模型技巧而是一种更一般的架构思想。3为什么它在LLM里特别重要因为LLM已经大到一个阶段大家既想继续扩大总容量又不想让每个token的计算无限增长。MoE正好卡在这个矛盾点上给了一个非常自然的解法总容量继续涨但每个token只付部分计算代价。10. 一点理解1MoE最漂亮的地方我觉得MoE最漂亮的地方就是它把一个原本绑在一起的东西拆开了总参数量和单次激活计算量。这件事一拆开很多原来做不到的规模就有机会做到了。2它为什么看起来简单但一直有人做表面上看MoE好像就是“复制很多个FFN再加个router”。但真正难的地方在于router怎么训负载怎么平衡通信怎么做专家怎么真正分工稀疏计算怎么落到系统里。所以MoE不是一个“概念一提出就结束”的方向而是一条很长的工程和算法路线。3怎么记MoE如果只是为了学习我觉得可以把MoE记成四句话稠密模型里所有token共享同一个FFN计算随着参数量一起涨MoE把一个FFN换成很多专家并用router选择少数专家因此总参数可以很大但每个token只激活部分参数真正的难点不在“有很多专家”而在“怎么路由、怎么平衡、怎么让专家真正专门化”。11. 参考鸣谢Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layerhttps://arxiv.org/abs/1701.06538GShard: Scaling Giant Models with Conditional Computation and Automatic Shardinghttps://arxiv.org/abs/2006.16668Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsityhttps://arxiv.org/abs/2101.03961ST-MoE: Designing Stable and Transferable Sparse Expert Modelshttps://arxiv.org/abs/2202.08906Mixture-of-Experts with Expert Choice Routinghttps://arxiv.org/abs/2202.09368Mixtral of Expertshttps://arxiv.org/abs/2401.04088DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Modelshttps://arxiv.org/abs/2401.0606612. 注这篇主要是个人学习整理重点放在主线理解文中主要写的是Transformer里最常见的稀疏MoE形式很多实现细节如capacity factor、token dropping、all-to-all通信优化等没有展开才疏学浅欢迎批评、指导和交流有错误望大家及时指正