MoE:为什么只让每个 Token 找少数专家?

从 FFN、Router 到负载均衡,看懂“参数多、计算少”如何成立

目录

  上一课说到:FFN 会分别加工每个 Token。普通的稠密模型让所有 Token 使用同一套 FFN。MoE(Mixture of Experts,混合专家)则准备多套 FFN,让一个小型 Router 为每个 Token 选其中少数几套。[1]

MoE 想增加可用的参数容量,同时控制每个 Token 实际执行的 FFN 数量。

Token 经过 Router 选择少数 FFN 专家,再合并输出

一个 Token 如何选择并经过专家,点击查看原始 SVG

先看问题:如果每个 Token 都用全部专家

  假设某层有 8 套 FFN。如果每个 Token 都要使用这 8 套,参数容量是增加了,但这一层的 FFN 计算也会跟着涨。上下文很长或并发很多时,代价会很高。

  MoE 的想法是:8 套 FFN 都可供选择,但每个 Token 只经过其中 1 套或 2 套。这样,模型拥有的总参数与一次计算中真正用到的激活参数就成了两个不同的数。Switch Transformer 选择每个 Token 只经过一套专家;其他 MoE 结构也常用 Top-2 或更大的 Top-K。[2]

专家是什么

  在常见的 Transformer MoE 中,专家一般是一套独立的 FFN。某些层的 FFN 被 MoE 层替换;Attention 和模型的其他部分仍按各自结构工作。专家之间有不同权重,但输入输出维度相同,所以它们的结果可以合并。

  这里的“专家”只是架构里的名称。它不保证有一位“数学专家”、一位“历史专家”,也不能看到某个 Token 走向专家 3,就断定专家 3 专管这个词。路由偏好和内部功能是训练形成的,需要专门分析才能解释。

Router 怎样给 Token 选专家

  仍用“我喜欢苹果”做例子。经过前面的 Attention 后,三个位置已经拥有不同的上下文表示。Router 读取每个位置的表示,为每位专家算分,再保留分数最高的 K 位。路由是逐 Token、逐 MoE 层发生的,因此一句话里的不同 Token 可以走向不同专家。

  以 4 位专家、Top-2 为例,某个 Token 的路由分数可能是:

专家 路由分数(示意) 是否执行
专家 1 0.08 否
专家 2 0.52 是
专家 3 0.31 是
专家 4 0.09 否

  这个 Token 只送进专家 2 和专家 3。两位专家都输出一个向量,模型再用路由权重加权合并,写回当前 Token 的表示:

\[y_t=\sum_{i\in\operatorname{TopK}(s_t)} g_{t,i}\,E_i(x_t)\]

  $s_t$ 是 Router 给出的分数,$E_i$ 是第 $i$ 位专家,$g_{t,i}$ 是被选中专家的合并权重。权重怎样归一化、是否另有共享专家,取决于具体模型。DeepSeek-V2 就同时使用共享专家和被路由选择的专家。[3]

动手试试:同一句话也会走不同路线

选择 Token:激活数量:
专家 1
专家 2
专家 3
专家 4

  交互里的分数是为了看清 Top-K 选择过程而设定的。真实 Router 会随着模型权重、上下文、层数变化;同一个词在不同句子里,也可能走不同路线。

为什么 K 远小于 E

  用 $E$ 表示可选专家总数,用 $K$ 表示每个 Token 激活的专家数。当 $K\ll E$,专家池可以扩大,而每个 Token 的专家计算大致只随 $K$ 增长。Router、共享层和合并操作仍有成本,因此这只是理解计算规模的近似。

方案 可选专家 每个 Token 执行 大致专家计算
全部执行 8 8 8 份 FFN
Top-2 路由 8 2 2 份 FFN
Top-1 路由 8 1 1 份 FFN

  这张表解释了 MoE 的“参数多、每 Token 计算相对少”。它不意味着模型权重只占 1/8 的显存:未被当前 Token 选中的专家仍要存放在某处。描述 MoE 时,最好同时给出总参数量和每 Token 激活参数量。例如 DeepSeek-V2 报告了 236B 总参数与每 Token 21B 激活参数;这两个数字不能混成一个。[3]

为什么路由还需要负载均衡

  假如 Router 总把 Token 送给专家 2,专家 2 会排队,其他专家却很闲。热门专家得到了更多训练样本,冷门专家更难学会有用的分工,拥堵可能进一步加重。

  因此 MoE 训练通常要考虑负载均衡,并为每位专家设置或管理容量。Switch Transformer 研究了辅助负载均衡损失和专家容量;实际系统还会处理超出容量的 Token。具体是跳过、改路由还是采取其他策略,必须看模型与实现,不能把其中一种行为当成 MoE 的统一规则。[2]

少算不一定更快

  当专家分布在不同 GPU 上,Token 要先被送到对应设备,专家算完后再把结果送回。这种分发与回收通常涉及跨设备通信。专家计算减少了,通信、调度、数据搬运和负载不均仍可能决定实际速度。[1]

  所以评价一个 MoE,至少要同时问:质量如何?总参数和激活参数各是多少?路由是否均衡?专家是否跨卡?端到端延迟与吞吐到底是多少?只看“激活参数少”还无法判断实际部署效果。

容易混淆的四件事

  1. 专家通常是一套 FFN。 它不等于一位被人工命名的学科专家。
  2. Router 是逐 Token 选路。 一个句子里不同位置可以经过不同专家。
  3. 稀疏激活节省本次专家计算。 全部专家的权重仍需存放。
  4. 计算少不保证延迟低。 路由、通信和拥堵都会影响速度。

80 秒讲解视频

参考文献(References)

  1. Shazeer, N., Mirhoseini, A., Maziarz, K., et al. “Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer.” International Conference on Learning Representations, 2017. arXiv:1701.06538.↩
  2. Fedus, W., Zoph, B., and Shazeer, N. “Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity.” Journal of Machine Learning Research, 23(120):1–39, 2022. arXiv:2101.03961.↩
  3. DeepSeek-AI, Liu, A., Feng, B., et al. “DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model.” arXiv preprint arXiv:2405.04434, 2024.↩

第十五课复习总图

MoE 第十五课复习总图

Router、Top-K、负载与通信的一页复习图,点击查看原始 SVG

  下一课转向 GPU 执行:FlashAttention 为什么数学结果不变,却能更快、更省显存?