MoE(Mixture of Experts,混合专家)
MoE(Mixture of Experts,混合专家)是目前大模型架构中非常重要的一个方向,GPT-4、Mixtral、DeepSeek、Qwen的部分版本都用了这个思路。我从核心思想讲到具体细节。
核心思想
传统的Transformer模型里,每一层的前馈网络(FFN)对所有输入token都用同一套参数计算。MoE的想法是:把这一个大的FFN换成很多个小的"专家"网络,每个token只激活其中一小部分专家,而不是全部。
打个比方:传统模型像一个全科医生,什么病人来了都要自己看一遍所有科室的知识;MoE则像一家医院,来了病人先经过分诊台(路由),分诊台判断这个病人该去哪几个科室,只有对应的专家医生会诊,其他科室的医生完全不参与。
这样做的好处是:模型的总参数量可以做得很大(知识容量大),但每次前向计算实际用到的参数量很小(计算成本低)。这就是"稀疏激活"。
基本架构组成
一个标准的MoE层通常替换Transformer中的FFN层,包含两个部分:
专家网络(Experts):一组结构相同、参数不同的前馈网络(FFN),数量可以是8个、16个、64个甚至上百个。
门控网络/路由器(Gating Network / Router):一个很小的网络,输入是token的隐藏状态,输出是每个专家的"打分"(通常经过softmax),决定这个token应该交给哪几个专家处理。MoE最有意思的地方就是"信息怎么从一个token流向被选中的专家"。
关键技术细节
1. Top-k路由
门控网络对每个token输出一个在所有专家上的概率分布(通常用softmax),然后只取分数最高的k个专家(常见k=1或k=2)。被选中专家的输出会按照各自的权重加权求和,作为这一层的最终输出。
2. 负载均衡问题
这是MoE训练中最大的坑。如果不加约束,路由器很容易"偷懒",把大部分token都送给少数几个专家,导致这些专家训练得很好、其他专家几乎没学到东西(即"专家坍塌")。为了解决这个问题,通常会加入:
- 辅助负载均衡损失(auxiliary load balancing loss):惩罚专家之间负载不均衡的情况,鼓励token更均匀地分配到各个专家。
- 专家容量限制(expert capacity):给每个专家设置一个最大处理token数上限,超出的token会被"丢弃"(drop)或路由到下一个候选专家。
3. 通信开销
在分布式训练中,不同专家往往分布在不同的GPU上,token需要跨设备发送给对应的专家,这会带来额外的通信成本(all-to-all通信),是MoE工程实现上的一大难点。
MoE vs 稠密(Dense)模型对比
| 维度 | 稠密模型 | MoE模型 |
|---|---|---|
| 每次前向计算的参数量 | 等于总参数量 | 远小于总参数量(仅激活的专家) |
| 总参数量/知识容量 | 相对较小 | 可以做得很大 |
| 训练/推理速度 | 参数量越大越慢 | 计算成本与激活参数量成正比,速度更快 |
| 训练稳定性 | 相对简单 | 需要处理负载均衡、路由不稳定等问题 |
| 显存占用 | 与计算量匹配 | 总参数都要加载到显存,占用大 |
简单说:MoE用"显存换计算",总参数量很大(需要更多显存存储所有专家),但实际推理时每个token只用一小部分参数计算,速度反而更快。
实际应用案例
- Switch Transformer(Google):早期大规模验证MoE可行性的代表工作,每个token只路由到1个专家(Top-1)。
- Mixtral 8x7B(Mistral AI):8个专家,每次激活2个,总参数约47B,但实际计算量只相当于约13B的稠密模型。
- DeepSeek-MoE / DeepSeek-V2/V3:引入了"细粒度专家"和"共享专家"的设计,把专家拆得更小更多,同时设置若干个所有token都必经的共享专家来捕捉通用知识,进一步提升效率。
- Qwen、Grok等也都采用了MoE架构的版本。
优缺点总结
优点
- 在相同计算预算下可以获得远大于稠密模型的参数量,理论知识容量更强
- 推理时激活参数少,速度和成本优于同等总参数量的稠密模型
- 不同专家可以隐式地"专业化",学习到不同类型的模式
缺点
- 训练不稳定,需要精心设计负载均衡机制
- 显存占用高(所有专家参数都要加载)
- 分布式训练/推理工程复杂度高(通信开销、专家并行)
- 微调时容易出现某些专家过拟合的问题