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

Abstract

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.

arXiv commentJMLR

References in corpus (22)

• Distilling the Knowledge in a Neural Network
• Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer
• Language Models are Few-Shot Learners
• Teaching Machines to Read and Comprehend
• Longformer: The Long-Document Transformer
• Scaling Laws for Neural Language Models
• Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism
• SQuAD: 100,000+ Questions for Machine Comprehension of Text
• mT5: A massively multilingual pre-trained text-to-text transformer
• REALM: Retrieval-Augmented Language Model Pre-Training
• Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challenge
• Generating Long Sequences with Sparse Transformers
• MaskGAN: Better Text Generation via Filling in the______
• Big Bird: Transformers for Longer Sequences
• Learning Factored Representations in a Deep Mixture of Experts
• PipeDream: Fast and Efficient Pipeline Parallel DNN Training
• ZeRO: Memory Optimizations Toward Training Trillion Parameter Models
• Adversarial NLI: A New Benchmark for Natural Language Understanding
• Large Memory Layers with Product Keys
• Routing Networks: Adaptive Selection of Non-linear Functions for Multi-Task Learning
• Exponentially Increasing the Capacity-to-Computation Ratio for Conditional Computation in Deep Learning
• Scalable Transfer Learning with Expert Models

Cited by in corpus (123)

• Transformers in Vision: A Survey
• On the Opportunities and Risks of Foundation Models
• DeBERTa: Decoding-enhanced BERT with Disentangled Attention
• DeBERTaV3: Improving DeBERTa using ELECTRA-Style Pre-Training with Gradient-Disentangled Embedding Sharing
• Sustainable AI: Environmental Implications, Challenges and Opportunities
• R-Drop: Regularized Dropout for Neural Networks
• VLMo: Unified Vision-Language Pre-Training with Mixture-of-Modality-Experts
• Towards a Unified View of Parameter-Efficient Transfer Learning
• Revisiting ResNets: Improved Training and Scaling Strategies
• Restormer: Efficient Transformer for High-Resolution Image Restoration
• The NLP Cookbook: Modern Recipes for Transformer based Deep Learning Architectures
• Explaining Neural Scaling Laws
• Enable Deep Learning on Mobile Devices: Methods, Systems, and Applications
• A Practical Survey on Faster and Lighter Transformers
• Swin Transformer V2: Scaling Up Capacity and Resolution
• Characterization and Prediction of Deep Learning Workloads in Large-Scale GPU Datacenters
• COCO-LM: Correcting and Contrasting Text Sequences for Language Model Pretraining
• Exploring the Limits of Out-of-Distribution Detection
• A Survey on Deep Learning Hardware Accelerators for Heterogeneous HPC Platforms
• PanGu-$\alpha$: Large-scale Autoregressive Pretrained Chinese Language Models with Auto-parallel Computation
• Rethinking Search: Making Domain Experts out of Dilettantes
• Finetuned Language Models Are Zero-Shot Learners
• Practical Galaxy Morphology Tools from Deep Supervised Representation Learning
• BASE Layers: Simplifying Training of Large, Sparse Models
• Scale Efficiently: Insights from Pre-training and Fine-tuning Transformers
• SLAM: A Unified Encoder for Speech and Language Modeling via Speech-Text Joint Pre-Training
• Hash Layers For Large Sparse Models
• M6: A Chinese Multimodal Pretrainer
• LambdaNetworks: Modeling Long-Range Interactions Without Attention
• Dynamic Neural Networks: A Survey
• FastMoE: A Fast Mixture-of-Expert Training System
• AMMUS : A Survey of Transformer-based Pretrained Models in Natural Language Processing
• FlexMoE: Scaling Large-scale Sparse Pre-trained Model Training via Dynamic Device Placement
• From BERT to GPT-3 Codex: Harnessing the Potential of Very Large Language Models for Data Management
• Audio Transformers
• OneFlow: Redesign the Distributed Deep Learning Framework from Scratch
• Scalable and Efficient MoE Training for Multitask Multilingual Models
• Scaling Vision with Sparse Mixture of Experts
• Taming Sparsely Activated Transformer with Stochastic Experts
• Group Convolutional Neural Networks Improve Quantum State Accuracy
• Knowledge Enhanced Pretrained Language Models: A Compreshensive Survey
• Mind the Gap: Assessing Temporal Generalization in Neural Language Models
• A visual introduction to Gaussian Belief Propagation
• Handcrafted Backdoors in Deep Neural Networks
• Maximizing Parallelism in Distributed Training for Huge Neural Networks
• M6-10T: A Sharing-Delinking Paradigm for Efficient Multi-Trillion Parameter Pretraining
• AdapMoE: Adaptive Sensitivity-based Expert Gating and Management for Efficient MoE Inference
• DSelect-k: Differentiable Selection in the Mixture of Experts with Applications to Multi-Task Learning
• Pre-Trained Models: Past, Present and Future
• Accelerated Sparse Neural Training: A Provable and Efficient Method to Find N:M Transposable Masks
• M6-T: Exploring Sparse Expert Models and Beyond
• ERICA: Improving Entity and Relation Understanding for Pre-trained Language Models via Contrastive Learning
• Swin2-MoSE: A New Single Image Super-Resolution Model for Remote Sensing
• Adaptive Sparse Transformer for Multilingual Translation
• NILMFormer: Non-Intrusive Load Monitoring that Accounts for Non-Stationarity
• Primer: Searching for Efficient Transformers for Language Modeling
• Loss landscapes and optimization in over-parameterized non-linear systems and neural networks
• Augmenting Decompiler Output with Learned Variable Names and Types
• DExperts: Decoding-Time Controlled Text Generation with Experts and Anti-Experts
• On component interactions in two-stage recommender systems
• ActNN: Reducing Training Memory Footprint via 2-Bit Activation Compressed Training
• Dynamic Inference with Neural Interpreters
• Models Are Codes: Towards Measuring Malicious Code Poisoning Attacks on Pre-trained Model Hubs
• Deep Optimizer States: Towards Scalable Training of Transformer Models Using Interleaved Offloading
• Amazon SageMaker Model Parallelism: A General and Flexible Framework for Large Model Training
• Breaking the Memory Wall: A Study of I/O Patterns and GPU Memory Utilization for Hybrid CPU-GPU Offloaded Optimizers
• BaPipe: Exploration of Balanced Pipeline Parallelism for DNN Training
• The Benchmark Lottery
• Transformers with Competitive Ensembles of Independent Mechanisms
• Collaboration of Experts: Achieving 80% Top-1 Accuracy on ImageNet with 100M FLOPs
• Robust Optimization for Multilingual Translation with Imbalanced Data
• Attention Approximates Sparse Distributed Memory
• MoC-System: Efficient Fault Tolerance for Sparse Mixture-of-Experts Model Training
• Samoyeds: Accelerating MoE Models with Structured Sparsity Leveraging Sparse Tensor Cores
• Towards mental time travel: a hierarchical memory for reinforcement learning agents
• SpeechMoE: Scaling to Large Acoustic Models with Dynamic Routing Mixture of Experts
• Tournesol: A quest for a large, secure and trustworthy database of reliable human judgments
• Should we Stop Training More Monolingual Models, and Simply Use Machine Translation Instead?
• 8-bit Optimizers via Block-wise Quantization
• Distributed Deep Learning in Open Collaborations
• Are Larger Pretrained Language Models Uniformly Better? Comparing Performance at the Instance Level
• Language Models as a Knowledge Source for Cognitive Agents
• IrEne: Interpretable Energy Prediction for Transformers
• ERNIE-Tiny : A Progressive Distillation Framework for Pretrained Transformer Compression
• Layered gradient accumulation and modular pipeline parallelism: fast and efficient training of large language models
• Staircase Attention for Recurrent Processing of Sequences
• Go Wider Instead of Deeper
• DS-Net++: Dynamic Weight Slicing for Efficient Inference in CNNs and Transformers
• Moshpit SGD: Communication-Efficient Decentralized Training on Heterogeneous Unreliable Devices
• Graph Classification by Mixture of Diverse Experts
• Generating Synthetic Text Data to Evaluate Causal Inference Methods
• Detecting Hate Speech with GPT-3
• Towards More Effective and Economic Sparsely-Activated Model
• Smart Scissor: Coupling Spatial Redundancy Reduction and CNN Compression for Embedded Hardware
• Switch Spaces: Learning Product Spaces with Sparse Gating
• X-Cross: Dynamic Integration of Language Models for Cross-Domain Sequential Recommendation
• GroupBERT: Enhanced Transformer Architecture with Efficient Grouped Structures
• Strategyproof Learning: Building Trustworthy User-Generated Datasets
• COMET: Learning Cardinality Constrained Mixture of Experts with Trees and Local Search
• Towards Structured Dynamic Sparse Pre-Training of BERT
• The MIS Check-Dam Dataset for Object Detection and Instance Segmentation Tasks
• Building Compact and Robust Deep Neural Networks with Toeplitz Matrices
• Metric Transforms and Low Rank Matrices via Representation Theory of the Real Hyperrectangle
• Search Spaces for Neural Model Training
• I've Got 99 Problems But FLOPS Ain't One
• Achieving Human Parity on Visual Question Answering
• Federated Dynamic Neural Network for Deep MIMO Detection
• Sparse MoEs meet Efficient Ensembles
• DistIR: An Intermediate Representation and Simulator for Efficient Neural Network Distribution
• Beyond Distillation: Task-level Mixture-of-Experts for Efficient Inference
• A Policy Efficient Reduction Approach to Convex Constrained Deep Reinforcement Learning
• Can the Transformer Be Used as a Drop-in Replacement for RNNs in Text-Generating GANs?
• The Efficiency Misnomer
• On Anytime Learning at Macroscale
• Domain Conditional Predictors for Domain Adaptation
• When in Doubt, Summon the Titans: Efficient Inference with Large Models
• A Survey on Green Deep Learning
• Matryoshka Model Learning for Improved Elastic Student Models
• Transfer training from smaller language model
• Pinpointing the Memory Behaviors of DNN Training
• TransfoRNN: Capturing the Sequential Information in Self-Attention Representations for Language Modeling
• CARLS: Cross-platform Asynchronous Representation Learning System
• Unbiased Gradient Estimation with Balanced Assignments for Mixtures of Experts