BASE Layers: Simplifying Training of Large, Sparse Models

Abstract

We introduce a new balanced assignment of experts (BASE) layer for large language models that greatly simplifies existing high capacity sparse layers. Sparse layers can dramatically improve the efficiency of training and inference by routing each token to specialized expert modules that contain only a small fraction of the model parameters. However, it can be difficult to learn balanced routing functions that make full use of the available experts; existing approaches typically use routing heuristics or auxiliary expert-balancing loss functions. In contrast, we formulate token-to-expert allocation as a linear assignment problem, allowing an optimal assignment in which each expert receives an equal number of tokens. This optimal assignment scheme improves efficiency by guaranteeing balanced compute loads, and also simplifies training by not requiring any new hyperparameters or auxiliary losses. Code is publicly released at https://github.com/pytorch/fairseq/

References in corpus (13)

• Adam: A Method for Stochastic Optimization
• Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer
• Language Models are Few-Shot Learners
• Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism
• Generating Long Sequences with Sparse Transformers
• Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity
• Sparse Networks from Scratch: Faster Training without Losing Performance
• Parameter Efficient Training of Deep Convolutional Neural Networks by Dynamic Sparse Reparameterization
• On Layer Normalization in the Transformer Architecture
• Generalization through Memorization: Nearest Neighbor Language Models
• Nearest Neighbor Machine Translation
• Adaptively Sparse Transformers
• Simple, Scalable Adaptation for Neural Machine Translation

Cited by in corpus (14)

• Alexa Teacher Model: Pretraining and Distilling Multi-Billion-Parameter Encoders for Natural Language Understanding Systems
• M6: A Chinese Multimodal Pretrainer
• Hash Layers For Large Sparse Models
• FlexMoE: Scaling Large-scale Sparse Pre-trained Model Training via Dynamic Device Placement
• Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation
• Taming Sparsely Activated Transformer with Stochastic Experts
• FSMoE: A Flexible and Scalable Training System for Sparse Mixture-of-Experts Models
• MPipeMoE: Memory Efficient MoE for Pre-trained Models with Adaptive Pipeline Parallelism
• 8-bit Optimizers via Block-wise Quantization
• Staircase Attention for Recurrent Processing of Sequences
• COMET: Learning Cardinality Constrained Mixture of Experts with Trees and Local Search
• Towards Structured Dynamic Sparse Pre-Training of BERT
• Beyond Distillation: Task-level Mixture-of-Experts for Efficient Inference
• Unbiased Gradient Estimation with Balanced Assignments for Mixtures of Experts