Source-linked AI summary
BASE Layers: Simplifying Training of Large, Sparse Models
Mike Lewis, Shruti Bhosale, Tim Dettmers, Naman Goyal, Luke Zettlemoyer
TL;DR
Sparse expert models can improve capacity and efficiency, but balancing their routing is difficult and prior solutions require heuristics, auxiliary losses, or tuning. BASE instead solves token allocation as a balanced linear assignment problem and combines it with softly mixed expert modules. Experiments report strong performance against dense and sparse baselines, with a favorable cost–accuracy tradeoff and no new hyperparameters or training-loss modifications.
Problem
Sparse expert models require balanced routing for specialization and efficiency, while prior approaches use heuristics, auxiliary losses, or tuning that do not guarantee balance.
Method
BASE formulates token-to-expert routing as a linear assignment maximizing affinities while assigning equal token counts to experts, with soft expert mixing.
Results
BASE shows strong performance relative to dense and previously proposed sparse models, including similar performance to Sparsely Gated MoE and better validation perplexity than Switch.
Takeaways & Limitations
BASE provides a drop-in conditional-compute layer that increases model capacity without new hyperparameters or training-loss modifications.
Takeaways & Limitations
Results may differ on other networking hardware because communication between workers is a significant overhead for model-parallel and sparse-expert approaches.
Abstract
from arXiv · showhide
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/
1. Introduction
BASE replaces heuristic or loss-based expert balancing with a linear assignment that maximizes token–expert affinity while giving every expert an equal token load. The resulting sparse layer avoids new balancing hyperparameters and training losses while achieving strong efficiency and performance.
- Sparse expert models use only a small subset of parameters per input, but training requires carefully balancing experts for specialization.
- Existing mixtures-of-experts methods route tokens through learned gating functions, requiring tuning and auxiliary losses to avoid imbalanced or degenerate experts.
- BASE formulates routing as a linear assignment problem that maximizes token–expert affinities while assigning each expert an equal number of tokens.
- Softly mixing expert contributions helps experts specialize without routing tokens to multiple experts or adding another loss term.
- Models up to 110B parameters show large gains over standard parallel training, while matching or exceeding prior sparse approaches at equal computation budgets.
2. Background: Training with Multiple Workers
Large language models require distributed training because they exceed single-device capacity. Data and model parallelism distribute work differently, while sparse expert layers add capacity but make communication and expert balancing central challenges.
- Large language models are distributed across workers because training them on a single device would take infeasibly long.
- Dense models use data or model parallelism, with data parallelism replicating models and model parallelism distributing parameters and computation.
- Sparse models route each token to a small subset of worker-held experts, then return the resulting representation to the original worker.
- Expert-layer training can require four routing communications per layer, allowing workers to idle while waiting for communication.
- Balanced expert loads prevent memory exhaustion, reduce slow-worker bottlenecks, and help rarely used experts receive adequate training.
- Prior balancing losses require careful weighting and still do not guarantee balance, so systems may impose hard token-capacity limits that discard excess tokens.
3. BASE Layers
BASE layers compute token–expert affinities, solve a balanced assignment during training, route tokens across workers, and softly combine expert outputs with the input. Training uses globally balanced routing, whereas inference greedily selects the best expert to avoid future-context leakage.
- BASE Layers: BASE computes token–expert scores, finds a balanced assignment maximizing those scores, routes features to experts, and returns outputs to the original workers.
- BASE Layers: The implementation computes affinities by multiplying shuffled token features by transposed expert centroids before determining token destinations.
- Parameterization: Each expert is a position-wise function with an embedding in the model dimension, and a modified residual connection softly mixes expert and input representations.
- Assignment During Training: During training, every expert receives an equal number of tokens so workers remain fully utilized and finish with similar loads.
- Assignment During Training: The assignment solves a linear assignment problem over token representations and expert embeddings, using an auction algorithm chosen for GPU parallelizability.
- Sharding: Assignment computation is sharded across workers, with each worker solving a smaller problem and sending T/E tokens to every other worker through all-to-all communication.
- Shuffling and Inference: Random shuffling distributes correlated within-worker tokens before assignment, while inference greedily chooses each token’s highest-scoring expert because balanced routing would leak future context.
4. Experiments
Experiments evaluate BASE on language modeling under matched compute budgets, comparing dense and sparse baselines while varying BASE layer size, position, and routing strategy. BASE matches dense training at 8 GPUs, gains increasingly at higher budgets, and matches or exceeds prior sparse approaches with limited tuning sensitivity.
- 4.1. Experimental Setup: Experiments compare language models using compute efficiency: perplexity achieved for a fixed number of GPUs and runtime.Models are trained for approximately 2.5 days with matched evaluation budgets.
- 4.2. Comparison with Dense Models: With larger compute budgets, BASE layers outperform both data parallel and model parallel training by a wide margin.BASE matches data parallel training in the 8-GPU experiments and gains increasingly at higher compute regimes.
- 4.2. Comparison with Dense Models: Model parallelism first outperforms data parallel training after 128 GPUs for 2 days, partly because it requires a reduced batch size.This comparison uses the same computational resources.
- 4.3. Comparison with Sparse Experts Models: BASE achieves similar performance to Sparsely Gated MoE and converges to better validation perplexity than Switch.The experiments use 128 experts per expert layer; BASE has 44.4B parameters versus 52.5B for the other sparse models.
- 4.4. Ablations: BASE performance remains strong across small, standard, and large expert configurations, indicating that shared-versus-expert capacity allocation needs little tuning.The tested configurations range from 18.8B to 117B total parameters.
- 4.4. Ablations: Three BASE placements show similar performance, while interleaving three BASE layers throughout the transformer stack achieves the best performance.The strong performance of BASE Top may support adding capacity to pretrained language models.
5. Analysis
BASE learns reasonably balanced expert assignments without an explicit balancing loss, while expert specialization often reflects local token and syntactic signals. Its reduced communication overhead also improves training token throughput.
- Expert Balancing: BASE learns reasonably balanced routing without an auxiliary balancing loss, although its top-1 assignment is less balanced than explicitly regularized models.It is more balanced than the second expert in Sparsely Gated MoE.
- Expert Specialization: Expert assignments often depend on local syntactic information, with the token at timestep t frequently indicating its assigned expert.Frequent selected-token clusters include quantities, numbers, possessives, subword fragments, verbs, nouns, and adjectives.
- Expert Balancing: Figure 8 measures the percentage of tokens assigned to each of 128 experts on the validation set across sparse expert approaches.Sparsely Gated MoE and Switch results are averaged across expert layers.
- Efficiency: BASE computes updates faster than other multi-worker approaches because reduced communication lowers overhead.The paper notes that slower token processing can improve sample efficiency at equal compute efficiency in lower-data regimes.
6. Related Work
Related work uses sparse experts, sparse training, memory or nearest-neighbor layers, and sparse attention to increase model capacity or efficiency. BASE further simplifies sparse expert layers by removing balancing losses and demonstrating effectiveness with a single expert layer.
- Sparse Expert Models: Sparse expert models route inputs to specialist workers, while Switch simplifies this design by routing each token to only one worker.
- Sparse Expert Models: BASE extends sparse expert work by eliminating balancing loss functions and showing that a single expert layer can be effective.
- Sparse Training: Fine-grained sparse training reduces FLOPS but can limit runtime benefits on GPUs, whereas contiguous expert blocks support efficient GPU processing.
- Specialized Layers: Task- or language-specific layers hard-code expert assignment, while BASE explores learnable routing when such structure is unavailable.
- High-Capacity Layers: Alternative high-capacity approaches include efficient sparse memory layers and nearest-neighbor classifiers over training data.
- Sparse Attention: Sparse attention improves transformer attention efficiency, whereas BASE develops sparsity for the feed-forward network.
7. Conclusion
The paper presents BASE as a simple sparse layer that increases neural-model capacity with little added training cost or complexity, while achieving strong performance relative to dense and prior sparse models.
- Conclusion: BASE increases the capacity of any neural model with little increase in training cost or complexity.The paper reports strong performance relative to both dense models and previously proposed sparse models.