训练稳定性、实验结论与局限
Training Stability, Results & Limitations
前两期我们把 Switch Transformer 的机制讲完了：
- EP1：把 FFN 拆成 N 个专家，router 让每个 token 只选 1 个（top-1 routing）
- EP2：用 Expert Capacity 兜底，用 Load Balancing Loss 治本
这一期是这个系列的收尾，回答三个问题：
- 论文为了让它"能训练起来"，还做了哪些事？（只讲在解决什么问题，不讲怎么实现）
- 论文用实验证明了什么？
- 它没有证明什么、有哪些局限？
读完这一期，你应该能用几句话，向别人讲清楚这篇论文干了什么。
一、先给结论
| # | 结论 | 证据来源 |
|---|---|---|
| 1 | 在每个 token 计算量（FLOPs）不变的前提下，只增加专家数量（也就是增加参数量），模型质量持续变好 | 专家数从 2 翻倍到 256，参数从 223M 增长到 14.7B，质量持续提升，但边际收益递减 |
| 2 | top-1 路由够用，而且比 top-2 的 MoE 在"速度–质量"上更划算 | 与 MoE（top-2）的正面对比 |
| 3 | 同等算力下，预训练到同样质量，Switch-Base（64 专家）用时约为 T5-Base 的 1/7 | 时间口径的预训练曲线 |
| 4 | 预训练的收益能迁移到下游任务，但不是所有任务都涨 | 微调结果表 |
| 5 | 大稀疏模型可以蒸馏回小的 dense 模型，保留约 30% 的质量增益 | 蒸馏实验 |
| 6 | 可以训到万亿参数规模，预训练达到固定质量的速度约为 T5-XXL 的 4 倍 | Switch-C（1.6T 参数） |
注意这张表里每一条的前提：都是预训练阶段、同等计算预算下的比较。这一点后面会反复强调，因为它决定了这些结论能推广到哪里。
二、让它"能训练起来"：三个稳定性技巧
EP2 的辅助损失和容量上限，解决的是"路由不均匀"。但 MoE 还有另一类麻烦：训练不稳定。
论文给出的原因有两个：
- 每一层都有 router 的硬切换（argmax 选专家），这是个离散决策
- 低精度格式（bfloat16）会加剧 router 里 softmax 计算的问题
针对这个，论文提了三个办法。这一节只讲每个办法在解决什么问题，不讲工程实现。
2.1 Selective Precision：只有 router 用高精度
问题：全程用 float32 训练最稳，但慢，而且 token 在设备间传输（all-to-all 通信）的数据量翻倍。全程用 bfloat16 快，但论文实验里直接训崩了。
做法：不是"全部用高精度"，而是只把 router 内部的计算临时转成 float32，算完再转回 bfloat16。因为转回之后才进入跨设备的通信，所以通信量没有增加。
论文的实验结果（32 专家的 Switch-Base，训练早期）：
| 精度方案 | 质量（Neg. Log Perp.，越高越好） | 速度（examples/sec） |
|---|---|---|
| 全程 float32 | -1.718 | 1160 |
| 全程 bfloat16 | -3.780（发散） | 1390 |
| Selective precision | -1.716 | 1390 |
质量和 float32 几乎一样，速度和 bfloat16 一样。
为什么偏偏是 router 对精度敏感？一个直观的理解
这一段是我为了帮助理解写的直观解释。论文原文只说"bfloat16 会加剧 router 中 softmax 计算的问题"，没有展开成下面这个例子。
bfloat16 只保留大约 3 位十进制有效数字。在 4 到 8 这个区间里，它能表示的相邻两个数之间的间隔是 0.03125。
假设"猫"这个 token 的 router 打分（logits）里，E0 是 4.01，E1 是 4.00。float32 下 E0 明显领先，bfloat16 下：
4.01 → 4.0
4.00 → 4.0
两者变成了同一个数，这个 token 该去哪个专家，变成了平局。而 router 的输出是硬选择：差一点点，token 就被送去了完全不同的专家，后面的计算全变了。一个普通的线性层，数值误差只会带来输出的小幅偏移，而 router 的 argmax 会把小误差放大成"去哪个专家"的跳变。
所以 router 是整个网络里对数值精度最敏感的位置之一，值得单独用高精度算。
2.2 更小的初始化
问题：专家多、层数深，初始阶段的随机性更容易把训练带偏。
做法：权重从截断正态分布中采样，标准差 σ = √(s / n)，其中 n 是输入维度，s 是缩放系数。常规 Transformer 默认 s = 1.0，论文建议缩小 10 倍，取 s = 0.1。
用我们猫吃鱼的设定（d_model = 4，即 n = 4）手算：
s = 1.0：σ = √(1.0 / 4) = 0.5
s = 0.1：σ = √(0.1 / 4) ≈ 0.158
论文的实验结果（32 专家，训练 3.5k 步，3 个随机种子）：
| 初始化 | 平均质量 | 不同种子间的标准差 |
|---|---|---|
| s = 0.1 | -2.72 | 0.01 |
| s = 1.0 | -3.60 | 0.68 |
不仅平均质量好得多，不同随机种子之间的波动也小了很多。后者对稳定性尤其重要：它说明 s = 1.0 时，训练结果很大程度上靠运气。
一个直观的理解（同样是我的解释，不是论文原话）
如果 router 的输入是经过 LayerNorm 的（每个元素量级约为 1），那么 router 打分的标准差大约是 √s，与 n 无关。
假设 3 个专家的打分在两种初始化下大致是这样（取"典型值"）：
| 打分 | softmax 后的概率 | |
|---|---|---|
| s = 1.0 | [1.0, 0, -1.0] | [0.665, 0.245, 0.090] |
| s = 0.1 | [0.316, 0, -0.316] | [0.442, 0.323, 0.235] |
s = 1.0 时，router 在还没学到任何东西的时候就已经很"有主见"，大多数 token 一开始就被推向少数专家。s = 0.1 时，初始路由更柔和，给后面的训练留出了调整空间。
2.3 Expert Dropout：微调时单独加强专家层的正则
问题：Switch 模型的参数量远超 FLOP 相同的 dense 模型。预训练之后，要在很小的下游数据集上微调，参数多更容易过拟合。
做法：微调时，普通层的 dropout 保持 0.1，只把专家内部 FFN 的 dropout 提到 0.4。
论文的结果：把所有层的 dropout 一起调高反而变差，只调高专家层才有提升。
| 模型 | GLUE | SuperGLUE |
|---|---|---|
| T5-Base（dropout 0.1） | 82.9 | 72.4 |
| Switch-Base（dropout 0.1） | 84.7 | 73.0 |
| Switch-Base（所有层 0.3） | 83.9 | 70.7 |
| Switch-Base（普通层 0.1，专家层 0.4） | 85.2 | 73.0 |
这个思路很朴素：过拟合风险主要来自"多出来的那部分参数"，所以正则就加在那部分上。
小结：三个技巧的共同点
| 技巧 | 针对的问题 |
|---|---|
| Selective precision | router 的硬决策对数值精度敏感 |
| 更小的初始化 | 初始路由过于"有主见"，训练结果依赖运气 |
| Expert dropout | 参数多、微调数据少，容易过拟合 |
它们都不改变模型结构，是围绕"怎么把这个结构训好"的经验性处理。如果你以后想复现，这些是要记得的坑。
三、论文用实验证明了什么
3.1 参数和计算量解耦：更多专家，质量持续提升
这是整篇论文的核心假设：除了计算量，参数量本身是一个独立的、值得扩展的维度。
实验设计：保持每个 token 的 FLOPs 不变，把专家数从 2 开始不断翻倍，一直到 256。对应的参数量从 T5-Base 的 223M 一路增长到 14.7B。
结果：
- 质量随专家数增加而持续提升
- 但提升幅度逐渐变小（论文自己也提到，无限增加专家数会遇到收益递减）
另外一个值得说的点：论文强调，即使只有 2 个专家，也观察到了改进。这意味着这个思路并非只适用于超算规模。
3.2 top-1 比 top-2 更划算
论文把 Switch（top-1）和 MoE（top-2）做了正面对比。所有模型都是 128 个专家，每隔一层 FFN 换成专家层，在 32 个 TPUv3 核上训练，用同样的步数。
| 模型 | 容量因子 | 100k 步后质量 | 达到质量阈值所需时间（小时） | 速度（examples/sec） |
|---|---|---|---|---|
| T5-Base | — | -1.731 | 未达到 | 1600 |
| T5-Large | — | -1.550 | 131.1 | 470 |
| MoE-Base（top-2） | 1.25 | -1.559 | 80.7 | 790 |
| Switch-Base（top-1） | 1.25 | -1.553 | 65.0 | 910 |
| MoE-Base（top-2） | 1.0 | -1.572 | 80.1 | 860 |
| Switch-Base（top-1） | 1.0 | -1.561 | 62.8 | 1000 |
（质量阈值是论文任意选定的 -1.50。）
三个观察：
- 在容量因子 1.25 和 1.0 下，Switch 的质量不比 MoE 差，而且更快到达同一质量
- Switch 在较小的容量因子（1.0、1.25）下表现更好。这对大模型很重要，因为大模型里显存紧张，容量因子要尽量小
- 论文还指出，这里的速度同时取决于算法和底层实现，不完全是算法本身的差异
这也回答了 EP1 里的疑问：早期 MoE 认为必须路由到至少 2 个专家，才能给 router 提供有用的梯度。Switch 的实验说明，top-1 就够了（EP1 讲过：gate 值乘在输出上，router 照样有梯度）。
3.3 预训练速度：同样的计算量，快多少
论文给了两个口径：
| 口径 | 结果 |
|---|---|
| 按训练步数（样本效率） | Switch-Base（64 专家）达到 T5-Base 的同样质量，所需步数约为 1/7.5 |
| 按墙钟时间 | Switch-Base（64 专家）用时约为 T5-Base 的 1/7 |
两个口径要分开看，因为 Switch 多了路由计算和跨设备通信，步数上的加速不一定能等价换算成时间上的加速。论文两个口径都做了，结论方向一致，这才是可信的地方。
还有一个对比：假如不用 Switch，而是把同样的资源花在一个更大的 dense 模型上（T5-Large，每个 token 的 FLOPs 是 Switch-Base 的 3.5 倍）。结果是 Switch-Base 仍然更样本高效，按时间算约有 2.5 倍的加速。
3.4 下游任务：涨了，但不是全涨
论文把 FLOP 匹配的 Switch 和 T5 在多个任务上微调，对比如下（部分任务）：
| 任务 | T5-Base | Switch-Base | T5-Large | Switch-Large |
|---|---|---|---|---|
| SuperGLUE | 75.1 | 79.5 | 82.7 | 84.7 |
| Winogrande | 66.6 | 73.3 | 79.1 | 83.0 |
| 闭卷问答 Trivia QA | 24.5 | 30.7 | 29.5 | 36.9 |
| ARC Challenge | 35.5 | 32.8 | 35.5 | 35.5 |
| ARC Easy | 56.7 | 61.3 | 68.8 | 66.0 |
大部分任务 Switch 更好，其中闭卷问答的提升尤其明显。但表里加粗的两处，Switch 反而更差或没有提升：
- ARC Challenge：Switch-Base 比 T5-Base 低
- ARC Easy：Switch-Large 比 T5-Large 低
论文自己也承认，只有 ARC 这两个数据集没有观察到收益。
关于闭卷问答提升明显，一个常见的解读是"更多参数能存更多知识"，这和 Switch 的设计思路是吻合的。但要说明，这是一种合理的解释，这篇论文并没有专门设计实验去证明它。
3.5 蒸馏：把大稀疏模型压回小 dense 模型
稀疏大模型部署不方便。论文研究了把它蒸馏回 FLOP 相同的 dense 模型（T5-Base，223M 参数）。
两个有效的做法：
- 用教师模型的非专家层权重来初始化学生（因为两者 FLOP 匹配，非专家层的维度一样，可以直接复制）
- 训练损失里，0.75 用真实标签，0.25 用教师的输出概率混合
结果（3.8B 参数的 Switch-Base 作为教师）：
| 方案 | 质量 | 保留的质量增益 |
|---|---|---|
| T5-Base（学生基线） | -1.636 | — |
| Switch-Base（教师） | -1.444 | 100% |
| 直接蒸馏 | -1.631 | 3% |
| + 非专家权重初始化 | -1.598 | 20% |
| + 0.75 混合损失 | -1.580 | 29% |
也就是用大约 1/20 的参数，保留了约 30% 的质量增益。进一步压缩到 99%，仍能保留约 28%。
注意措辞：是**"质量增益的 30%"**，不是"质量的 30%"。学生比教师差，只是把教师相对于 T5-Base 多出来的那部分提升，拿回来了三成。
3.6 多语言
在 101 种语言的预训练中，FLOP 匹配的 Switch 在所有 101 种语言上都优于 mT5-Base。平均步数加速约 5 倍，91% 的语言达到 4 倍及以上。
3.7 万亿参数
论文最后展示了两个大模型：
| 模型 | 参数量 | 专家数 | 每序列 FLOPs | 备注 |
|---|---|---|---|---|
| T5-XXL（基线） | 11B | — | 6.3T | — |
| Switch-XXL | 395B | 64 | 6.3T | 与 T5-XXL FLOP 匹配 |
| Switch-C | 1.6T | 2048 | 890B | 专家多、单次计算量小 |
关键结果：
- Switch-C 达到同一困惑度的速度，约为 T5-XXL 的 4 倍（相同计算预算下）
- 论文报告 Switch-C 训练中没有出现不稳定
- 训练 250k 步后，两个 Switch 变体的质量都比 T5-XXL 高出 0.061 以上；500k 步时，Switch-XXL 领先 0.087
摘要里那句"4x speedup over T5-XXL"，指的是 Switch-C 这个对比。
四、它没有证明什么：局限和边界
一篇论文最容易被误读的，往往是它没说的部分。
1. "同样计算量"不等于"同样成本"。 Switch 的 FLOPs 和对应的 dense 模型一样，但所有专家的参数都要放进内存。参数量大 100 倍，显存和存储也要大那么多。蒸馏是一种缓解思路，但论文自己也说明，被蒸馏的学生只能保留一部分收益。
2. 所有"快几倍"，指的都是预训练速度。 论文的核心结论是"达到同样的预训练质量，需要的步数或时间更少"。它不是在说推理更快，也不是在说部署成本更低。
3. 速度结论受实现影响。 论文自己在对比 Switch 和 MoE 的速度时就注明了：速度既取决于算法，也取决于底层实现。实验是在 TPU、Mesh-TensorFlow 这套环境下做的，换成别的硬件和框架，具体倍数会变。
4. 下游任务不是全面领先。 前面表里 ARC 两个数据集就是例子。另外，据媒体对论文的报道，参数最多的 Switch-C 在 SQuAD 上的成绩还不如 Switch-XXL。论文把这归因于"微调质量、计算量和参数量之间的关系并不清晰"。简单说就是：参数越多，不保证每个下游任务都越好。
5. 稳定性问题并没有被彻底解决。 Switch-C 没有出现不稳定，但这是在"很多专家、每个专家算得很少"的配置下。后续工作 ST-MoE 回顾时指出，Switch 在高 FLOPs 的大稀疏模型上，在他们的设置里仍然不稳定，并围绕稳定性做了新的设计。所以"三个技巧"是缓解，不是根治。
6. 它基于 T5，是 encoder-decoder。 我们这个系列实现的是 decoder-only 的 GPT。MoE 层替换的是 Block 里的 FFN，这在两种结构里是通用的，但论文里的具体数字，不能直接套到 decoder-only 模型上。
五、后续影响（只点到为止）
Switch Transformer 的价值，更多在于证明了一个方向可行、并把实现简化了，而不是定下了最终形态：
- top-1 并没有成为终局。 后来的主流开源 MoE 并非都沿用 top-1，比如 Mixtral 8x7B 是每个 token 选 2 个专家
- 负载均衡的做法还在演化。 DeepSeek 系列引入了共享专家和更细粒度的专家划分，并在后续版本里尝试不依赖辅助损失的均衡方式
- 但"参数量与每个 token 的计算量解耦"这个核心思路，是这一整条路线的共同起点
这些内容我们不在本系列展开，留给以后单独讲。
六、如果要在我们的 GPT 里用 Switch，怎么改
这是一个很小的改动，EP2 的代码已经足够：
- 保持 Pre-LN 的 Block 结构不变
- 把每隔一个 Block 的 FFN 换成
SwitchFFN（论文的主要实验就是每隔一层一个专家层） - 总损失 = 交叉熵 + 所有 Switch 层的辅助损失之和，α 取 0.01
- 初始化用较小的 s（0.1），router 的计算用 float32
但别指望在小模型、小数据上看到论文里的结果。论文的收益来自大规模预训练，我们用随机初始化的小模型做玩具实验，只能验证机制是否正确，不能复现"快 7 倍"这样的结论。
七、一句话讲清楚这篇论文
把 Transformer 每隔一层的 FFN，换成 N 个并行的专家，每个 token 只走其中 1 个。这样参数量扩大 N 倍，而每个 token 的计算量几乎不变。
靠容量上限和负载均衡损失保证专家不被饿死、也不被压垮；靠float32 router、更小的初始化、专家 dropout 让它能在低精度下稳定训练。
实验表明，在同等计算预算下，预训练达到同样质量，比 dense 的 T5 快 4 到 7 倍。代价是显存和部署更重，并且下游收益并不总是均匀的。
系列小结
| 篇 | 主题 | 核心问题 |
|---|---|---|
| EP01 | MoE 与单专家路由 | 如何增加参数量，同时保持每个 token 的计算量基本不变？ |
| EP02 | 专家容量与负载均衡 | 路由不均匀、专家容量不足怎么办？ |
| EP03 | 训练稳定性、实验结论与局限 | 如何稳定训练，实验验证了什么，又有哪些局限？ |
Switch Transformer 这一篇论文，到这里就解读完了。
接下来我们回到"零件替换"的主线：从现代大模型相对于 GPT-2 替换掉的组件里，挑一个开始。
参考：Fedus, Zoph, Shazeer. Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity. JMLR 2022 (arXiv:2101.03961)。文中数据均来自该论文；标注为"我的解释"的部分是帮助理解的直观说明，不是论文原话。