switch transformer

Switch Transformer is a sparse Mixture of Experts (MoE) model architecture introduced by Fedus et al. (2022) at Google that simplifies MoE routing by sending each token to exactly one expert (top-1 routing), demonstrating that this simpler approach achieves better scaling properties than previous multi-expert routing strategies while being easier to implement and more computationally efficient. The key insight of Switch Transformer is that routing each token to a single expert (k=1) rather than multiple experts works better than expected — previous MoE work like the Sparsely-Gated MoE (Shazeer et al., 2017) used top-2 routing, but Switch Transformer showed that simpler top-1 routing actually improves training stability and quality when combined with proper initialization and load-balancing. Architecture: Switch Transformer replaces the dense feedforward layers in a standard transformer with MoE layers, where each MoE layer contains multiple independent feedforward expert networks sharing the self-attention layer. A simple learned linear router computes expert scores for each token and routes it to the highest-scoring expert. Key innovations include: simplified routing (top-1 expert selection reduces computation and communication overhead), improved training stability through careful initialization (reducing expert output variance at initialization), auxiliary load-balancing loss (encouraging equal token distribution across experts — preventing expert collapse), selective precision (using FP32 for the router while using BFloat16 for experts — stabilizing routing decisions), and efficient expert parallelism (distributing experts across different devices with minimal cross-device communication). Switch Transformer demonstrated remarkable scaling: a Switch-C model with 1.6 trillion parameters (but only ~equivalent computation to a T5-Base model per token) achieved significant speedups over dense T5 models in pre-training. The paper showed that sparse MoE provides a "free lunch" — more parameters without proportional compute increase — validating the principle that parameter count and computational cost can be effectively decoupled.