大模型数学速成(18):MoE——稀疏 FFN 与路由
第 06 篇 的 FFN 对每个 token 做同一套「升维 → 非线性 → 降维」。模型想更大时,若继续把 FFN 做厚,参数量和每 token 算力一起涨。MoE(Mixture of Experts,混合专家) 换思路:总参数可以很多,但每个 token 只激活少数专家。
这是「大模型数学速成」续篇第 18 篇。建议先读第 06 篇 FFN。下一篇(收官)讲 DPO / RLHF 目标函数。
一、稠密 FFN 的成本
标准 Transformer 块里,FFN 往往占参数大头。稠密(dense)含义是:
每个 token 都完整跑同一套 FFN 权重。
于是:
| 放大方式 | 参数 | 每 token FLOPs |
|---|---|---|
| 加宽 / 加深稠密 FFN | ↑ | ↑(同步) |
| MoE 稀疏 FFN | ↑(多专家) | 主要随 激活的专家数 涨 |
MoE 追求的是:参数规模与单次前向算力部分解耦。
二、专家与门控
把原来的一个 FFN 换成 $N$ 个结构相同的 专家 $E_1,\ldots,E_N$,再加一个 路由器(router / gate)。
对 token 特征 $\mathbf{x}$:
- 路由器打分:$\mathbf{g} = \mathrm{Router}(\mathbf{x})$(常为线性 + Softmax 或稀疏变体)
- 选 top-$k$ 专家(常见 $k=1$ 或 $2$)
- 只跑被选中的专家,按门控权重混合:
\[\mathbf{y} = \sum_{i \in \mathcal{T}(\mathbf{x})} \tilde{g}_i(\mathbf{x})\, E_i(\mathbf{x})\]
| 符号 | 含义 |
|---|---|
| $E_i$ | 第 $i$ 个专家 FFN |
| $\mathcal{T}(\mathbf{x})$ | 被选中的专家下标集合,$\|\mathcal{T}\|=k$ |
| $\tilde{g}_i$ | 选中专家上的归一化门控权重 |
1 | token x ──► Router ──► 选 top-k |
稀疏:未选中的 $E_i$ 本步不算——参数仍占显存(要加载),但算力按激活集走。
三、手算:2 专家、top-1
设路由器对某 token 打分(未归一化)$z = [2.0,\ 0.5]$,专家输出(标量示意)$E_1=1.0$,$E_2=3.0$。
Softmax:
\[g \approx [0.818,\ 0.182]\]
top-1 只留专家 1,并在选中集上重归一(此处仅 1 个,权重为 1):
\[y = 1 \cdot E_1 = 1.0\]
若 top-2 都保留:
\[y = 0.818\cdot 1.0 + 0.182\cdot 3.0 \approx 1.364\]
同一套专家,$k$ 不同,算力与输出都不同。
四、负载均衡(为何需要)
若路由器总把 token 塞给少数「热门」专家:
- 这些专家过载,其它专家吃不饱、学不好
- 分布式推理时还造成计算不均
训练时常加 负载均衡辅助损失(形式因论文而异):鼓励专家被选频率接近均匀。
本篇只需记住:稀疏路由要防塌缩到单一专家。
五、与稠密模型对照
| 稠密 FFN | MoE FFN | |
|---|---|---|
| 参数量 | 与宽度同步 | 可随专家数 $N$ 放大 |
| 每 token 算力 | 全 FFN | 约 $k$ 个专家 |
| 实现复杂度 | 低 | 路由、均衡、并行 |
| 典型叙事 | 「类标准 Transformer」 | 「类 Mixtral 的稀疏 FFN」(公开架构举例,非内部配置) |
注意力层仍可稠密;许多开源 MoE 主要是 把 FFN 换成专家混合。
六、和推理的关系(点到为止)
- 显存:专家权重要驻留(或卸载),总参数大 ≠ 单步算力大
- 并行:专家可分布在多卡,路由带来通信
- 与 第 15 篇量化、第 10 篇 GQA 正交,可叠用
部署选型属推理工程路线;此处只建立 稀疏激活 的数学图像。
七、小结
| 概念 | 要点 |
|---|---|
| 动机 | 参数↑ 而每 token FLOPs 不线性↑ |
| 结构 | 多专家 FFN + 路由器 |
| 公式 | $y=\sum_{i\in\mathcal{T}} \tilde g_i E_i(x)$ |
| top-$k$ | 控制激活宽度 |
| 均衡 | 防路由塌缩 |
大模型数学速成续篇第 18 篇完。下一篇 DPO / RLHF 目标函数速览——从 CE 接到对齐目标,系列收官。
系列导航
| 篇号 | 标题 | 状态 |
|---|---|---|
| 17 | Flash Attention | ✅ |
| 18 | MoE:稀疏 FFN 与路由(本篇) | ✅ |
| 19 | DPO / RLHF 目标函数速览 | ✅ 收官 |
完整大纲见工作区 docs/MATH_SERIES_OUTLINE_V2.md。









