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

1. Introduction

• Switch Transformer simplifies Mixture of Experts, activating a subset of weights per example to pursue greater computational efficiency.It addresses complexity, communication costs, and training instability that hinder MoE adoption. Superior scaling spans diverse natural-language tasks across pre-training, fine-tuning, and multi-task training, with benefits even on few computational cores.Source: §1 · pp. 3–4
• 7x+ pre-training speedups over strongly tuned T5 (Raffel et al., 2019) are achieved with the same FLOPS per token.Improvements also hold with as few as two experts. Figure 1 compares negative log perplexity against T5 models using the same compute budget.Source: Figure 1 · p. 3, §1 · pp. 3–4
• Up to 99% model-size reduction preserves 30% of the large sparse teacher’s quality gains through distillation into small dense models.Distillation succeeds for both sparse pre-trained models and specialized fine-tuned models.Source: §1 · p. 4
• Selective precision enables bfloat16 training, improved initialization supports more experts, and increased expert regularization improves sparse-model fine-tuning and multi-task training.Source: §1 · p. 4
• 91% of languages achieve 4x+ pre-training speedups over mT5 (Xue et al., 2020), with improvements across all 101 languages.Source: §1 · p. 4
• 4x faster pre-training than a strongly tuned T5-XXL baseline is achieved with models containing up to a trillion parameters.These models efficiently combine data, model, and expert parallelism.Source: §1 · p. 4

2. Switch Transformer

• Distributing unique expert weights across devices increases parameter count while keeping per-example FLOPs constant and each device’s memory and computational footprint manageable.The implementation targets TPUs, although the authors suggest GPU clusters could also train these models. Mesh-Tensorflow (Shazeer et al., 2018) supports partitioning tensors and computations across a logical processor mesh.Source: §2 · p. 4, §2.2 · p. 6

2.1 Simplifying Sparse Routing

• Single-expert routing replaces top-k mixtures, preserving model quality while reducing routing computation and communication.Where top-k routing combines multiple gated expert outputs, Switch returns one gate-weighted expert output and retains router differentiability. Routing each token once also at least halves required expert capacity and simplifies implementation.Source: §2.1 · pp. 5–6, Figure 2 · p. 5

2.2 Efficient Sparse Routing

• Fixed expert capacities accommodate static tensor shapes, while overflow tokens bypass expert computation through residual connections.Capacity equals (batch tokens / experts) × capacity factor; larger factors buffer imbalance but waste computation and memory. Dropped-token rates were typically <1%, with no observed dependence on expert count.Source: Figure 3 · p. 6, §2.2 · p. 7
• An auxiliary load-balancing loss encourages uniform expert utilization through a scaled dot product of token-dispatch fractions and router-probability fractions.Router-probability fractions are differentiable, unlike dispatch fractions; multiplying by expert count keeps the loss constant as expert count changes. A coefficient of 10^-2 balanced load quickly without interfering with training loss.Source: §2.2 · pp. 7–8

2.3 Putting It All Together: The Switch Transformer

• Switch Transformers outperform tuned dense and top-2 MoE Transformers on speed–quality, while speed-matched Switch-Base+ also outperforms both per training step.On C4 masked-language-model pre-training, models use identical hardware and training steps; Switch is FLOP-matched to T5-Base, and quality is measured by negative log perplexity. Switch performs better at capacity factors 1.0 and 1.25, although measured speed also depends on implementation optimizations.Source: §2.3 · p. 8, §2.4 · p. 8, Table 1 · p. 9

2.4 Improved Training and Fine-Tuning Techniques

• Selective float32 routing stabilizes bfloat16 training while retaining nearly bfloat16 speed and avoiding float32 all-to-all communication.Hard routing decisions and low-precision router softmax can cause instability. Router inputs and local calculations use float32, but dispatch and combine tensors return to bfloat16 before communication.Source: §2.4 · pp. 8–10, Table 2 · p. 10
• 10-fold smaller initialization scale improves Switch Transformer quality and reduces training instability and variation across runs.The method reduces the default initialization scale from 1.0 to 0.1. The same initialization approach supports stable training from the 223M-parameter baseline to models exceeding one trillion parameters.Source: §2.4 · p. 10, Table 3 · p. 10
• 0.1 dropout in non-expert layers and 0.4 in expert layers improve Switch Transformer fine-tuning performance on four smaller downstream tasks.Expert dropout targets overfitting associated with the sparse model’s larger parameter count. Increasing dropout across all layers instead worsens performance.Source: §2.4 · p. 11

3. Scaling Properties

• Scaling expert count keeps computation approximately fixed because each token selects one expert, with only lightweight additional routing.Router computation costs O(dmodel × num experts), where dmodel is the embedding dimension. Pre-training uses over 180B C4 target tokens (Raffel et al., 2019) and continues until diminishing returns, targeting a regime without data or computational-budget bottlenecks.Source: §3 · p. 11

3.1 Scaling Results on a Step-Basis

• 7.5x step-based speedup is reported for 64-expert Switch-Base versus T5-Base at fixed FLOPs per token.Increasing expert count consistently improves perplexity and sample efficiency at equal computational budgets. The sweep reaches 256 experts and 14.7B parameters, compared with T5-Base’s 223M parameters.Source: Figure 4 · p. 12, §3.1 · p. 12

3.2 Scaling Results on a Time-Basis

• One-seventh the training time lets 64-expert Switch-Base reach similar perplexity to T5-Base under a fixed computational budget.All models use 32 TPUv3 cores and equal FLOPs per example. This wall-clock advantage persists despite additional inter-device communication and routing computation.Source: Figure 5 · p. 13, §3.2 · p. 13

3.3 Scaling Versus a Larger Dense Model

• 2.5x wall-clock speedup over T5-Large is achieved by Switch-Base, which is also more sample efficient despite T5-Large using 3.5x more FLOPs per token.The authors also introduce Switch-Large, FLOP-matched to T5-Large, for scaling and fine-tuning evaluation in the following section.Source: §3.3 · pp. 13–14, Figure 6 · p. 14

4.1 Fine-Tuning

• Fine-tuning evaluates language understanding, summarization, question answering, factual knowledge, commonsense reasoning, and adversarial natural language inference.Metrics include average subtask scores for GLUE and SuperGLUE, Rouge-2 for summarization, exact match for question answering, and generated-response accuracy for ARC, ANLI, and Winogrande.Source: §4.1 · p. 15
• 4.4 and 2 percentage-point SuperGLUE gains over T5-Base and T5-Large, respectively, accompany broader downstream improvements from FLOP-matched Switch models.Notable gains also occur on Winogrande, closed-book Trivia QA, and XSum. Exceptions are ARC Challenge, where T5-Base outperforms Switch-Base, and ARC Easy, where T5-Large outperforms Switch-Large.Source: §4.1 · p. 15, Table 5 · p. 16

4.2 Distillation

• The best distillation method combines non-expert weight initialization with a loss mixing 0.25 teacher probabilities and 0.75 ground-truth labels.FLOP matching gives non-expert layers matching dimensions, allowing the dense student to inherit trained parameters.Source: §4.2 · p. 16, Table 6 · p. 17
• 99% compression preserves 28% of the sparse teacher’s quality improvement, while 82% compression preserves 37% for the 1.1B-parameter teacher.Retained quality gain measures how much of the Switch-Base teacher’s improvement over T5-Base survives in the dense student.Source: §4.2 · p. 17
• 30% of the sparse model’s SuperGLUE gains survive distillation from a 7.4B-parameter Switch-Base teacher into 223M-parameter T5-Base.The resulting dense model is 97% compressed relative to its teacher.Source: §4.2 · p. 17, Table 8 · p. 18, §4.3 · p. 18

4.3 Multilingual Learning

• 5x mean speed-up over mT5-Base accompanies improved final negative log perplexity in all 101 languages, with 91% achieving at least 4x speed-up.The FLOP-matched mSwitch-Base and mT5-Base models were pre-trained for 1M steps on multilingual mC4.Source: §4.3 · pp. 17–18

5.1 Data Parallelism

• Data parallelism assigns all N cores to batch sharding, with n = N and m = 1, deferring gradient aggregation until forward and backward passes finish.No communication is needed before this final gradient aggregation.Source: §5.1 · p. 20

5.2 Model Parallelism

• Model parallelism keeps all B tokens on every core, shards weights across N cores, and requires all-reduce communication during forward and backward passes.With n = 1 and m = N, each core communicates a [B, dmodel] tensor because the partitioned dff dimension must be summed.Source: §5.2 · p. 20

5.3 Model and Data Parallelism

• Combining data and model parallelism assigns each core B/n tokens and dff/m weight and activation slices, with N = n × m.Each core communicates a tensor of shape [B/n, dmodel] through all-reduce during forward and backward passes.Source: §5 · p. 20, §5.3 · p. 21

5.4 Expert and Data Parallelism

• Expert–data parallelism places one expert per core, computes routing assignments locally, and uses all-to-all communication to dispatch tokens and receive expert outputs.A binary routing tensor of shape [n, B/n, E, C] gathers tokens into [n, E, C, dmodel], which communication reshards along the expert dimension. Each of the E experts can process up to C tokens.Source: §5 · p. 20, §5.4 · p. 22

5.5 Expert, Model and Data Parallelism

• Combining all three strategies trades parameter count, FLOPs, memory, batch size, and all-to-all and all-reduce communication, requiring empirically chosen mappings.Adding experts increases parameters without increasing FLOPs per token, whereas increasing dff raises FLOPs and can require more model sharding and smaller batches; Switch-C uses only expert parallelism, while FLOP-matched Switch-XXL adds model parallelism. Switch-C, with 1.6T parameters and 2048 experts, exhibits no training instability, whereas Switch-XXL, with nearly 10x larger FLOPs per sequence, is sometimes unstable and was not pre-trained for a full 1M steps.Source: §5.5 · p. 22, §5.6 · pp. 22–23

5.6 Towards Trillion Parameter Models

• 4x faster progress to a fixed perplexity distinguishes Switch-C from T5-XXL under the same compute budget, with the gap increasing during training.The two variants contain 395 billion and 1.6 trillion parameters; both improve C4 negative log perplexity over T5-XXL by over 0.061 after 250k steps, while Switch-XXL improves it by 0.087 after 500k steps. These differences are lower bounds because T5-XXL used an easier C4 dataset containing duplicated snippets.Source: §5.6 · pp. 22–23, Table 9 · p. 23
• 65.7 accuracy on ANLI (Nie et al., 2019) surpasses the prior best of 49.4, while other reasoning benchmarks remain below state-of-the-art.Multi-task evaluation uses a Switch-XXL checkpoint pre-trained on 503B tokens, approximately half T5-XXL’s text. SQuAD validation accuracy reaches 89.7 versus the state-of-the-art 91.3; SuperGLUE averages 87.5 on the test set versus T5’s 89.3 and the state-of-the-art 90.0 (Wang et al., 2019).Source: §5.6 · p. 24
• 34.4 versus 32.8 on Natural Questions, 41.0 versus 37.2 on WebQuestions, and 47.5 versus 42.9 on TriviaQA improve over T5-XXL without Salient Span Masking.Natural Questions measures exact match; these knowledge-task evaluations omit additional Salient Span Masking pre-training (Guu et al., 2020). Overall, upstream gains transfer better to knowledge-based tasks than reasoning tasks.Source: §5.6 · p. 24

6. Related Work

• Model parallelism scales networks to billions of parameters by splitting weights and tensors across cores, while pipeline parallelism distributes layers and pipelines micro-batches.Model-parallel examples include Shazeer et al. (2018), Rajbhandari et al. (2019), Raffel et al. (2019), Brown et al. (2020), and Shoeybi et al. (2019). Harlap et al. (2018); Huang et al. (2019) propose pipeline-based model parallelism.Source: §6 · p. 24
• Product Key networks (Lample et al., 2019) increase model capacity through learnable-embedding lookups keyed by incoming token representations.The lookup uses the token representations entering a given layer.Source: §6 · p. 24
• Conditional computation dynamically makes computation decisions based on inputs, encompassing adaptive weight selection, stacked expert layers, and semantic-class routing.Cho and Bengio (2014) selected weights using hidden-state bit patterns; Eigen et al. (2013) stacked dense, ReLU-based expert layers, obtaining promising jittered MNIST and monotone speech results. Puigcerver et al. (2020) manually routed tokens by semantic class during vision pre-training and selected relevant experts for downstream tasks.Source: §6 · p. 24
• Shazeer et al. (2017) routed tokens to expert combinations between LSTM layers, achieving state-of-the-art language modeling and machine translation results.Mesh Tensorflow (Shazeer et al., 2018) subsequently substituted MoE layers for Transformer FFN layers, without accompanying NLP results.Source: §6 · p. 25
• 100-language machine translation improved dramatically with GShard (Lepikhin et al., 2020), while Fan et al. (2021) deterministically partitioned parameters into non-overlapping language groups.GShard extended the XLA compiler to support its MoE Transformer.Source: §6 · p. 25
• Attention sparsity reduces O(L^2) sequence-length complexity and enables longer sequences; this Switch Transformer omits it, but combining the techniques could benefit long-context tasks.Prior attention-sparsity work includes Child et al. (2019), Correia et al. (2019), Sukhbaatar et al. (2019), Kitaev et al. (2020), Zaheer et al. (2020), and Beltagy et al. (2020).Source: §6 · p. 25

Parameter scaling

• Switch Transformers deliberately increase parameter counts independently of FLOPs, achieving greater sample efficiency and faster execution with the same computational resources.The discussion treats parameter count as a useful scaling axis distinct from total computation.Source: §7 · p. 25

Small-scale use

• Two experts can suffice to improve performance while fitting within the memory constraints of commonly available GPUs or TPUs.These results support using the techniques in small-scale settings.Source: §7 · p. 25

Speed–accuracy tradeoff

• Sparse models outperform dense models on the speed–accuracy Pareto curve across model sizes, both per step and in wall-clock time.Controlled experiments show better performance for fixed computation and time.Source: §7 · p. 25

Compression

• 10 to 100x compression is achievable by distilling sparse models into dense models while retaining ≈30% of the expert model’s quality gain.Distillation does not fully preserve model quality.Source: §7 · p. 25

Model parallelism

• Switch Transformers can be more time-efficient than dense models with sharded parameters, while also supporting model parallelism.Using model parallelism increases FLOPs per token but incurs the slowdown of conventional model parallelism.Source: §7 · pp. 25–26

Adoption barriers

• Switch Transformer alleviates complexity, training difficulties, and communication costs that have hindered sparse-model adoption.Dense-model scaling success has also discouraged exploration; Hooker (2020) attributes part of that success to co-adaptation with deep learning hardware.Source: §7 · p. 26

8. Future Work

• Training stability remains unsolved for Switch-XXL, although Switch-Base, Switch-Large, and Switch-C showed no observed instability with the proposed techniques.Early stabilization efforts include regularizers and adapted forms of gradient clipping.Source: §8 · p. 26
• 87.7 SQuAD exact match for 1.6T-parameter Switch-C trails Switch-XXL’s 89.6 despite similar C4 perplexities, highlighting poorly understood relationships between fine-tuning quality, computation, and parameter count.Improved pre-training quality generally improves downstream results, but anomalies remain. Switch-XXL uses ≈10x the FLOPS per token despite ≈4x fewer unique parameters than Switch-C (395B vs 1.6T).Source: §8 · p. 26
• Comprehensive scaling studies could guide architectures combining data, model, and expert parallelism around hardware computation, memory, and communication constraints, while also informing future hardware design.The goal is to identify an optimal model more rapidly from a hardware configuration’s specifications.Source: §8 · p. 26
• Heterogeneous experts could replace the current homogeneous experts, enabling routing to larger experts when more computation is desired, potentially for harder examples.Such designs would require more flexible infrastructure.Source: §8 · p. 26
• Expert layers beyond Transformer FFNs show preliminary quality improvements in self-attention, but bfloat16 training instability leaves this extension for future work.The self-attention variant replaces the weight matrices that produce Q, K, and V.Source: §8 · pp. 26–27
• Switch Transformer has only been studied in language; future work would test whether sparsity also benefits other modalities and multimodal networks.Benefits outside language remain a proposed direction rather than an established result.Source: §8 · p. 27

• Hundreds of billions to a trillion parameters are trainable with Switch Transformers, which achieve substantial speedups relative to dense T5 baselines.Source: §9 · p. 27
• Simplifying Mixture of Experts yields an architecture that is easy to understand, stable to train, and vastly more sample efficient than equivalently sized dense models.Source: §9 · p. 27
• Switch Transformers excel across diverse natural language tasks under pre-training, fine-tuning, and multi-task training.Source: §9 · p. 27

A. Switch for Attention

• Switch attention replaces query, key, and value projection matrices with expert layers, extending sparsity beyond the standard feedforward replacement.The design uses separate expert weights for queries and shared keys and values, testing both linear and feedforward experts.Source: A. Switch for Attention · p. 27, Figure 10 · p. 28
• Switch attention improved quality but was more unstable in bfloat16, so it was excluded from the final variant.The preliminary positive results from stable training runs suggest a promising future direction.Source: A. Switch for Attention · pp. 27–28, Table 10 · p. 28

B. Preventing Token Dropping with No-Token-Left-Behind

• Static TPU tensor shapes impose fixed expert capacities, making uneven routing require padding for underfilled experts and overflow handling for overloaded experts.The standard approach bypasses expert computation for overflow tokens through a residual connection, following Lepikhin et al. (2020).Source: B. Preventing Token Dropping with… · p. 29
• No-Token-Left-Behind iteratively reroutes overflow tokens to prevent almost all dropping during training and inference, but yielded no empirical benefits.The authors suspect that rerouting to lower-ranked experts can disrupt learned token–expert associations and degrade performance.Source: B. Preventing Token Dropping with… · p. 29

C. Encouraging Exploration Across Experts

• Deterministic top-expert routing exploits current preferences without counterfactual feedback about alternative experts, motivating exploration to seek better assignments.The router makes a discrete expert choice conditioned on each token’s representation.Source: C. Encouraging Exploration Across… · p. 29
• Multiplicative input jitter performed best among the tested exploration strategies and was used throughout the work.The alternatives were deterministic argmax routing, sampling from the softmax distribution, and dropout on the incoming representation.Source: C. Encouraging Exploration Across… · p. 29

D. Switch Transformers in Lower Compute Regimes

• 2 experts already produce compelling gains over a FLOP-matched counterpart.Many prior experiments instead used models with 10B+ parameters.Source: D. Switch Transformers in Lower… · pp. 29–30
• The authors typically recommend one expert per core, supporting small configurations when a supercomputer is unavailable.They identify configurations with 2, 4, or 8 experts as yielding solid improvements over dense T5 baselines.Source: D. Switch Transformers in Lower… · p. 30
• 2, 4, and 8 experts all yield Switch Transformer improvements over T5-Base.The gains extend the architecture’s effectiveness to small-scale configurations.Source: D. Switch Transformers in Lower… · pp. 29–30, Figure 12 · p. 31

E. Relation of Upstream to Downstream Model Performance

• Both dense and Switch models achieve better downstream results as C4 pre-training quality improves.Figure 13 compares validation performance on reasoning-oriented SuperGLUE and knowledge-heavy TriviaQA. At fixed upstream perplexity, dense and Switch models perform similarly at small to medium scales.Source: Figure 13 · p. 32, E. Relation of Upstream to Downstream… · p. 32
• Dense models often outperform Switch models on SuperGLUE at fixed upstream perplexity, particularly at large scales.Average SuperGLUE score is loosely linear in negative log perplexity, but the largest Switch models do not always translate upstream quality well in the T5-11B/T5-XXL regime. Expert-model fine-tuning depends on regularization, load-balancing, and fine-tuning hyper-parameters, warranting further investigation.Source: Figure 13 · p. 32, E. Relation of Upstream to Downstream… · p. 32
• Switch Transformers may outperform dense counterparts on TriviaQA at fixed upstream perplexity, suggesting an improved scaling relationship for factual knowledge.Further statistics, which are expensive to collect and left to future work, are needed to confirm these observations.Source: Figure 13 · p. 32

F. Pseudo Code for Switch Transformers

• Figure 14 presents pseudocode for the Switch Transformer load balance loss.Source: Figure 14 · p. 33
• Figure 15 presents pseudocode for the Switch Transformer router in Mesh Tensorflow.Source: Figure 15 · p. 34
• Figure 16 presents pseudocode for the Switch Transformer layer in Mesh Tensorflow.The appendix's pseudocode uses no model parallelism.Source: F. Pseudo Code for Switch Transformers · p. 33, Figure 16 · p. 35