[CS336-04] Mixture of Experts

这篇是 Stanford CS336 Spring 2025 第 4 讲 Course Materials 的学习笔记。主题是 Mixture of Experts，也就是 MoE。

1. MoE 的核心想法#

MoE 最常见的做法是把 Transformer block 里的 dense FFN 换成多个 expert FFN，再用 router 为每个 token 选择少数几个 expert。

它的关键优势是：可以增加总参数量，但每个 token 只激活一小部分参数。因此在相近 FLOPs 下，MoE 可以拥有更大的模型容量。

dense FFN: 每个 token 都走同一个大 FFN
MoE FFN: 每个 token 只走 top-k 个 expert

这也是 MoE 近年来重新流行的主要原因：同等训练计算下，更多参数和更稀疏的激活可能带来更好的效果。

2. 为什么 MoE 适合大模型时代#

大模型训练同时受 compute、memory、network 约束。MoE 的吸引力在于它把“参数规模”和“每 token 计算量”部分解耦。

这带来几个收益：

• 可以在不线性增加 FLOPs 的情况下扩展参数量；
• expert 可以自然分布到多卡或多机上；
• 对同等训练预算，MoE 可能比 dense model 更快达到目标 loss；
• 在推理时也可以只激活少数 expert，保持较低单 token 计算量。

但 MoE 不是免费午餐。它把简单 dense matmul 变成了动态路由、稀疏计算和跨设备通信问题。

3. Router：token 到 expert 的分配#

Lecture 4 把 routing 作为 MoE 的核心设计。主流方法是 token choice top-k：每个 token 经过 gating network 后选择得分最高的 k 个 expert。

常见设置包括：

• Switch Transformer：top-1；
• GShard、Mixtral、Grok：top-2；
• Qwen、DBRX、DeepSeek 等使用更多 active experts 或 shared experts。

Routing 也可以反过来由 expert 选择 token，或者通过全局优化做 assignment，但主流大模型里 token choice top-k 更常见，因为实现和扩展更直接。

4. Load balancing 和训练稳定性#

MoE 的一个问题是 expert collapse：router 可能总是选择少数 expert，导致负载不均、部分 expert 学不到东西、硬件利用率下降。

因此 MoE 训练常引入额外 loss 或约束，例如 load balancing loss、router z-loss、capacity factor 等。这些目标通常比较 heuristic，但在工程上很重要。

系统层面还要处理 token dispatch 和 combine：token 被送到不同 expert，expert 计算后结果再聚合回来。跨设备 MoE 会引入 all-to-all 通信，这对网络拓扑和实现库提出更高要求。

5. Fine-grained experts 与 shared experts#

近年的 MoE 模型还探索了更细粒度的 expert、更大的 expert 总数、少量 shared experts 等设计。

Shared expert 的直觉是：有些通用能力不应该完全依赖 router 选择，可以让部分 expert 对所有 token 始终可用。Fine-grained expert 则把容量切得更细，让 router 有更多组合空间。

不同论文的 ablation 并不总是一致，这说明 MoE 仍然是强经验驱动领域。

6. takeaway#

MoE 可以理解成一种“稀疏激活的容量扩展方法”：

• 它用更多总参数换更强表达能力；
• 用 top-k routing 控制每 token FLOPs；
• 用 expert parallelism 把模型自然分布到多设备；
• 代价是训练目标、路由稳定性和系统通信都更复杂。

如果 dense Transformer 的主要问题是参数和计算强绑定，那么 MoE 的价值就是把二者拆开。但拆开之后，新的瓶颈会转移到 routing、load balance 和 all-to-all 通信。

参考#

• Stanford CS336 Spring 2025 Course Materials: lecture 4.pdf