Source-linked AI summary
Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity
William Fedus, Barret Zoph, Noam Shazeer
TL;DR
Dense language-model scaling is effective but computationally intensive, while sparse models face complexity, communication, and training-stability barriers. Switch Transformer simplifies sparse expert routing and training, enabling models that outperform dense baselines across settings, including 7x+ pre-training speedups at the same FLOPs per token.
Problem
Dense scaling is computationally intensive, while sparse models have been hindered by complexity, communication costs, and training difficulties.
Method
Switch Transformer simplifies Mixture-of-Experts routing and combines sparse activation with precision, initialization, regularization, and distributed-parallelism techniques.
Results
Sparse models outperform dense models across model sizes on speed-accuracy, per-step, and wall-clock comparisons, with 7x+ pre-training speedups at the same FLOPs per token.
Takeaways & Limitations
The approach enables effective natural-language learners with hundreds of billions to trillion parameters and substantial speedups relative to dense T5 baselines.
Takeaways & Limitations
Further improving training stability for the largest models remains unsolved because the reported techniques were insufficient for Switch-XXL.
Abstract
from arXiv · showhide
In deep learning, models typically reuse the same parameters for all inputs. Mixture of Experts (MoE) defies this and instead selects different parameters for each incoming example. The result is a sparsely-activated model -- with outrageous numbers of parameters -- but a constant computational cost. However, despite several notable successes of MoE, widespread adoption has been hindered by complexity, communication costs and training instability -- we address these with the Switch Transformer. We simplify the MoE routing algorithm and design intuitive improved models with reduced communication and computational costs. Our proposed training techniques help wrangle the instabilities and we show large sparse models may be trained, for the first time, with lower precision (bfloat16) formats. We design models based off T5-Base and T5-Large to obtain up to 7x increases in pre-training speed with the same computational resources. These improvements extend into multilingual settings where we measure gains over the mT5-Base version across all 101 languages. Finally, we advance the current scale of language models by pre-training up to trillion parameter models on the "Colossal Clean Crawled Corpus" and achieve a 4x speedup over the T5-XXL model.
1. Introduction
Switch Transformer addresses the adoption barriers of sparse Mixture-of-Experts models by simplifying routing and improving training and systems efficiency. The paper reports faster pre-training, multilingual gains, large-scale trillion-parameter models, and effective distillation into small dense models.
- 1. Introduction: Switch Transformer simplifies Mixture-of-Experts routing while reducing routing computation and preserving model quality.The architecture routes each token to a single expert and is designed to reduce complexity, communication, and training difficulties.
- 1. Introduction: 7x+ pre-training speedups are measured against strongly tuned T5 models with the same FLOPS per token, including settings with as few as two experts.The reported improvements hold under limited computational resources.
- 1. Introduction: Distillation reduces model size by up to 99% while preserving 30% of the quality gains of the large sparse teacher.The paper studies distillation from sparse pre-trained and specialized fine-tuned models into small dense models.
- 1. Introduction: Selective precision, expert initialization, and increased expert regularization improve training, fine-tuning, and multi-task training for sparse models.These techniques enable lower-precision bfloat16 training and address scaling and optimization challenges.
- 1. Introduction: 91% of languages benefit from 4x+ speedups over the mT5 baseline, with universal improvement across all 101 languages.The multilingual evaluation reports gains across the full language set.
- 1. Introduction: 4x faster pre-training than T5-XXL is achieved for models with up to a trillion parameters.The models combine data, model, and expert parallelism to increase parameter scale while maintaining efficient training.
2. Switch Transformer
Switch Transformer replaces dense feed-forward layers with sparsely activated expert routing, increasing parameter count while keeping per-example computation manageable. Its single-expert routing and training techniques improve efficiency and stability, although measured speed also depends on implementation details.
- Design principle: Switch Transformer scales parameter count while keeping floating-point operations per example constant by distributing unique expert weights across devices.This preserves a manageable memory and computational footprint on each device.
- Design principle: Each token is independently routed across FFN experts, and the selected expert output is multiplied by the router gate value.The Switch FFN replaces the dense Transformer FFN and operates independently on sequence tokens.
- Simplifying sparse routing: Switch routing selects one expert instead of the top-k experts used in conventional MoE routing, preserving model quality while reducing routing computation.The authors describe this k = 1 strategy as a Switch layer.
- Simplifying sparse routing: Single-expert routing can at least halve each expert’s batch capacity and simplifies routing while reducing communication costs.These benefits arise because each token is sent to only one expert.
- Efficient sparse routing: Capacity factors buffer uneven token assignment, but higher values increase computation and memory while lower capacity can cause tokens to be dropped.Dropped tokens bypass the expert layer through the residual connection.
- Empirical comparison: Switch Transformers outperform dense and MoE baselines on a speed-quality basis, while lower capacity factors perform better in the reported comparison.The comparison is FLOP-matched to T5-Base and uses identical hardware and training steps.
- Caveat: Final speed differences reflect both Switch’s reduced algorithmic computation and low-level implementation optimizations.The paper explicitly cautions that speed measurements depend on the algorithm and implementation details.
- Training and fine-tuning: Selective precision, reduced initialization scale, and expert dropout address instability and regularization challenges in sparse training and fine-tuning.The reported initialization scheme supports stable training from a 223M-parameter baseline to models exceeding one trillion parameters.
3. Scaling Properties
Switch Transformers scale parameter count through more experts while keeping computation approximately fixed, improving sample efficiency and training speed. These gains persist against larger dense baselines, although routing introduces communication overhead.
- Step-basis scaling: One expert per token keeps computation approximately fixed as expert count increases, making experts the most efficient scaling dimension.The router’s additional computation grows as O(dmodel × num experts).
- Step-basis scaling: More experts improve test loss consistently at fixed FLOPS per token, revealing sparse parameter count as an additional scaling axis.The study scales from T5-Base’s 223M parameters to a 256-expert model with 14.7B parameters.
- Time-basis scaling: Increasing experts improves step-based performance, but communication and routing costs can prevent sample-efficiency gains from translating directly to wall-clock quality.This motivates evaluating sparse and dense models under fixed training duration and computational budget.
- Time-basis scaling: 1/7 the time: a 64-expert Switch-Base model reaches T5-Base quality under equal FLOPs per example.The comparison uses models trained on 32 TPUv3 cores.
- Time-basis scaling: 2.5x speedup: Switch-Base remains faster and more sample efficient than T5-Large despite T5-Large using 3.5x more FLOPs per token.A larger sparse Switch-Large model is also FLOP-matched to T5-Large for subsequent comparisons.
4. Downstream Results
Switch models improve downstream performance across many language tasks and can transfer part of their sparse-model gains into compact dense models. The downstream evaluation spans fine-tuning, distillation, and multilingual training.
- Fine-tuning: Fine-tuning evaluates question answering, summarization, knowledge, reasoning, sentiment, word sense, and sentence similarity across diverse benchmarks.GLUE and SuperGLUE are evaluated using average scores across their subtasks.
- Fine-tuning: FLOP-matched Switch variants improve SuperGLUE by 4.4 and 2 percentage points over T5-Base and T5-Large, respectively.Other notable gains occur on Winogrande, closed-book Trivia QA, and XSum; ARC is the stated exception.
- Distillation: ≈30% of the sparse teacher’s quality gain transfers to a small dense student using non-expert initialization and mixed teacher-ground-truth losses.The distilled student uses approximately 1/20th of the teacher’s parameters.
- Distillation: 99% compression: dense distillation still preserves 28% of the quality gain from the largest sparse models studied.For a 1.1B-parameter model, 82% compression preserves 37% of its quality gain.
- Multilingual training: 5x mean step speedup: Switch models outperform mT5-Base across all 101 languages, with at least 4x speedup in 91% of languages.The multilingual corpus spans 101 languages and 107 tasks because of script variants.
5. Designing Models with Data, Model, and Expert-Parallelism
The paper combines data, model, and expert parallelism to implement increasingly large Switch models and evaluate their scaling, speed, and downstream behavior. This design reaches trillion-parameter scale but retains a training-stability boundary at larger computational regimes.
- Data and model parallelism: Mixed data and model parallelism shards B/n tokens and dff/m weights per core, requiring all-reduce communication during forward and backward passes.The total core count is represented as N = n × m.
- Trillion-parameter models: 395 billion and 1.6 trillion parameters: the two large Switch models combine expert, model, and data parallelism.They are evaluated for upstream pre-training and downstream fine-tuning.
- Trillion-parameter models: The 1.6T-parameter Switch-C shows no training instability, whereas the higher-FLOP Switch-XXL is sometimes unstable and was not trained for the full 1M steps.The instability limits the duration of the reported Switch-XXL pre-training.
6. Related Work
Switch Transformer builds on prior work in model scaling, conditional computation, Mixture of Experts, and sparse attention. Its specific focus is input-dependent expert computation within Transformer models.
- Model scaling: Prior scaling work increased dense Transformer size through model parallelism, pipeline parallelism, or larger computational budgets.These approaches established routes to models with billions of parameters.
- Conditional computation: Conditional-computation methods dynamically select weights or expert layers based on model inputs or hidden-state patterns.Earlier examples include adaptive weight selection and stacked expert layers.
- Mixture of Experts: Modern MoE layers route tokens to combinations of experts and have achieved state-of-the-art results in language modeling and machine translation.Switch Transformer simplifies this broader MoE approach for Transformer architectures.
- Sparse attention: Attention sparsity reduces attention complexity from O(L^2) and enables longer sequences, but this Switch Transformer version does not use attention sparsity.The paper describes attention sparsity as complementary future work.
7. Discussion
The discussion positions Switch Transformers as useful across hardware scales and deployment settings, while highlighting efficiency, compression, and unresolved adoption challenges.
- Models with as few as two experts improve performance while fitting within the memory constraints of commonly available GPUs or TPUs.
- Sparse models outperform dense models for fixed computation and time across a wide range of model sizes.Controlled experiments report better per-step and wall-clock speed-accuracy trade-offs.
- 10 to 100x compression is achievable through distillation, while retaining approximately 30% of the sparse model’s quality gain.The dense student does not fully preserve the teacher’s quality.
- Switch Transformers can be more time-efficient than dense models with sharded parameters, although combining both approaches incurs conventional model-parallel slowdown.
- Sparse-model adoption remains constrained by complexity, training difficulties, and communication costs, which Switch Transformer aims to alleviate.
8. Future Work
The paper identifies unresolved limitations and future directions involving stability, scaling behavior, expert heterogeneity, attention-layer sparsity, and modalities beyond language.
- Training stability remains unsolved for Switch-XXL, although the proposed techniques were effective for Switch-Base, Switch-Large, and Switch-C.The authors identify regularizers and adapted gradient clipping as early stabilization steps.
- Similar perplexities can yield different downstream results: Switch-C scores 87.7 exact match on SQuAD versus 89.6 for Switch-XXL.The authors note that Switch-XXL uses approximately 10x the FLOPs per token despite having approximately 4x fewer unique parameters.
- A comprehensive study of scaling relationships could guide architectures and hardware that jointly balance computation, memory, and communication.
- Heterogeneous experts could enable more flexible adaptation by routing harder examples to larger experts when more computation is desired.The current approach uses identical, homogeneous experts.
- Attention-layer expert replacements showed preliminary quality gains but were excluded from the final variant because bfloat16 training was unstable.
- Future work should examine Switch Transformers across new and multimodal settings, since the paper considers only language.
9. Conclusion
The conclusion presents Switch Transformers as scalable, effective, and sample-efficient sparse language models across multiple tasks and training regimes.
- Switch Transformers simplify Mixture of Experts into an architecture that is easy to understand, stable to train, and more sample-efficient than equivalently sized dense models.
- The models perform strongly across diverse natural-language tasks in pre-training, fine-tuning, and multi-task training.
- The advances enable training models with hundreds of billions to trillions of parameters.
A. Switch for Attention
The appendix explores replacing self-attention projection matrices with Switch layers, alongside routing and capacity mechanisms, but reports greater bfloat16 instability than the standard FFN setup.
- Switch for Attention: Switch attention layers replace the trainable matrices producing queries, keys, and values within Transformer self-attention.The design uses Switch layers in place of the Q, K, and V projections.
- Switch for Attention: Attention-layer variants improved quality in preliminary experiments but were more unstable with bfloat16 and therefore omitted from the final model.
- Switch for Attention: Experts FF is the standard setup, whereas Experts FF + Attention replaces both feedforward and self-attention layers; the latter diverges with bfloat16.
- Switch for Attention: Static tensor shapes impose finite expert capacity, so uneven routing can cause padding and inefficient hardware use.
- Switch for Attention: No-Token-Left-Behind iteratively reroutes overflowing tokens to guarantee that almost no tokens are dropped during training and inference.
- Switch for Attention: The router makes a discrete expert choice from each token representation without counterfactual feedback about alternate experts.This creates an exploration-exploitation dilemma for routing.
- Switch for Attention: Input jitter is used to inject routing noise because it performed best empirically among the tested exploration approaches.
D. Switch Transformers in Lower Compute Regimes
Switch Transformer remains effective in lower-compute settings, where a small number of experts improves over dense or FLOP-matched baselines. Its routing and exploration variants support practical sparse-model configurations without material speed differences across randomness strategies.
- No-Token-Left-Behind Routing first selects each token’s highest-probability expert, then reroutes overflowed tokens to their second choice.The process can be iterated to route virtually all tokens.
- Router exploration strategies have no material speed performance difference, while quality is measured by negative log perplexity.Lower negative log perplexity is better.
- Training with 2, 4, or 8 experts results in solid improvements over T5 dense baselines, even without a supercomputer.The paper typically recommends one expert per core.
- Switch Transformer improves over T5-Base with 2, 4, and 8 experts at very small scales.
E. Relation of Upstream to Downstream Model Performance
Upstream C4 pre-training quality consistently correlates with downstream SuperGLUE and TriviaQA performance for both dense and Switch models. The relationship is similar at fixed perplexity in smaller models, but the largest Switch models do not always transfer upstream quality to SuperGLUE fine-tuning.
- For SuperGLUE, negative log perplexity has a loosely linear relation with average SuperGLUE score.
- Improved C4 pre-training leads to better downstream results for both baseline and Switch models.The downstream measures are average SuperGLUE performance and TriviaQA score.
- At fixed upstream perplexity, Switch and dense models perform similarly in the small to medium model size regime.
- In the largest model regime, T5-11B/T5-XXL Switch models do not always translate upstream perplexity well to downstream SuperGLUE fine-tuning.
F. Pseudo Code for Switch Transformers
The appendix provides Mesh TensorFlow pseudocode for Switch Transformer components and notes that the presented code uses no model parallelism. The listed components cover load balancing, routing, and the Switch Transformer layer.
- The presented Mesh TensorFlow code uses no model parallelism; further details are referenced in Section 5.4.
- The load balance loss pseudocode specifies the balancing component for Switch Transformers.
- The router pseudocode specifies the routing component used by Switch Transformers in Mesh TensorFlow.
- The appendix provides pseudocode for the Switch Transformer layer in Mesh TensorFlow.