← All papers

Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity

Fedus, Zoph, Shazeer (Google) · 2022 · JMLR 2022

ArchitectureRead on arXiv

Simplified Mixture-of-Experts by routing each token to a single expert (top-1), enabling trillion-parameter models that train faster than dense counterparts at the same compute budget.

Key Idea

Switch Transformers simplify the MoE routing mechanism by sending each token to exactly one expert (top-1 routing) instead of the top-2 used in prior work (Shazeer et al., 2017). This reduces communication costs and makes MoE practical at massive scale: the paper trained models up to 1.6 trillion parameters.

How It Works

  • Switch layer replaces the FFN in selected transformer layers with N expert FFNs + a learned router
  • Top-1 routing: router produces logits over experts, each token goes to the single highest-scoring expert
  • Capacity factor: each expert has a fixed buffer size = (tokens_in_batch / num_experts) × capacity_factor. Tokens exceeding capacity are dropped (passed through via residual)
  • Auxiliary load-balancing loss: encourages uniform routing across experts to prevent expert collapse
  • Selective precision: router computation in FP32 for stability, experts in BF16 for speed

Why It Matters

  • 7× speedup over dense T5-Base at equivalent compute, with the same or better quality
  • Proved that sparse models scale better than dense models, more parameters without proportional compute increase
  • The simplicity of top-1 routing (vs top-2) reduced implementation complexity and communication overhead
  • Directly inspired Mixtral, DeepSeek-MoE, and modern sparse architectures
  • Established key MoE training practices: load balancing loss, capacity factors, selective precision

Key Takeaways for Interviews

  • Switch Transformer = top-1 MoE routing, the paper that made MoE practical at scale
  • Top-1 vs top-2: top-1 is simpler and faster but potentially less robust; modern models (Mixtral) returned to top-2 for quality
  • Load balancing is critical, without the auxiliary loss, most tokens route to a few experts (expert collapse)
  • Capacity factor controls the tradeoff: higher = fewer dropped tokens but more memory; typical value is 1.0-1.25
  • In interviews: cite Switch Transformer as the foundational modern MoE paper, Mixtral/DeepSeek as the successors