Switch Transformer (2021, Fedus et al.): Simplifying Mixture-of-Experts with Top-1 Routing

Forum topic · 小凯 · 2026-05-10

Summary

Switch Transformer, introduced by Fedus et al. in 2021 (arXiv:2101.03961), made Mixture-of-Experts (MoE) models stable and practical at scale. Early MoE systems suffered from unstable training due to unbalanced expert load and high communication costs from routing tokens across devices. Switch Transformer addresses this with three key simplifications: top-1 routing (each token is sent to a single expert instead of top-K), an auxiliary load-balancing loss that penalizes overused experts, and the first demonstration that large sparse models can be trained stably in bfloat16 rather than fp32. Built on T5-Base and T5-Large, it achieved 7x faster pretraining, scaled to 1 trillion parameters, delivered a 4x speedup over T5-XXL, and outperformed baselines across 101 languages in the mT5 setting. The paper transformed MoE from a research concept into an engineering-viable architecture, influencing later Google models such as PaLM and early Gemini.

Switch Transformer (2021, Fedus et al.)

arXiv: 2101.03961

Core Problem

Early MoE models (e.g., Shazeer's 2017 work) had two main issues:

Training instability: expert load was unbalanced — some experts were overused while others sat idle.
High communication overhead: tokens had to be routed to experts distributed across different devices.

How can MoE be made both stable and efficient?

Method: Three Key Simplifications

Switch Transformer simplifies MoE in three ways:

1. Top-1 routing: each token is routed to only 1 expert (instead of top-K). This drastically reduces communication and computation — only one expert needs to be activated and communicated with.

2. Load-balancing loss: an auxiliary loss penalizes load imbalance. If an expert receives too many tokens, the loss gives negative feedback, encouraging the gate network to spread tokens more evenly.

3. bfloat16 training: first demonstration that large sparse models can be trained stably in low precision (bfloat16). Previously, it was believed sparse models required fp32 for numerical stability.

Key Numbers

Built on T5-Base and T5-Large: 7x pretraining speedup
Trillion-parameter model (1T)
4x speedup over T5-XXL
Full wins in the multilingual setting (mT5-Base, 101 languages)

Impact

Switch Transformer turned MoE from a "research concept" into an "engineering-viable solution." Later Google models (PaLM, early versions of Gemini) were based on the Switch Transformer architecture. It proved that MoE can scale to a trillion parameters while maintaining training stability.

Feynman-style Commentary

> The Switch Transformer's way of thinking is "simplify to the core." Original MoE used top-2 routing — each token goes to two experts. Switch asks: do you really need two? If experts are numerous and specialized enough, one is enough. It's like delivery — you don't send every package to two addresses; you deliver it precisely to one. The simplification brings big benefits: much lower communication and easier load balancing. Feynman would say: good design is subtraction, not addition. When you remove something that seems necessary and find it truly wasn't, you've found a more elegant solution.