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

Problem

Framing

Dense Transformer scaling raises both quality and compute cost together. The paper argues that parameter count can scale independently by replacing dense FFNs with single-expert routed sparse FFNs, yielding up to $7\times$ faster pre-training at matched FLOPs and scaling to trillion-parameter language models.

Currently Used Methods

Foundational

• @vaswaniAttentionAllNeed2017 — dense Transformer backbone with shared FFN parameters for every token.
  • Limitation in context: compute grows with parameter count, so scaling is expensive.
• @raffelT5_2020 — strongly tuned text-to-text Transformer baseline for large-scale pre-training.
  • Limitation in context: dense scaling needs far more FLOPs to gain quality.
• "Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer" — top-$k$ expert routing for sparse conditional computation.
  • Limitation in context: routing is complex, communication-heavy, and unstable at scale.
• "GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding" — large MoE Transformers for multilingual translation.
  • Limitation in context: top-2 routing still adds dispatch cost and training fragility.
• @brownGPT3_2020 — brute-force dense language-model scaling with strong in-context behavior.
  • Limitation in context: dense trillion-scale training is prohibitively expensive.

Proposed Method

Architecture

Switch Transformer replaces each dense Transformer FFN with a bank of $N$ FFN experts plus a router. Each token is sent to exactly one expert, then scaled by its router probability; self-attention and residual structure stay unchanged.

Loss / Objective

Training adds a load-balancing auxiliary term to the language-model loss.

Sampling Rule / Algorithm

Routing is top-1 expert selection with fixed per-expert capacity.

Training Procedure

• Pre-training data: C4 masked language modeling.
• Router precision: cast routing softmax computations to float32; keep the rest in bfloat16.
• Initialization: expert and dense weights use reduced $0.1\times$ standard initialization scale.
• Experts: many core experiments use 32 or 128 experts; largest models scale to trillion parameters.
• Capacity factor: swept as a routing hyperparameter; larger values reduce dropped tokens.

Evaluation

Datasets

• C4 pre-training corpus.
• GLUE.
• SQuAD.
• SuperGLUE.
• Winogrande.
• Multilingual mC4 across 101 languages.

Metrics

• Negative log perplexity for pre-training quality.
• Time to a fixed quality threshold.
• Examples per second.
• Task validation accuracy or score for downstream benchmarks.
• Step speedup over dense baselines.

Headline results

• C4 pre-training, matched FLOPs: Switch-Base reaches the target quality in $62.8$ hours vs $131.1$ for T5-Large.
• C4 pre-training, matched FLOPs: Switch-Base reports up to $7\times$ speedup over dense T5 variants.
• Selective precision: float32 routing + bfloat16 elsewhere gives $-1.716$ neg-log-perplexity at $1390$ examples/s; pure bfloat16 diverges to $-3.780$.
• Initialization ablation: $0.1\times$ init gives $-2.72 \pm 0.01$ after $3.5$k steps; standard init gives $-3.60 \pm 0.68$.
• Multilingual mT5 setting: mean $5\times$ step speedup across 101 languages; $91\%$ of languages exceed $4\times$.

Results table

Table 1: Benchmarking Switch versus MoE.

Ablations

• Routing choice: top-1 routing matches or beats top-2 MoE while simplifying dispatch.
• Router precision: selective float32 on routing stabilizes training with nearly pure-bfloat16 speed.
• Initialization scale: smaller initialization sharply improves early stability and variance.
• Fine-tuning dropout: lower non-expert dropout plus higher expert dropout performs best.

Method Strengths and Weaknesses

Strengths

• Top-1 routing cuts MoE complexity without hurting pre-training quality.
• Matched-FLOP comparisons show large wall-clock gains, not just larger parameter counts.
• Selective router precision enables stable bfloat16-era sparse training.
• Reduced initialization scale supports runs from 223M to trillion-parameter models.

Weaknesses

• Fixed expert capacity can drop overflowed tokens during routing.
• Load balancing needs an auxiliary loss and careful coefficient tuning.
• Sparse speedups depend on distributed all-to-all communication efficiency.
• Fine-tuning needs expert-specific regularization to avoid overfitting.

Suggestions from the authors

• Route sparse experts inside self-attention, not only FFN layers.
• Reduce token dropping with multi-stage No-Token-Left-Behind routing.
• Explore stochastic routing to improve expert exploration.
• Support heterogeneous experts for adaptive computation on harder examples.

Links

Prior Papers

• @vaswaniAttentionAllNeed2017 — Switch Transformer keeps the Transformer skeleton and sparsifies the FFN sublayer.
• @raffelT5_2020 — the paper builds directly on T5 architectures and compares against T5-Base, T5-Large, and T5-XXL.
• @brownGPT3_2020 — it responds to dense large-model scaling by proposing a cheaper sparse route to extreme parameter counts.

Further Papers

• @jiangMixtral2024 — extends sparse MoE language modeling with modern open-weight expert routing.
• @hoffmannChinchilla2022 — sharpens the compute-optimal scaling question that Switch attacks from the sparsity side.
• @alayracFlamingo2022 — large multimodal few-shot models inherit the same systems pressures around scaling efficient Transformers.