Source-linked AI summary
Efficient Large Scale Language Modeling with Mixtures of Experts
Mikel Artetxe, Shruti Bhosale, Naman Goyal, Todor Mihaylov, Myle Ott, Sam Shleifer, Xi Victoria Lin, Jingfei Du, Srinivasan Iyer, Ramakanth Pasunuru, Giri Anantharaman, Xian Li, Shuohui Chen, Halil Akin, Mandeep Baines, Louis Martin, Xing Zhou, Punit Singh Koura, Brian O'Horo, Jeff Wang, Luke Zettlemoyer, Mona Diab, Zornitsa Kozareva, Ves Stoyanov
TL;DR
The paper examines whether sparse Mixture-of-Experts models offer better performance–computation trade-offs than dense language models across scales and evaluation settings. Through a comprehensive empirical comparison, it finds that sparse models remain competitive while using less computation, although their advantage narrows at scale and varies across tasks and domains.
Problem
Growing computational and environmental costs motivate evidence on whether more efficient model designs can retain language-model and downstream-task performance.
Method
The authors train dense and MoE language models from hundreds of millions to over one trillion parameters and compare them across language modeling, zero-shot, few-shot, and supervised settings.
Results
Sparse models offer better performance–computation trade-offs than dense counterparts for language modeling, zero-shot learning, and few-shot learning; at modest budgets, they can match a dense model using almost four times less compute.
Takeaways & Limitations
Sparse MoE models provide an alternative to dense architectures that saves computation and can reduce model energy consumption, while generalization differences across tasks and domains merit further study.
Takeaways & Limitations
The study evaluates one MoE configuration and does not explore MoE-specific hyperparameters such as the number of experts; factors favoring MoEs on particular tasks remain unclear.
Abstract
from arXiv · showhide
Mixture of Experts layers (MoEs) enable efficient scaling of language models through conditional computation. This paper presents a detailed empirical study of how autoregressive MoE language models scale in comparison with dense models in a wide range of settings: in- and out-of-domain language modeling, zero- and few-shot priming, and full-shot fine-tuning. With the exception of fine-tuning, we find MoEs to be substantially more compute efficient. At more modest training budgets, MoEs can match the performance of dense models using $\sim$4 times less compute. This gap narrows at scale, but our largest MoE model (1.1T parameters) consistently outperforms a compute-equivalent dense model (6.7B parameters). Overall, this performance gap varies greatly across tasks and domains, suggesting that MoE and dense models generalize differently in ways that are worthy of future study. We make our code and models publicly available for research use.
1 Introduction
The paper studies whether sparse MoE models can match dense-model accuracy with substantially less computation across language modeling and downstream learning settings. It finds competitive performance at lower compute and differences in how sparse and dense models generalize at scale.
- Scaling language models improves accuracy and generalization but increases computational resource requirements, creating research and environmental concerns.
- Sparse MoE models increase the number of learnable parameters without proportionally increasing computational costs.
- The study trains dense and MoE models from hundreds of millions to more than one trillion parameters and compares them in zero-shot, few-shot, and fully supervised settings.
- MoE models can match dense-model downstream performance at a fraction of the compute, with modest-budget models requiring almost four times less compute.
- Sparse MoE models show competitive zero- and few-shot performance while using a fraction of the computation for training and inference.
- Dense and sparse models exhibit different generalization behavior at scale, suggesting complementary behavior and a future research direction.
2 Background and Related Work
The background contrasts conventional dense scaling and fine-tuning with conditional-computation approaches and the emergence of zero- and few-shot learning. It also situates the study’s approximately FLOP-matched dense and MoE model comparisons and large-scale training techniques.
- Transformer language models are pretrained by predicting, masking, or perturbing text and can be adapted to tasks through supervised fine-tuning.
- GPT-3 showed that large language models can perform zero- and few-shot learning through in-context examples without fine-tuning, with these behaviors emerging or strengthening at scale.
- Conditional-computation strategies activate only a subset of a model for each input, reducing the computational cost of increasing model capacity.
- Dense and MoE models are compared in roughly FLOP-matched configurations, with MoE layers using alternating dense and expert layers plus top-2 expert selection.
- The paper contrasts zero-shot, few-shot, and fully supervised fine-tuning performance between dense and MoE models.
- Large models require techniques including FP16 training, activation checkpointing, and fully sharded data parallel training because parameter storage can exceed one GPU’s usable memory.
3 Experimental Setup
The study compares dense and MoE autoregressive Transformers under matched training designs across language modeling and downstream evaluation settings. It uses in-domain and out-of-domain perplexity, zero- and few-shot priming, and fine-tuning benchmarks with GPT-3-based comparisons.
- 3.1 Models: MoE models mirror dense configurations for approximately FLOPs-matched comparisons, using alternating dense and expert layers with top-2 routing across 512 experts.Each expert has capacity 2·B/E tokens; overflowed tokens pass through the residual connection.
- 3.1 Models: Models are autoregressive decoder-only Transformers trained for 300B tokens with 2048-token contexts and Adam optimization.Training uses learning-rate warm-up and decay, weight decay of 0.01, and dropout of 0.1.
- 3.3 Evaluation: Language modeling is evaluated with next-token perplexity on held-out in-domain data and The Pile's 22 out-of-domain subsets.The study reports each Pile subset's official test perplexity and their average.
- 3.3 Evaluation: Perplexity evaluation chunks inputs into nonoverlapping 2048-token blocks, so the first tokens in each block lack preceding-block context.A sliding-window approach could produce better results but would be more expensive.
- 3.3 Evaluation: Downstream evaluation compares zero-shot, few-shot, and fully supervised fine-tuning approaches on selected tasks from Brown et al. (2020).Few-shot results average 25 runs, while task selection emphasizes settings where GPT-3 showed scaling or zero-to-few-shot gains.
4 Results and Analysis
Across language modeling and downstream evaluation, MoE models outperform dense counterparts, but their compute advantage varies by domain, task, and scale. Few-shot gains remain positive for larger MoEs, whereas supervised fine-tuning can underperform dense models.
- Language modeling perplexity: MoEs outperform dense models across all language-modeling datasets, with larger gains in-domain than out-of-domain.They match dense models trained with 8–16x more compute in-domain, versus a 2–4x speedup on the Pile out-of-domain.
- Language modeling perplexity: MoE efficiency decreases with scale, from roughly 4x less compute at 2–6 ZFLOPs to roughly 2x at 30 ZFLOPs.
- Zero-shot learning: MoEs outperform dense models on every zero-shot dataset and model size, with speedups largest on HellaSwag and PIQA and more modest on ReCoRD and WinoGrande.
- Few-shot learning: Few-shot improvements are smaller for MoEs than dense models: the 6.7B dense model gains 3.6 points to 69.3, while the 1.1T MoE gains 2.3 points to 70.1.
- Supervised Fine-Tuning: Supervised fine-tuning substantially improves dense models, but MoE fine-tuning worsens HellaSwag, PIQA, and WinoGrande while improving several other datasets.Where MoE fine-tuning helps, its accuracy approaches that of corresponding dense models; the largest models were omitted because of resource needs.
5 Conclusion
The study scales sparse language models to 1.1T parameters and finds better performance–computation trade-offs than dense counterparts for language modeling and zero- and few-shot learning. The gap narrows at scale, but the largest sparse model still outperforms its dense counterpart with half the computation.
- Sparse models up to 1.1T parameters offer better performance–computation trade-offs than dense models for language modeling, zero-shot learning, and few-shot learning.
- At larger scale, the biggest sparse model outperforms its dense counterpart even though the dense model requires twice as much computation.The results support sparse MoE models as an alternative that saves computation and reduces model energy consumption.
Ethical considerations
The paper examines bias, stereotyping, energy use, and environmental impact associated with scaling language models, while emphasizing compute-efficient MoEs and transparency practices.
- Bias and toxicity: Bias and stereotyping are comparable between dense and MoE models, especially at scale, while larger models generally show worse performance.The authors identify further mitigation research as necessary and suggest that designed experts might offer greater controllability.
- Energy and environmental impact: MoEs may reduce the energy and environmental costs associated with future language-model scaling by requiring less computation.The paper discusses training energy use separately and frames compute efficiency as a potential way to alleviate future environmental impact.
- Transparency and reproducibility: The authors release code, pretrained models, data cards, and model cards to support research, replication, accountability, and more efficient exploration.These resources are presented as reducing the exploration burden for researchers and the environment.
Limitations
The study evaluates only one MoE configuration, leaving open whether alternative architectures or MoE-specific hyperparameters would scale better and why task advantages vary.
- Configuration scope: The study uses a single MoE configuration with dense and sparse models sharing GPT-3-like hyperparameters and model structure.This configuration may be suboptimal for scaling MoE models.
- Configuration scope: The experiments do not vary MoE-specific choices such as the number of experts.The authors therefore do not establish how these design choices affect scaling performance.
- Open questions: The factors making some tasks and domains more favorable for MoEs remain unclear despite observed variation in performance gaps.This limits the study’s ability to explain when MoEs are most advantageous.
A Full perplexity results
This section reports full perplexity results across the Pile subsets and describes the fine-tuning procedure used elsewhere in the evaluation.
- Full perplexity results: Table 5 reports complete perplexity results for all different subsets of the Pile.The passage identifies the table’s coverage but does not provide individual perplexity values.
- Fine-tuning procedure: Fine-tuning uses fixed epoch counts by dataset and selects models according to validation-set accuracy.The schedules range from 6 to 100 epochs, depending on the task.
- Fine-tuning procedure: For datasets lacking validation sets, the training data are randomly divided into 80% fine-tuning and 20% per-epoch validation subsets.The same split is used when evaluation occurs on the validation set.
C Understanding Potential Harms
The paper evaluates potential harms using StereoSet and CrowS-Pairs, finding that bias generally worsens with scale and that MoE–dense differences are usually small or nonsignificant.
- Evaluation benchmarks: The models are evaluated on StereoSet and CrowS-Pairs to examine whether sparse models behave differently from dense models regarding bias.StereoSet measures bias across multiple domains, while the reported metrics include LMS, SS, and ICAT.
- Evaluation metrics: StereoSet defines LMS as meaningful-association preference, SS as stereotypical-association preference, and ICAT as a combined metric where higher values are better.For SS, scores near 50 indicate better performance.
- Interpretation caveat: The two bias benchmarks have limitations, including unclear explanations of how some biases are measured, so their results require careful interpretation.This caveat applies to conclusions drawn from both benchmarks.
- Results by scale: Both dense and MoE models generally obtain worse CrowS-Pairs scores as scale increases, with a statistically significant difference between the best and worst scores.The comparison uses the Stereotype Score metric and a bootstrap test with p < 0.05.
- Dense versus MoE models: On StereoSet, comparable-FLOPs dense and sparse models perform comparably, while MoEs are slightly less biased on CrowS-Pairs without a statistically significant difference.Both model types also generally become more biased with increasing scale.
D CO2 Emission Related to Experiments
The paper estimates training time, energy use, and carbon emissions for large dense and MoE models, while noting important omissions and infrastructure overheads. Pilot experimentation may substantially increase total training costs beyond the reported estimates.
- Emission estimates: Each GPU-day of training is estimated to produce 3.24 kgCO2e under the stated Azure A100 and regional carbon-efficiency assumptions.The estimate assumes a 400W GPU, carbon efficiency of 0.3 kgCO2e/kWh, and PUE of 1.125.
- Cost estimation: Training-time estimates use 160 TFLOP/s per A100 GPU for dense models and 115 TFLOP/s for MoE models.The rates are based on observed training speeds for the largest models.
- MoE overhead: MoE models have lower GPU utilization than dense models because all-to-all communication introduces additional overhead.The paper notes that implementation optimization could reduce this overhead.
- Scope and omissions: The estimates omit infrastructure-manufacturing costs and exploratory pilot experiments.These exclusions limit the estimates’ coverage of total research costs.
- Scope and omissions: Pilot experimentation is estimated to multiply total training cost by 2, because exploration and tuning commonly add substantial training runs.The authors trained and discarded pilot 6.7B dense and 1.1T MoE models during the project’s early stages.
E Knowledge Distillation
The paper investigates knowledge distillation as a way to transfer capabilities from large dense or MoE teachers into small dense students. Distillation improves the dense baseline and can transmit some sparse-training efficiency advantages to the student.
- Approach: Knowledge distillation trains small dense students to mimic larger dense or MoE teacher models.The approach is motivated by combining dense-model and sparse-model benefits.
- Training setup: Students use a weighted objective combining 25% standard cross entropy with 75% soft distillation loss.The distillation loss encourages students to reproduce the teacher’s logits; sequence length is reduced to 1024 tokens for faster experimentation.
- Results: Distilled students outperform a well-tuned dense baseline for both dense and sparse teacher models.The results are reported in Table 9, whose PPL metric is in-domain validation perplexity.
- Results: Students distilled from a 52B-parameter MoE teacher outperform students distilled from a 1.3B-parameter dense teacher.The dense teacher is twice as costly to train, while the MoE teacher’s advantages are partly transmitted to the dense student.
F Techniques for Large-scale Training
The paper combines memory-saving and distributed-training techniques to make very large models feasible to train, then analytically counts their training FLOPs. MoE FLOP estimates include extra feed-forward computation but are independent of the number of experts.
- Scaling techniques: Pure FP16 training, activation checkpointing, and fully sharded data parallelism make models beyond single-GPU memory feasible to train.These techniques are used because parameter storage for many models exceeds usable single-GPU memory.
- FP16 training: Maintaining only 16-bit weights, optimizer states, and gradients reduces memory requirements by 50% without reducing model accuracy.The method discards 32-bit weights and dynamically rescales Adam states stored in 16-bit.
- Activation checkpointing: Activation checkpointing saves memory by discarding selected forward-pass activations and recomputing them during backpropagation.The technique increases computation by 33%.
- Fully Sharded Data Parallel: FSDP shards model parameters, optimizer state, and gradients across workers, enabling training of models with more than one trillion parameters.The paper wraps every transformer layer to balance communication efficiency against peak resident memory.
- FLOP accounting: Training FLOPs are counted analytically with an additional forward pass assumed for activation checkpointing.Dense-model FLOPs depend on training tokens, layers, hidden dimension, sequence length, and vocabulary size.
- FLOP accounting: MoE FLOPs add a feed-forward network at every other layer for top-2 routing, while routing-projection FLOPs are treated as negligible.The resulting quantity is independent of the number of experts.