Source-linked AI summary

A Review of Sparse Expert Models in Deep Learning

William Fedus, Jeff Dean, Barret Zoph

arXiv:2209.01667v1cs.LGcs.CL

TL;DR

Large deep-learning models have become expensive to train and serve, motivating architectures that increase parameter capacity without proportionally increasing computation per example. This paper reviews sparse expert models, their routing and systems trade-offs, and their reported gains across tasks and domains. It concludes that sparsity can produce massive models with better accuracy than dense counterparts, while important design and generalization questions remain.

  • Problem

    Scaling modern neural models requires extensive accelerator resources, making training expensive and energy-intensive.

  • Method

    The paper surveys sparse expert models in deep learning, describing common routing algorithms, contextualizing recent advances, and identifying future research directions.

  • Results

    Sparse expert models have achieved significant gains across domains; GLaM outperformed GPT-3 in zero- and one-shot performance while using 49% fewer inference FLOPs per token and 65% lower power.

  • Takeaways & Limitations

    Sparsity can yield massive models with better accuracy than dense counterparts while keeping training and inference costs lower.

  • Takeaways & Limitations

    Sparse expert systems incur additional distributed communication overhead, and their optimal expert configuration and out-of-domain generalization remain incompletely understood.

Abstract

from arXiv · show

Sparse expert models are a thirty-year old concept re-emerging as a popular architecture in deep learning. This class of architecture encompasses Mixture-of-Experts, Switch Transformers, Routing Networks, BASE layers, and others, all with the unifying idea that each example is acted on by a subset of the parameters. By doing so, the degree of sparsity decouples the parameter count from the compute per example allowing for extremely large, but efficient models. The resulting models have demonstrated significant improvements across diverse domains such as natural language processing, computer vision, and speech recognition. We review the concept of sparse expert models, provide a basic description of the common algorithms, contextualize the advances in the deep learning era, and conclude by highlighting areas for future work.

1 INTRODUCTION

Sparse expert models address the rising cost of scaling deep learning by routing each example through only part of a larger parameter set. This survey reviews their deep-learning-era development and applications across multiple domains.

  • Motivation: Scaling computation, data, and model size has improved machine learning but now requires costly, energy-intensive accelerator systems.State-of-the-art models may use thousands of interconnected accelerators for weeks or months.
  • Core concept: Sparse expert models partition parameters into experts with distinct weights and route each example to specific experts.Each example interacts with only a subset of the network parameters.
  • Core concept: Using only a fraction of experts can keep per-example computation small relative to total model size.This separates the amount of computation from the full network’s parameter count.
  • Applications: Sparse expert models have been used successfully in natural language processing, computer vision, speech recognition, and multimodal learning.Their popularity increased alongside Transformer language models.
  • Scope: The survey focuses on sparse expert models from the deep-learning era, approximately 2012 onward, while directing readers to broader historical and efficiency surveys.It recounts recent advances and discusses future research avenues.

2 SPARSE EXPERT MODELS

Sparse expert models evolved from early Mixture-of-Experts systems into routed neural components that select a small subset of experts for each input. Their efficiency depends not only on routing, but also on distributed execution, load balancing, and communication costs.

  • History: Mixture-of-Experts originated at least three decades ago, when experts were entire neural networks and the architecture resembled an ensemble.Later work treated MoE as a component within a larger neural network.
  • Routing: Top-k routing computes router logits from token representations, normalizes them across experts, and selects the top-k experts for each token.The selected experts’ computations are combined using their gate values.
  • Routing: Top-k selection is more computationally efficient than continuous soft selection because it restricts computation to a subset of experts.Soft selection incurs the full computational cost even for experts with very small routing weights.
  • Distributed execution: Sparse expert models fit distributed parallelism by placing experts on different accelerators and dynamically dispatching inputs among them.This complements data, tensor-model, and pipeline parallelism strategies.
  • Hardware trade-offs: Dynamic routing introduces all2all communication overhead, while larger capacity factors can improve quality but increase communication, memory, and compute costs.Load-balancing methods encourage better accelerator utilization when routed inputs are unevenly distributed.
  • Systems advances: Recent systems work has improved sparse expert training and deployment, including reported 37× inference speedups for a special case of unbatched inference.Other work dynamically recompiles computations to optimize resource use.

3 SCALING PROPERTIES OF SPARSE EXPERT MODELS

Sparse expert scaling can improve efficiency and performance across upstream training and downstream inference, but gains depend on scale, transfer setting, expert configuration, and hardware. The evidence also shows important scope limits, including reduced transfer performance and analyses tied to specific training-token budgets or expert-layer designs.

  • 3.1 UPSTREAM SCALING: Sparse models achieved 4–7× wall-time speed-ups over T5 using the same compute resources, while upstream gains diminished with 256+ experts.Sparse scaling was also associated with a 13.5 BLEU gain on a 600B-parameter translation model.
  • 3.1 UPSTREAM SCALING: In-domain language modeling scaled significantly better for MoE models than for dense models, while out-of-domain transfer remained more difficult.This finding corroborated earlier concerns about transferring upstream improvements to different distributions.
  • 3.1 UPSTREAM SCALING: The benefit of sparsity decreased with scale and, under an analysis using 130B training tokens, extrapolated to no further benefit beyond 900B parameters or FLOPs.The authors note that this conclusion should be revisited using token budgets closer to compute-optimal training recommendations.
  • 3.2 DOWNSTREAM SCALING: Sparse models improved few-shot performance over dense counterparts, including GLaM outperforming GPT-3 with 49% fewer inference FLOPs per token and 65% lower power.BIG-Bench measured a 2× improvement on 161 contributed JSON tasks for sparse over dense models.
  • 3.2 DOWNSTREAM SCALING: Sparse models achieved calibration roughly matching a dense model using 10× more FLOPs, although Expected Calibration Error improved with scale for both model types.The calibration comparison was reported on multiple-choice BIG-Bench tasks.
  • 3.3 SCALING THE NUMBER, SIZE AND FREQUENCY OF EXPERT LAYERS: Expert count, expert size, and expert-layer frequency materially affect scaling, with fewer larger experts favored for transfer and hardware-specific sharding needed for efficient execution.The optimal configuration depends on the application and hardware; the discussed trade-offs concern experts-as-a-layer designs, not independently trained full-model experts.

4 ROUTING ALGORITHMS

Routing algorithms determine which experts process each input, using token-wise, expert-wise, or globally assigned selections. The survey also covers static, reinforcement-learning, and load-balancing approaches alongside these main categories.

  • Routing taxonomy: Routing algorithms determine where each example is sent, but discrete expert selection is typically non-differentiable.Some work recasts expert selection as a bandit problem and uses reinforcement learning.
  • Routing taxonomy: The main taxonomy includes token-chosen top-k experts, expert-chosen top-k tokens, and globally determined expert assignments.The categories correspond to choosing along the Experts axis, choosing along the Tokens axis, and solving assignments without a greedy approach.
  • Token-chosen routing: Top-k routing sends each token to selected experts, including the original top-2 formulation and successful top-1 variants.A reinforcement-learning variant uses REINFORCE with negative cross entropy as the reward instead of scaling expert outputs by router probability.
  • Token-chosen routing: Extensions include expert prototyping, annealing from soft gating to top-1 routing, smooth DSelect-k routing, and shared dense computation with one selected expert.These methods modify selection structure or training dynamics while retaining sparse expert computation.
  • Expert-chosen routing: Expert-chosen routing lets each expert select its top-k tokens, removing auxiliary load-balancing losses and ensuring equal expert token counts.Some tokens may reach no expert or all experts under this design.
  • Global assignment: BASE layers formulate routing as a linear assignment problem that maximizes routing scores while assigning a fixed number of tokens to each expert.Distributed implementations may shuffle tokens before solving assignments locally, introducing additional communication.

5 SPARSE EXPERT MODELS ACROSS DOMAINS

Sparse expert models originated in language modeling and translation, then spread across vision, speech, and multimodal learning. Across domains, the sparse architecture remains similar while the routed input changes, with reported quality and efficiency gains.

  • Cross-domain overview: Sparse expert models expanded from natural language processing into computer vision, speech recognition, and multimodal applications.The underlying sparse architectures and algorithms remain roughly similar across domains, while the routed inputs differ.
  • Natural language processing: In early language-modeling work, experts were inserted between LSTM layers; Transformer-based models typically replace dense layers with expert layers.Lepikhin et al. introduced MoE layers into Transformers for machine translation.
  • Natural language processing: 2048 experts per expert layer produced state-of-the-art translation results across 100 languages, while a sparse 1.6T-parameter model achieved state-of-the-art pre-training quality.Sparse self-attention q/k/v activations were less stable in the reported study.
  • Natural language processing: ST-MoE achieved state-of-the-art results across reasoning and generation tasks while using roughly 20× less pre-training FLOPs and 40× less inference FLOPs than PaLM-540B when fine-tuned on SuperGLUE.Another MoE decoder-only model achieved state-of-the-art few-shot results using one-third of GPT-3’s training compute.
  • Computer vision: V-MoE matched prior state-of-the-art image-classification performance using half the inference compute, while another sparse MoE MLP improved ImageNet and CIFAR performance over its dense counterpart.These approaches embed sparse experts into vision architectures such as ViT and MLP-Mixer.
  • Computer vision: Residual Mixture-of-Expert layers reduced training cost by 30% with comparable segmentation and object-detection quality, while another system achieved 1.5x-2x speedups in training and inference.The speedups were measured against a previous MoE implementation in a Swin Transformer V2 system.
  • Speech and multimodal learning: SpeechMoE models improved character error rates across four speech-recognition datasets through auxiliary losses and new routing architectures.LIMoE routed image patches and word tokens and outperformed CLIP under a comparable training strategy.

6 WHEN TO USE A SPARSE VERSUS DENSE MODEL

Sparse models are attractive when distributed hardware can host their expanded parameters while keeping per-example computation low, but downstream memory and deployment constraints can favor dense models.

  • Sparse models increase parameter count by adding experts while keeping FLOPs per example approximately constant.A single-expert sparse model is roughly equivalent to a dense model.
  • Sparse models fit training or serving setups with many machines that can host additional parameters in parallel.They are especially suitable for data-parallel training and high-throughput serving.
  • The number of experts should match downstream memory when pre-training has more machines than fine-tuning or serving.This is a practical design consideration for transferring models across deployment stages.
  • For a fixed hardware size, sparse models can be worse than dense models on a per-parameter basis.This comparison assumes all parameters remain in accelerator memory.
  • Sparse models can remain practical in memory-restrictive settings with two experts or dynamic CPU-GPU memory swapping.Other parameter-footprint reduction methods are also discussed as ways to ease memory requirements.

7 SPARSE MODEL TRAINING IMPROVEMENTS

Sparse models require specialized training and transfer strategies because they can be unstable and can lag dense models after adaptation to new distributions.

  • Sparse models often have different training dynamics from dense models and are frequently reported to be more unstable at larger scales.Reported instabilities include diverging loss and failures in high-scale models.
  • Higher precision, lower initialization scales, selective router precision, gradient scaling, and router z-loss have been used to improve stability.Higher precision can increase memory use and slow training, while router z-loss reduces floating-point roundoff by keeping router logits small.
  • 7.2 TRANSFER TO NEW DISTRIBUTIONS: For a given pre-training perplexity, sparse models fine-tune worse on reasoning tasks but better on knowledge-heavy tasks.Other work also reports worse out-of-domain language modeling and lower fine-tuning performance on HellaSwag, PIQA, and Winogrande.
  • 7.2 TRANSFER TO NEW DISTRIBUTIONS: Using fewer experts with more FLOPs can improve fine-tuning performance compared with maximizing sparsity.A 128-expert model matched the FLOPs of an 11B dense model, while a 2048-expert model matched a 2B dense model.
  • 7.2 TRANSFER TO NEW DISTRIBUTIONS: Sparse models need independent fine-tuning hyperparameter studies because optimal learning rates and batch sizes can differ dramatically from dense models.Applying dense-model hyperparameters to sparse models can mask pre-training improvements.
  • Task-level routing can reduce inference costs for machine translation by selecting weights at the task rather than token level.This approach reduces the inference-time parameter subset needed for each task.

8 INTERPRETABILITY

Sparse expert models support interpretability by exposing discrete routing choices, revealing specialization across textual and visual inputs while leaving contextual mechanisms underexplored.

  • Sparse models make interpretation easier by identifying the discrete experts selected for each input.This replaces analysis of potentially trillions of weights with a small set of expert indices.
  • BASE-layer analyses found experts specializing in punctuation, conjunctions and articles, verbs, visual descriptions, proper names, and counting.The table reports the most frequent preceding top-five tokens for selected experts.
  • ST-MoE analyses similarly identified specialization in punctuation, articles, verbs, visual descriptions, proper names, and numbers.The reproduced table presents encoder expert specializations.
  • Full encoder-decoder ST-MoE analysis found clearer specialization in the encoder than the decoder.The absence of evident specialization may reflect difficult-to-discern patterns or no useful patterns.
  • LIMoE experts specialized across textual and visual data, including textures, plants, eyes, and words.Visual examples included textures, natural objects, and man-made objects such as wheels and door handles.
  • Token- and patch-based interpretability heuristics may miss nuanced specialization because embeddings incorporate surrounding context.Transformers provide this contextual information through self-attention or encoder-decoder attention.

9 FUTURE DIRECTIONS AND CONCLUSIONS

The review points toward combining sparse experts with adaptive computation and retrieval while identifying unresolved questions about architecture, granularity, generalization, and expert design.

  • LIMoE visual experts specialize in textures, natural objects, and man-made objects, illustrating multimodal specialization.Examples include plants, hands, eyes, wheels, door handles, and words.
  • Adaptive Computation: Sparse expert models and adaptive computation can be combined by varying expert count, layer usage, or expert architecture across inputs.Heterogeneous experts could differ in depth or width and therefore incur different computation.
  • Retrieval Methods: Retrieval methods and sparse experts both increase model capacity, but retrieval accesses external information while experts store additional parametric knowledge.The passage frames these approaches as having an overlapping capacity goal.
  • Conclusions: Open questions include how task demands determine expert number and size, how to improve out-of-domain generalization, and how to diversify architectures.The review also highlights unresolved choices about sparsity granularity and modular expert structure.
  • Conclusions: The review concludes that sparsity reduces training and inference costs while producing massive models with better accuracy than dense counterparts.It also reports significant gains alongside continuing engineering and research problems.
Loading 2209.01667v1…