# BASE Layers: Simplifying Training of Large, Sparse Models

arXiv: https://arxiv.org/abs/2103.16716 · HTML: https://mlsanity.com/paper/2103.16716/

## Key facts

- Cited by 320 papers in the map (data through 14 Oct 2026).
- Builds on 28 of its 32 references that are in the map.
- First posted on arXiv 30 Mar 2021.

## Builds on (most cited first, 28 of 28)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [RoBERTa: A Robustly Optimized BERT Pretraining Approach](https://mlsanity.com/paper/1907.11692/) (2019)
- [BART: Denoising Sequence-to-Sequence Pre-training for Natural Language Generation, Translation, and Comprehension](https://mlsanity.com/paper/1910.13461/) (2019)
- [Layer Normalization](https://mlsanity.com/paper/1607.06450/) (2016)
- [Parameter-Efficient Transfer Learning for NLP](https://mlsanity.com/paper/1902.00751/) (2019)
- [Unsupervised Cross-lingual Representation Learning at Scale](https://mlsanity.com/paper/1911.02116/) (2019)
- [Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer](https://mlsanity.com/paper/1701.06538/) (2017)
- [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://mlsanity.com/paper/2101.03961/) (2021)
- [Neural Machine Translation of Rare Words with Subword Units](https://mlsanity.com/paper/1508.07909/) (2015)
- [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://mlsanity.com/paper/1909.08053/) (2019)
- [GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding](https://mlsanity.com/paper/2006.16668/) (2020)
- [Generating Long Sequences with Sparse Transformers](https://mlsanity.com/paper/1904.10509/) (2019)
- [Energy and Policy Considerations for Deep Learning in NLP](https://mlsanity.com/paper/1906.02243/) (2019)
- [On Layer Normalization in the Transformer Architecture](https://mlsanity.com/paper/2002.04745/) (2020)
- [Generalization through Memorization: Nearest Neighbor Language Models](https://mlsanity.com/paper/1911.00172/) (2019)
- [Beyond English-Centric Multilingual Machine Translation](https://mlsanity.com/paper/2010.11125/) (2020)
- [Rigging the Lottery: Making All Tickets Winners](https://mlsanity.com/paper/1911.11134/) (2019)
- [Efficient Content-Based Sparse Attention with Routing Transformers](https://mlsanity.com/paper/2003.05997/) (2020)
- [Simple, Scalable Adaptation for Neural Machine Translation](https://mlsanity.com/paper/1909.08478/) (2019)
- [Sparse Networks from Scratch: Faster Training without Losing Performance](https://mlsanity.com/paper/1907.04840/) (2019)
- [Nearest Neighbor Machine Translation](https://mlsanity.com/paper/2010.00710/) (2020)
- [Parameter Efficient Training of Deep Convolutional Neural Networks by Dynamic Sparse Reparameterization](https://mlsanity.com/paper/1902.05967/) (2019)
- [Adaptively Sparse Transformers](https://mlsanity.com/paper/1909.00015/) (2019)
- [Large Memory Layers with Product Keys](https://mlsanity.com/paper/1907.05242/) (2019)

## Built on by (most cited first, 30 of 320)

- [Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation](https://mlsanity.com/paper/2108.12409/) (2021)
- [GLaM: Efficient Scaling of Language Models with Mixture-of-Experts](https://mlsanity.com/paper/2112.06905/) (2021)
- [Scaling Vision with Sparse Mixture of Experts](https://mlsanity.com/paper/2106.05974/) (2021)
- [Mixture-of-Experts with Expert Choice Routing](https://mlsanity.com/paper/2202.09368/) (2022)
- [8-bit Optimizers via Block-wise Quantization](https://mlsanity.com/paper/2110.02861/) (2021)
- [ST-MoE: Designing Stable and Transferable Sparse Expert Models](https://mlsanity.com/paper/2202.08906/) (2022)
- [OLMoE: Open Mixture-of-Experts Language Models](https://mlsanity.com/paper/2409.02060/) (2024)
- [Unified Scaling Laws for Routed Language Models](https://mlsanity.com/paper/2202.01169/) (2022)
- [From Sparse to Soft Mixtures of Experts](https://mlsanity.com/paper/2308.00951/) (2023)
- [Hash Layers For Large Sparse Models](https://mlsanity.com/paper/2106.04426/) (2021)
- [MegaBlocks: Efficient Sparse Training with Mixture-of-Experts](https://mlsanity.com/paper/2211.15841/) (2022)
- [Sparse Upcycling: Training Mixture-of-Experts from Dense Checkpoints](https://mlsanity.com/paper/2212.05055/) (2022)
- [A Survey on Efficient Inference for Large Language Models](https://mlsanity.com/paper/2404.14294/) (2024)
- [Multimodal Contrastive Learning with LIMoE: the Language-Image Mixture of Experts](https://mlsanity.com/paper/2206.02770/) (2022)
- [Tutel: Adaptive Mixture-of-Experts at Scale](https://mlsanity.com/paper/2206.03382/) (2022)
- [BlenderBot 3: a deployed conversational agent that continually learns to responsibly engage](https://mlsanity.com/paper/2208.03188/) (2022)
- [Branch-Train-Merge: Embarrassingly Parallel Training of Expert Language Models](https://mlsanity.com/paper/2208.03306/) (2022)
- [OpenMoE: An Early Effort on Open Mixture-of-Experts Language Models](https://mlsanity.com/paper/2402.01739/) (2024)
- [Efficient Large Language Models: A Survey](https://mlsanity.com/paper/2312.03863/) (2023)
- [Efficient Large Scale Language Modeling with Mixtures of Experts](https://mlsanity.com/paper/2112.10684/) (2021)
- [Scaling Laws for Fine-Grained Mixture of Experts](https://mlsanity.com/paper/2402.07871/) (2024)
- [On the Representation Collapse of Sparse Mixture of Experts](https://mlsanity.com/paper/2204.09179/) (2022)
- [MoEfication: Transformer Feed-forward Layers are Mixtures of Experts](https://mlsanity.com/paper/2110.01786/) (2021)
- [Pushing Mixture of Experts to the Limit: Extremely Parameter Efficient MoE for Instruction Tuning](https://mlsanity.com/paper/2309.05444/) (2023)
- [A Review of Sparse Expert Models in Deep Learning](https://mlsanity.com/paper/2209.01667/) (2022)
- [Kimi K3: Open Frontier Intelligence](https://mlsanity.com/paper/2607.24653/) (2026)
- [AdaMix: Mixture-of-Adaptations for Parameter-efficient Model Tuning](https://mlsanity.com/paper/2210.17451/) (2022)
- [DEMix Layers: Disentangling Domains for Modular Language Modeling](https://mlsanity.com/paper/2108.05036/) (2021)
- [LLaMA-MoE: Building Mixture-of-Experts from LLaMA with Continual Pre-training](https://mlsanity.com/paper/2406.16554/) (2024)
- [Taming Sparsely Activated Transformer with Stochastic Experts](https://mlsanity.com/paper/2110.04260/) (2021)

## Abstract (arXiv)

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/