精读 06 | Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity

论文基本信息

阅读地图（先看这里）

本文篇幅很长（约 40 页），但核心思想集中在前 25 页。建议按如下顺序理解：

Abstract（摘要）
    ↓
Introduction（动机：为什么需要稀疏模型？）
    ↓
Section 2.1：MoE 基础机制——路由器与专家
    ↓
Section 2.2：Switch Routing——简化到只选 1 个专家
    ↓
Section 2.3：负载均衡损失——防止所有 token 挤向同一专家
    ↓
Section 2.4：容量因子与丢弃 token 的处理
    ↓
Section 2.5：训练技巧——解决训练不稳定
    ↓
Section 3：规模化实验——和 T5 对比速度与质量
    ↓
Section 4：下游任务表现

核心一句话：Switch Transformer 把"混合专家（MoE）"机制简化到极致——每次只选 1 个专家，结果训练速度提升 7 倍，还能扩展到万亿参数。

术语速查表（首次阅读必看）

在正式精读之前，先把核心术语搞清楚：

一、Abstract（摘要）精读

原文

"In deep learning, models typically reuse the same parameters for all inputs. Mixture of Experts (MoE) models defy this and instead select different parameters for each incoming example. The result is a sparsely-activated model—with an outrageous number of parameters—but a constant computational cost. However, despite several notable successes of MoE, widespread adoption has been hindered by complexity, communication costs, and training instability. We address these with the introduction of the Switch Transformer. We simplify the MoE routing algorithm and design intuitive improved models with reduced communication and computational costs. Our proposed training techniques mitigate the instabilities, and we show large sparse models may be trained, for the first time, with lower precision (bfloat16) formats. We design models based off T5-Base and T5-Large to obtain up to 7x increases in pre-training speed with the same computational resources. These improvements extend into multilingual settings where we measure gains over the mT5-Base version across all 101 languages. Finally, we advance the current scale of language models by pre-training up to trillion parameter models on the 'Colossal Clean Crawled Corpus', and achieve a 4x speedup over the T5-XXL model."

翻译

在深度学习中，模型通常对所有输入复用相同的参数。混合专家（MoE）模型打破了这一惯例——它为每个输入样本选择不同的参数子集。这样得到的是一个稀疏激活的模型：参数量可以极其庞大，但每次推理的计算量却保持不变。然而，尽管 MoE 取得了若干显著成果，其大规模落地一直受到三大障碍的阻碍：实现复杂、通信开销高、训练不稳定。我们通过引入 Switch Transformer 来解决上述问题。我们简化了 MoE 路由算法，设计出直觉上更合理、通信与计算开销更低的改进模型。我们提出的训练技巧缓解了不稳定性，并首次证明大型稀疏模型可以用低精度（bfloat16）格式进行训练。以 T5-Base 和 T5-Large 为基础，我们在相同计算资源下实现了高达 7 倍的预训练加速。这些改进同样适用于多语言场景，在全部 101 种语言上均超越了 mT5-Base。最后，我们在"超大规模爬取语料"上预训练了万亿参数级别的模型，相比 T5-XXL 实现了 4 倍加速。

新手讲解

"稀疏激活"是什么意思？

想象你去医院看病。传统的大型神经网络就像每次看病都把全院所有医生（神经元）叫来会诊——不管你是感冒还是骨折，所有人都要参与，既费时又费资源。

MoE 模型则像一个聪明的分诊台：你进门时，分诊护士（路由器/Router）先判断你的病情，然后只把你送给对口的一位或几位专科医生（专家网络/Expert）。其他专科医生继续休息，不消耗任何资源。

关键矛盾：你可以雇很多专科医生（参数量很大），但每次看诊只动用其中一两个（计算量不变）。这就是"参数量巨大，计算量恒定"的核心逻辑。

MoE 之前为什么没普及？

• 路由算法复杂，难以实现
• 多个 GPU 之间需要频繁传递数据（通信成本高）
• 训练过程容易崩溃（训练不稳定）

Switch Transformer 做了什么？

把路由决策从"选 Top-K 个专家"极简化为"只选 1 个专家"，同时配套一系列训练稳定技巧，结果不仅更简单，速度还快了 7 倍。

二、Introduction（引言）精读

核心段落 1：规模化是方向，但"稠密扩展"代价太高

原文

"Large scale training has been an effective path towards flexible and powerful neural language models."

翻译

大规模训练已经被证明是通向灵活、强大的神经语言模型的有效路径。

原文

"We propose here a sparsely-activated expert model: the Switch Transformer, where sparsity comes from activating a subset of the neural network weights for each incoming example."

翻译

我们提出了一种稀疏激活的专家模型：Switch Transformer。其稀疏性体现在：对于每个输入样本，只激活神经网络权重中的一个子集。

新手讲解

GPT-3 那样的"密集"模型（Dense Model）：每处理一个词，所有 1750 亿参数都参与计算，代价极高。

Switch Transformer 的路线：参数可以有 1 万亿，但每处理一个词，只激活其中极小一部分（比如 1/64）。这样，参数容量大幅提升，计算量却几乎不变。

核心段落 2：MoE 的历史与障碍

原文

"Machine learning libraries and hardware accelerators still cater to dense matrix multiplications."

翻译

机器学习库和硬件加速器（如 GPU、TPU）目前仍主要针对稠密矩阵乘法进行优化。

原文

"We start with the Mixture-of-Expert (MoE) paradigm and simplify it to yield training stability and computational benefits. While MoE models had notable successes in machine translation, widespread adoption was hindered by complexity, communication costs, and training instabilities."

翻译

我们以混合专家（MoE）范式为出发点，对其加以简化，从而获得训练稳定性和计算效率上的双重收益。尽管 MoE 模型在机器翻译任务上取得了显著成功，但其大规模推广仍受制于实现复杂度、通信成本以及训练不稳定性。

新手讲解

MoE 并非 Switch Transformer 的发明——它的历史可以追溯到 1990 年代，近年 Google 的 GShard 等工作也在用。但它一直没有成为主流，原因有三：

1. 实现复杂：路由逻辑很难写正确
2. 通信开销大：不同专家在不同 GPU 上，token 需要跨 GPU 传输
3. 训练不稳定：模型经常发散（loss 爆炸或变成 NaN）

Switch Transformer 的贡献就是把这三个问题全部解决，或至少显著缓解。

核心段落 3：本文的核心贡献清单

原文（意译提炼）

论文的主要贡献包括：

1. Switch Transformer 架构——简化并改进了 MoE，便于理解、易于实现
2. 在相同 FLOPs 约束下实现 7 倍以上预训练加速（对比强调调优的 T5）
3. 首次成功用 bfloat16 格式训练大型稀疏模型，解决训练不稳定问题
4. 稀疏模型到稠密小模型的蒸馏——可保留约 30% 的稀疏模型质量增益
5. 在全部 101 种语言上超越 mT5-Base
6. 成功预训练 1.6 万亿参数模型，比 T5-XXL 快 4 倍

新手讲解

这份贡献清单可以理解为："我们不只是把 MoE 做得更快，我们还把它做得更稳定、更可用、更通用。" 每一条贡献都对应论文中的一个实验或技术章节。

三、Section 2：Switch Transformer 模型（核心章节，逐段精读）

这是整篇论文最重要的部分，务必逐段理解。

3.1 混合专家（MoE）基础机制

背景：什么是 MoE 路由？

原文

"The router variable Wr produces logits h(x) = Wr · x which are normalized via a softmax distribution over the N experts at that layer."

翻译

路由器变量 $W_r$ 生成 logits（未归一化的分数）：$h(x) = W_r \cdot x$，然后通过 softmax 归一化为 N 个专家上的概率分布。

对于第 i 个专家，其"门控值"（gate value）为：

$$p_i(x) = \frac{e^{h(x)i}}{\sum{j=1}^{N} e^{h(x)_j}}$$

最终输出是所有被选中专家的加权组合：

$$y = \sum_{i \in \mathcal{T}} p_i(x) \cdot E_i(x)$$

其中 $\mathcal{T}$ 是被选中的 Top-K 个专家的下标集合，$E_i(x)$ 是第 i 个专家对输入 x 的输出。

翻译补充

路由器是一个小型线性层，把 token 的特征向量 x 映射为 N 个专家的"得分"，再用 softmax 把得分转成概率。传统 MoE 会选分数最高的 K 个专家（通常 K=2），然后把这 K 个专家的输出按概率加权求和。

新手讲解（医院类比展开）

把整个过程用医院比喻讲清楚：

病人（token）进门
    ↓
分诊护士（Router 路由器）
    扫描病人信息（token特征向量 x）
    计算"适合去哪个科室"的分数（logits）
    归一化成概率（softmax）：比如骨科 60%，内科 30%，眼科 10%
    ↓
传统 MoE（Top-2）：
    送去骨科（得分最高）+ 送去内科（得分第二）
    最终结果 = 骨科诊断×0.6 + 内科诊断×0.3（加权融合）
    ↓
Switch Transformer（Top-1，本文方法）：
    只送去骨科（得分最高的那一个）
    最终结果 = 骨科诊断（单一专家）

为什么之前大家觉得 K=1 不行？

"Shazeer et al. conjectured that routing to k > 1 experts was necessary in order to have non-trivial gradients to the routing functions. The authors intuited that learning to route would not work without the ability to compare at least two experts."

翻译

Shazeer 等人猜测，路由到 K > 1 个专家是必要的，这样才能为路由函数提供非平凡的梯度。他们的直觉是：如果不能同时比较至少两个专家，路由机制就无法通过梯度学习。

新手讲解

这是个很重要的历史背景：之前的研究者认为，如果每次只选 1 个专家，路由器在反向传播时拿不到"其他专家表现如何"的信号，就没法学习"应该把 token 发给谁"。

但 Switch Transformer 的实验发现，这个担忧是多余的——K=1 完全可行，而且更简单更快。

3.2 Switch Routing：只选 1 个专家（论文核心创新）

原文

"We simplify the MoE routing algorithm by instead using a simplified strategy where we route to only a single expert. We show this simplification preserves model quality, reduces routing computation and performs better."

翻译

我们通过只将 token 路由到单个专家来简化 MoE 路由算法。我们证明这一简化方案不仅保持了模型质量，还降低了路由计算量，并在性能上表现更好。

原文（三大优势）

"The benefits for the Switch layer are three-fold:
(1) The router computation is reduced as we are only routing a token to a single expert.
(2) The batch size (expert capacity) of each expert can be at least halved since each token is only being routed to a single expert.
(3) The routing implementation is simplified and communication costs are reduced."

翻译

Switch 层带来的三重好处：
1. 路由计算量降低：每个 token 只需路由到一个专家，计算量减半。
2. 每个专家的批次大小（容量）至少可以减半：因为每个 token 只去一个专家，专家的负担降低。
3. 路由实现更简单，通信成本降低：逻辑更清晰，跨设备通信减少。

新手讲解

这就是论文标题中"Simple"（简单）的体现：

简单类比：去医院只看一个专科医生，比同时看两个医生、再把两份诊断书加权平均，效率高得多，结论也往往更清晰。

Switch 路由的数学表达

路由仍然使用 softmax 计算概率（和传统 MoE 完全一样），区别只是最终只保留概率最高的那 1 个专家：

$$\text{选中专家} = \arg\max_i p_i(x)$$

$$y = p_{\text{选中}}(x) \cdot E_{\text{选中}}(x)$$

注意：路由概率 $p_i(x)$ 仍然是可微的（softmax 输出），所以路由器可以通过反向传播正常学习，这就消除了"K=1 无法学习"的顾虑。

3.3 负载均衡损失（Load Balancing Loss）

这是 Switch Transformer 必须解决的核心工程问题。

问题：为什么需要负载均衡？

如果不加任何约束，路由器会自然地把大量 token 都发给"最受欢迎"的 1-2 个专家，而其他专家几乎没有工作可做。这个现象叫做"路由崩溃"（routing collapse）。

用医院类比：如果医院里有一位特别出名的专家，所有病人都想挂他的号，结果他的诊室门口排了几千人，而其他专家的诊室空空如也——这对整个医院的效率是灾难性的。

解决方案：辅助损失函数

原文

"To encourage a balanced load across experts we add an auxiliary loss. Switch Transformers simplifies the original design which had separate load-balancing and importance-weighting losses."

翻译

为了鼓励各专家之间的负载均衡，我们增加了一个辅助损失。Switch Transformer 简化了原始设计（原设计有两个独立的损失：负载均衡损失和重要性加权损失），合并为统一形式。

负载均衡损失的完整公式

$$\mathcal{L}{\text{aux}} = \alpha \cdot N \cdot \sum{i=1}^{N} f_i \cdot P_i$$

各符号含义：

$f_i$（专家 i 实际收到的 token 比例）：

$$f_i = \frac{1}{T} \sum_{x \in \mathcal{B}} \mathbf{1}{\arg\max_j p_j(x) = i}$$

含义：在当前这批（batch）T 个 token 中，有多少比例被路由到了专家 i。这是一个离散的统计量（用指示函数 $\mathbf{1}{\cdot}$ 计算）。

$P_i$（路由器分配给专家 i 的概率总量）：

$$P_i = \frac{1}{T} \sum_{x \in \mathcal{B}} p_i(x)$$

含义：路由器对这批 token 计算出的"发给专家 i 的概率"的平均值。这是一个连续可微的量。

$\alpha$：超参数，论文设定为 $10^{-2}$（作者在 $10^{-1}$ 到 $10^{-5}$ 范围内做了超参搜索，发现 $10^{-2}$ 最佳）。

$N$：专家总数。

翻译

$\alpha$ 是一个可调整的系数乘子，$f_i$ 是当前批次中被路由到专家 i 的 token 分数，$P_i$ 是路由器为专家 i 分配的概率分数。

为什么这个损失有效？

关键洞察：当所有专家的负载均匀时，损失最小。

如果 N 个专家完全均匀地分担所有 token：
- 每个专家的 $f_i = 1/N$
- 每个专家的 $P_i = 1/N$
- 损失 $= \alpha \cdot N \cdot \sum_{i=1}^{N} \frac{1}{N} \cdot \frac{1}{N} = \alpha \cdot N \cdot N \cdot \frac{1}{N^2} = \alpha$（最小值）

如果某个专家"垄断"了所有 token：
- 某一个 $f_i \approx 1$，其余 $\approx 0$
- 损失会远大于均匀分布的情况
- 梯度会驱使路由器把 token 分散开

新手讲解（类比）

这个损失就像"医院管理层对每位医生下达的工作量考核"：
- 如果某位医生接诊了全院 80% 的病人，管理层会扣分（辅助损失增大）
- 梯度反向传播后，路由器会被"迫"把更多病人送给其他医生
- 最终趋向均衡：每位医生接诊约 1/N 的病人

技巧细节：$f_i$ 是离散的（不可微），而 $P_i$ 是连续的（可微）。两者相乘后，梯度通过 $P_i$ 流向路由器，这是一个精心设计的"曲线救国"——用可微项代替不可微项来传导梯度。

3.4 容量因子（Capacity Factor）与 Token 溢出处理

为什么需要"容量"的概念？

当 N 个 token 要被路由到 N 个专家时，理想情况是每个专家恰好收到相同数量的 token。但实际上路由是动态的，不可能完美均匀。

硬件处理要求所有专家同时并行处理固定大小的批次（因为 TPU/GPU 的矩阵运算需要固定形状的张量），所以必须预先分配每个专家能处理的 token 数量上限。

专家容量公式

$$\text{Expert Capacity} = \left(\frac{\text{tokens per batch}}{\text{number of experts}}\right) \times \text{capacity factor}$$

翻译

专家容量 = （批次中的 token 总数 / 专家数量）× 容量因子

示例：
- 批次中有 512 个 token
- 有 8 个专家
- 理想均分：每个专家处理 64 个 token
- 容量因子 = 1.5 时：每个专家最多处理 96 个 token（为不均匀情况留出缓冲）

原文

"A capacity factor greater than 1.0 creates additional buffer to accommodate for when tokens are not perfectly balanced across experts."

翻译

容量因子大于 1.0 会创建额外缓冲，以应对 token 在专家之间分配不完全均匀的情况。

当专家"超载"时发生什么？

原文

"If too many tokens are routed to an expert (referred to later as dropped tokens), computation is skipped and the token representation is passed directly to the next layer through the residual connection."

翻译

如果太多 token 被路由到某个专家（这些 token 被称为"丢弃的 token"），它们的计算会被跳过，其 token 表示直接通过残差连接传递到下一层。

新手讲解

这就像医院的诊室座位有限：
- 每个诊室预设最多接诊 X 位病人
- 如果涌入的病人超过 X，超出的病人不看诊，直接去下一个环节（残差连接）
- 这样不会导致系统崩溃，但这些 token 在该层没有被专家处理

容量因子的权衡

原文

"Increasing the expert capacity is not without drawbacks, however, since high values will result in wasted computation and memory."

翻译

然而，增大专家容量并非没有代价——容量过高会导致计算和内存的浪费。

论文建议：实际训练中保持 token 溢出率（drop rate）低于 1%，这样对质量的影响可以忽略不计。

3.5 训练稳定技巧（Section 2.4~2.5）

MoE 模型历史上一直有训练不稳定的问题（loss 突然爆炸，或路由崩溃）。Switch Transformer 用了三个关键技巧解决这些问题：

技巧 1：选择性精度转换（Selective Precision）

原文

"Sparse expert models introduce training difficulties over a vanilla Transformer. Instability can result because of the hard-switching (routing) decisions at each of these layers. Further, low precision formats like bfloat16 can exacerbate issues in the softmax computation for our router."

翻译

稀疏专家模型相比普通 Transformer 引入了训练难度。在路由层的"硬切换"（每次只选一个专家）决策会导致不稳定。此外，bfloat16 等低精度格式会加剧路由器 softmax 计算中的数值问题。

解决方案：只把路由器的输入提升到 float32 精度计算，其他所有地方保持 bfloat16。

类比：路由器是"精密仪器"，需要高精度。其他地方是"重型机械"，用低精度就够了。

技巧 2：更小的参数初始化（Smaller Initialization）

将默认初始化尺度从 1.0 降低到 0.1，显著提升了训练稳定性。

论文数据：负对数困惑度从 -3.60 提升到 -2.72，且不同随机种子间的方差大幅降低。

类比：新建的专家网络刚开始"说话声音要小一点"，这样不会一上来就主导整个系统，给整个网络的调整留出余地。

技巧 3：专家层 Dropout

微调（Fine-tuning）时，专家层的 dropout 率设为 0.4（而普通层只有 0.1）。

类比：专家医生能力很强，容易"过拟合"到特定病人群体。微调时让他们"随机休息"得更频繁，强迫整个系统不过度依赖少数几个专家。

四、Section 3：规模化实验结果精读

4.1 核心对比：稀疏 vs 稠密，同等 FLOPs

关键设计原则

论文中最重要的一个思想：

"Increase the parameter count while keeping the floating point operations (FLOPs) per example constant."

翻译

在保持每个样本浮点运算量（FLOPs）不变的前提下，增加参数量。

新手讲解

这是比较稀疏模型和稠密模型的正确方式：

• 不公平的比较：稀疏模型参数多 → 计算也多 → 当然更好
• 公平的比较：稀疏模型参数多，但每次只激活一小部分 → 总计算量和小模型相当 → 看谁学得更好

论文用"相同 FLOPs"作为约束，比较的是"在相同计算预算下，谁的质量更高"。

4.2 关键实验数字

速度对比（预训练阶段）

参数规模

原文

"The parameter count, independent of total computation performed, is a separately important axis on which to scale."

翻译

参数量——独立于总计算量——是一个重要的、独立的规模化维度。

新手讲解

这句话是整篇论文的哲学核心：参数量和计算量是两个可以独立调节的旋钮。

传统理解：更多参数 = 更多计算 = 更贵。

Switch Transformer 的发现：可以通过稀疏激活，让参数量暴增，而计算量保持不变。这就像图书馆里有一百万本书（参数量），但每次查阅只翻其中几本（计算量）。

4.3 样本效率（Sample Efficiency）

稀疏模型不仅训练更快，还更会学习：

"Larger models are also more sample efficient—learning more quickly for a fixed number of observed tokens."

翻译

更大的模型也具有更高的样本效率——在处理相同数量的 token 时，学习速度更快。

意义：在预训练数据有限的场景下，Switch Transformer 能从同样多的文本中学到更多。

五、Section 4：下游任务表现精读

5.1 微调（Fine-tuning）结果

经过大量稀疏预训练后，Switch Transformer 在具体任务上同样表现优异：

5.2 蒸馏（Distillation）

稀疏大模型很难部署（占用内存极大），论文提出将其"蒸馏"成小的稠密模型：

"Models can be compressed by 95%+ while preserving approximately 30% of the sparse teacher's quality gains."

翻译

模型可以被压缩 95% 以上，同时保留约 30% 的稀疏教师模型相对于基线的质量增益。

新手讲解

蒸馏就像让一位经验丰富的老专家（大型稀疏模型）去训练一位刚入职的年轻医生（小型稠密模型）。年轻医生学不到老专家的全部经验，但能学到其中 30% 的精华——这在资源受限的部署场景中非常有价值。

5.3 多语言表现

"Across 101 languages, Switch Transformers show 5x mean speedup, with 91% of languages achieving 4x+ improvements over mT5-Base."

翻译

在 101 种语言上，Switch Transformer 平均实现 5 倍加速，91% 的语言达到了相比 mT5-Base 4 倍以上的改进。

六、Switch Transformer 整体架构图解

输入 Token 序列
    ↓
[Transformer Encoder/Decoder 标准层]
    ↓
[Switch 层（替换部分 FFN 层）]
    ┌─────────────────────────────────┐
    │                                 │
    │  每个 Token                     │
    │      ↓                          │
    │  Router（路由器）                │
    │  softmax(W_r · x)               │
    │      ↓                          │
    │  选出概率最高的 1 个专家          │
    │      ↓                          │
    │  Expert_i（专家 FFN 网络）       │
    │      ↓                          │
    │  p_i(x) · E_i(x)               │
    │      +                          │
    │  残差连接（溢出 token 直接通过） │
    └─────────────────────────────────┘
    ↓
[下一个标准 Transformer 层]

稀疏性的体现：在一个 Switch 层中，有 N 个专家网络，但对于每个 token，只有 1 个专家被激活（运算），其余 N-1 个专家的参数不参与计算。

七、关键思想总结与联系

Switch Transformer 解决的三大问题

与后续工作的联系

Switch Transformer 的思想直接影响了后续大模型的设计：

• Mixtral 8x7B（Mistral AI, 2023）：使用 Top-2 MoE 路由，延续了 Switch 的路线
• GPT-4（OpenAI, 2023，据推测）：疑似采用了 MoE 架构
• Gemini 1.5（Google, 2024）：MoE 架构，可处理 100 万 token 上下文
• DeepSeek-V2/V3（DeepSeek, 2024）：使用精细化的 MoE 路由机制

Switch Transformer 是 2021 年之后"MoE 复兴"的直接推动者。

八、一句话理解 Switch Transformer

"把神经网络里最费参数的前馈层（FFN）拆成 N 个'专科医生'（专家），每次只让一个专家处理当前的词，这样参数可以增加 N 倍，但计算量完全不变，结果训练速度快了 7 倍，还能做到万亿参数级别。"

九、值得深思的问题（延伸学习）

1. 路由学到了什么？ 论文没有深入分析路由器学到的模式，但后续研究发现不同专家确实专注于不同语义类型（实体名词、动词、代词等）。

2. K=1 真的够吗？ 论文证明在预训练中够用，但后续研究（如 Mixtral）发现 K=2 在推理质量上有时更好。这是一个仍在探索的开放问题。

3. 通信瓶颈如何解决？ 当专家分布在不同 TPU 核心时，token 需要跨设备传输。Switch Transformer 通过减少到 K=1 缓解了这个问题，但在更大规模时仍是挑战。

4. 稀疏模型的"遗忘"问题？ 如果某些专家长期不被激活，它们可能会"退化"。负载均衡损失正是为了防止这种情况。

精读完成。涵盖章节：Abstract、Introduction（全部核心段落）、Section 2（MoE 基础、Switch Routing、负载均衡损失、容量因子、训练技巧，共 5 个小节全部逐段精读）、Section 3（规模化实验核心段落）、Section 4（下游任务结果）。

字数约：9,000 字。