[2021] Switch Transformer — Fedus et al.

24. Switch Transformer (2021, Fedus et al.)

arxiv: 2101.03961

核心问题：早期的 MoE（如 2017 Shazeer 的）有两个问题：训练不稳定（不同专家负载不均衡，有的专家被过度使用，有的闲置）和通信开销大（需要把 token 路由到分布在不同设备的专家）。怎么让 MoE 既稳定又高效？

方法创新： Switch Transformer 做了三个关键简化：

1. Top-1 路由：每次只路由到 1 个专家（而不是 top-K）。这 drastically 减少了通信和计算——只有一个专家需要被激活和通信。

2. 负载均衡损失：加一个辅助损失函数，惩罚负载不均衡。如果某个专家被分配了太多 token，损失函数会给负面反馈，鼓励 gate 网络把 token 分散开。

3. bfloat16 训练：首次证明大型稀疏模型可以用低精度（bfloat16）稳定训练。之前人们认为稀疏模型需要 fp32 来保持数值稳定。

关键数字：

• 基于 T5-Base 和 T5-Large，预训练速度提升 7x
• 万亿参数模型（1T）
• 比 T5-XXL 4x 加速
• 多语言设置（mT5-Base，101 种语言）全面超越

影响评估： Switch Transformer 让 MoE 从"研究概念"变成"工程可行方案"。Google 的后续模型（PaLM、Gemini 早期版本）都基于 Switch Transformer 的架构。它证明了 MoE 可以 scale 到万亿参数，同时保持训练稳定性。

费曼点评： > Switch Transformer 的思维方式是"简化到核心"。原始 MoE 用 top-2 路由——每个 token 去两个专家。Switch 问：真的需要两个吗？如果专家足够多、足够特化，一个就够了。这就像快递——不是每个包裹送两个地址，而是精确送到一个地址。简化带来的好处是通信大幅降低、负载均衡更容易。费曼会说：好的设计不是加法，是减法。当你去掉一个看似必需的东西后发现它确实不需要，你就找到了更优雅的解。

arxiv: 2101.03961

#论文深度研究 #小凯 #SwitchTransformer #MoE #负载均衡