In deep learning, models typically reuse the same parameters for all inputs. Mixture of Experts (MoE) defies this and instead selects different parameters for each incoming example. The result is a sparsely-activated model -- with outrageous numbers 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 Switch Transformer. We simplify the MoE routing algorithm and design intuitive improved models with reduced communication and computational costs. Our proposed training techniques help wrangle 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.
核心贡献 · Key contributions
将 MoE 路由简化为 top-1 专家，降低计算和通信成本。 Simplifies MoE routing to top-1 expert, reducing computation and communication costs.
在相同每 token FLOPs 下，预训练速度相比 T5 提升高达 7 倍。 Achieves up to 7x pre-training speedup over T5 with same FLOPs per token.
引入选择性精度（bfloat16）训练大型稀疏模型。 Introduces selective precision (bfloat16) training for large sparse models.
展示扩展到万亿参数，相比 T5-XXL 实现 4 倍加速。 Demonstrates scaling to trillion parameters with 4x speedup over T5-XXL.
在多语言设置中，所有 101 种语言均获得改进。 Shows universal improvements across 101 languages in multilingual setting.
提出专家 dropout 和蒸馏，以 99% 压缩率保留 30% 的质量提升。 Proposes expert dropout and distillation to preserve 30% quality gain with 99% compression.
局限 · Limitations
最大模型（如 Switch-XXL）仍存在训练不稳定性。 Training instability persists for largest models like Switch-XXL.
下游任务收益未能完全匹配上游困惑度改进。 Downstream gains do not fully match upstream perplexity improvements.
微调质量与每 token FLOPs 和参数数量的关系尚不明确。 Fine-tuning quality depends on FLOPs per token and parameter count in unclear ways.
专家容量因子在 token 丢弃和效率之间存在权衡。 Expert capacity factor introduces trade-off between token dropping and efficiency.
蒸馏仅保留约 30% 的稀疏模型质量提升。 Distillation preserves only ~30% of sparse model quality gain.
论文章节 · Sections（共 32）
摘要Abstract
1 引言1 Introduction
2 Switch Transformer2 Switch Transformer
2.1 简化稀疏路由2.1 Simplifying Sparse Routing
2.2 高效稀疏路由2.2 Efficient Sparse Routing
2.3 综合运用：Switch Transformer2.3 Putting It All Together: The Switch Transformer
2.4 改进的训练与微调技术2.4 Improved Training and Fine-Tuning Techniques
3 Scaling 属性3 Scaling Properties
3.1 基于步骤的扩展结果3.1 Scaling Results on a Step-Basis
3.2 基于时间的缩放结果3.2 Scaling Results on a Time-Basis
3.3 Scaling（规模扩张）与更大的稠密模型3.3 Scaling Versus a Larger Dense Model
4 下游结果4 Downstream Results
4.1 微调4.1 Fine-Tuning
4.2 蒸馏4.2 Distillation
4.3 多语言学习4.3 Multilingual Learning
5 结合数据、模型和专家并行设计模型5 Designing Models with Data, Model, and Expert-Parallelism
5.1 数据并行5.1 Data Parallelism
5.2 模型并行5.2 Model Parallelism
5.3 模型并行与数据并行5.3 Model and Data Parallelism
5.4 专家并行与数据并行5.4 Expert and Data Parallelism
5.5 专家、模型和数据并行5.5 Expert, Model and Data Parallelism
5.6 迈向万亿参数模型5.6 Towards Trillion Parameter Models
6 相关工作6 Related Work
7 讨论7 Discussion
8 未来工作8 Future Work
9 结论9 Conclusion
附录 A 注意力开关Appendix A Switch for Attention
附录 B 使用无令牌遗留防止令牌丢失Appendix B Preventing Token Dropping with No-Token-Left-Behind
附录 C 鼓励跨专家探索Appendix C Encouraging Exploration Across Experts
附录 D 低计算量下的开关变换器Appendix D Switch Transformers in Lower Compute Regimes
附录 E 上游与下游模型性能的关系Appendix E Relation of Upstream to Downstream Model Performance
附录 F 开关变换器的伪代码Appendix F Pseudo Code for Switch Transformers