PaperlayerView exact revision on arXiv

Source-linked AI summary

ST-MoE: Designing Stable and Transferable Sparse Expert Models

Barret Zoph, Irwan Bello, Sameer Kumar, Nan Du, Yanping Huang, Jeff Dean, Noam Shazeer, William Fedus

arXiv:2202.08906v2cs.CLcs.LG

TL;DR

Sparse models face training instability and uncertain fine-tuning quality despite their efficient scaling. This paper studies stability and transferability, introduces design guidance including the router z-loss, and scales ST-MoE-32B to 269B parameters. The resulting sparse model achieves state-of-the-art performance across diverse NLP benchmarks, while the authors identify limitations on some tasks and multilingual training.

  • Problem

    Sparse models have shown pre-training speedups but can lag smaller models after fine-tuning, while larger-scale pre-training is hindered by instability.

  • Method

    The paper studies stability–quality trade-offs, introduces the router z-loss, analyzes fine-tuning and routing, and jointly optimizes pre-training and fine-tuning metrics.

  • Results

    ST-MoE-32B, a 269B-parameter sparse model, achieves state-of-the-art performance across a diverse set of natural-language benchmarks.

  • Takeaways & Limitations

    The findings provide design guidance for practical sparse models and support their adoption across competitive public benchmarks.

  • Takeaways & Limitations

    The model fails to improve on some tasks, including CB and WSC, and does not surpass GPT-3 or GLAM on closed-book TriviaQA.

Abstract

from arXiv · show

Scale has opened new frontiers in natural language processing -- but at a high cost. In response, Mixture-of-Experts (MoE) and Switch Transformers have been proposed as an energy efficient path to even larger and more capable language models. But advancing the state-of-the-art across a broad set of natural language tasks has been hindered by training instabilities and uncertain quality during fine-tuning. Our work focuses on these issues and acts as a design guide. We conclude by scaling a sparse model to 269B parameters, with a computational cost comparable to a 32B dense encoder-decoder Transformer (Stable and Transferable Mixture-of-Experts or ST-MoE-32B). For the first time, a sparse model achieves state-of-the-art performance in transfer learning, across a diverse set of tasks including reasoning (SuperGLUE, ARC Easy, ARC Challenge), summarization (XSum, CNN-DM), closed book question answering (WebQA, Natural Questions), and adversarially constructed tasks (Winogrande, ANLI R3).

1 INTRODUCTION

The paper targets sparse-model training instability and unreliable fine-tuning by studying stability–quality trade-offs, introducing design guidance, and evaluating a 269B-parameter model across diverse NLP benchmarks.

  • Sparse models can expand parameter counts while keeping FLOPs per token roughly constant, but prior models lagged smaller systems after fine-tuning.Training instabilities also hampered necessary pre-training at larger scales.
  • The study jointly examines upstream pre-training and downstream fine-tuning to improve sparse-model practicality and reliability.It includes a large-scale study of stability techniques and their quality trade-offs.
  • The router z-loss resolves instability issues while slightly improving model quality.The paper recommends this auxiliary loss after comparing several stability approaches.
  • Fine-tuning quality depends differently on batch size and learning rate for sparse and dense models.Bad hyperparameters can yield virtually no fine-tuning gain over dense models despite large pre-training speed-ups.
  • The work develops architectural, routing, and model-design principles plus qualitative analyses of token routing across expert layers.These analyses support designing Pareto-efficient sparse models in distributed settings.
  • The 269B-parameter ST-MoE-32B achieves state-of-the-art performance across a diverse set of natural-language benchmarks.

2 BACKGROUND

Sparse expert models replace a shared neural-network layer with routed experts, processing each input through only a subset. Routing and capacity choices balance computation, load distribution, memory, and communication costs.

  • Sparse expert layers contain multiple experts with unique weights, and a router sends each input to a selected subset.Experts are typically homogeneous, though heterogeneous expert types are possible.
  • A token representation x is routed to the best-matched top-k experts from a set of N experts.The router computes logits with Wr · x and normalizes them using a softmax over experts.
  • The batch is divided into G groups across the data-parallel dimension, with each group containing B/G tokens.Expert capacity is determined using the capacity factor, token count, and number of experts.
  • Increasing the capacity factor reduces token dropping from load imbalance but increases memory and computational costs.The optimal capacity factor and routing choice depend on the hardware-software system’s communication speeds.
  • An auxiliary load-balancing loss encourages roughly even token distribution across experts, improving hardware efficiency.Alternative approaches formulate balanced allocation as an assignment problem and remove the auxiliary loss.

3 STABILIZING TRAINING OF SPARSE MODELS

Sparse models often suffer training instabilities, and many stabilizing changes trade improved stability for worse quality. The router z-loss instead stabilizes training without quality degradation, while multiplicative components improve quality but can worsen stability.

  • Sparse models often suffer training instabilities worse than standard densely activated Transformers.
  • Many stability techniques improve training stability at the expense of model quality.The study evaluates generic and sparse-specific approaches, including noise injection, activation and gradient constraints, and architectural changes.
  • 3.1 Stability and quality tradeoffs when removing multiplicative interactions: Multiplicative interactions improve quality but can destabilize training, with added expert-layer interactions improving quality with virtually no slowdown in model step time.
  • 3.1 Stability and quality tradeoffs when removing multiplicative interactions: Removing GEGLU layers or RMS scale parameters improves stability but significantly worsens model quality.The GEGLU replacement matches FLOPs and parameters, isolating the effect of the multiplicative component.
  • 3.3 Stability and quality tradeoffs when constraining activations and gradients: The router z-loss stabilizes all three runs without quality degradation, whereas tighter update clipping significantly hurts quality.The router z-loss penalizes large logits entering the gating network and is selected for its combined stability and quality.
  • 3.4 Selecting a precision format: trading efficiency and stability: Sparse models are sensitive to roundoff errors because router exponential functions can amplify perturbations and alter second-expert routing.The router z-loss encourages smaller logits, reducing roundoff errors before exponential functions.

4 FINE-TUNING PERFORMANCE OF SPARSE MODELS

Sparse models can optimize and converge quickly during fine-tuning, yet their generalization depends strongly on task size, regularization, parameter updates, and protocol choices. The paper studies these sensitivities and identifies practices that improve or preserve fine-tuning quality.

  • 4.1 HYPOTHESIS: A GENERALIZATION PROBLEM: Sparse models reach 100% training accuracy faster than dense models, but outperform them on ReCoRD validation while lagging on smaller CB validation.CB uses 250 training sequences, whereas ReCoRD uses 138k.
  • 4.1 HYPOTHESIS: A GENERALIZATION PROBLEM: A global dropout rate of 0.1 maximizes SuperGLUE fine-tuning quality, while higher rates over-regularize and expert dropout provides further generalization benefits.The regularization study averages three random seeds to convergence.
  • 4.2 FINE-TUNING SUBSETS: Updating non-MoE parameters performs about as well as updating all parameters, while updating only MoE parameters significantly degrades fine-tuning performance.Updating only non-MoE parameters can also reduce fine-tuning memory and computation.
  • 4.3 HYPERPARAMETER SENSITIVITY: The best fine-tuning batch size and learning rate differ between sparse and dense models, so dense-model hyperparameters can conceal sparse-model pre-training gains.The experiments pre-train both models on 500B C4 tokens and fine-tune them on SuperGLUE.
  • 4.5 PRE-TRAINING AND FINE-TUNING OBJECTIVES: Making fine-tuning resemble span-corruption pre-training accelerates convergence but does not consistently improve held-out quality across tasks.Sentinel tokens improve GEC but not SuperGLUE.

5 DESIGNING SPARSE MODELS

The paper presents hardware- and scale-aware principles for selecting sparse-model routing, capacity factors, and architecture. Quality generally improves with capacity and routing choices, but efficiency trade-offs make the best configuration system-dependent.

  • 5 DESIGNING SPARSE MODELS: The paper recommends top-2 routing with a 1.25 capacity factor and at most one expert per core.Capacity factors can also be changed during evaluation to match available memory and compute.
  • 5.1 NUMBER OF EXPERTS: Using more than one expert per core increases memory transfer and lowers compute-to-memory ratio without changing computation per input.The paper therefore uses at most one expert per data-parallelism row on its TPU system.
  • 5.2 CHOOSING THE CAPACITY FACTOR AND ROUTING ALGORITHM: Increasing capacity factors improves quality but also increases einsum, activation-memory, all2all, and allreduce costs.A top-1 increase from 1.0 to 1.25 gives +0.011 neg. log perp., while top-2 from 1.25 to 2.0 gives +0.009.
  • 5.2 CHOOSING THE CAPACITY FACTOR AND ROUTING ALGORITHM: Top-(n+1) routing yields only small gains over top-n routing at fixed capacity factor, including +0.004 for top-2 over top-1 at train CF 1.25.The reported gain is about 1/20th of the boost from tripling a dense model.
  • 5.2 CHOOSING THE CAPACITY FACTOR AND ROUTING ALGORITHM: The hardware-software stack determines the Pareto-efficient routing and capacity configuration because higher capacity factors slow larger models.On TPUs, raising train capacity factor from 1.25 to 2.0 increases step time by +7% for the 1B model and +14% for the 32B model.
  • 5.2 CHOOSING THE CAPACITY FACTOR AND ROUTING ALGORITHM: Batch Prioritized Routing significantly improves performance below capacity factor one, motivating its use when larger models favor lower-capacity configurations.The paper links this regime to more expensive all2all and allreduce communication.

6 EXPERIMENTAL RESULTS

The experiments validate sparse models at T5-Large scale and culminate in ST-MoE-32B, a 269B-parameter model FLOP-matched to a dense 32B Transformer that achieves broad state-of-the-art results. Fine-tuning outcomes depend strongly on task size and hyperparameter selection, with several documented weaknesses.

  • 6.1 FLOP-Matched Models: A FLOP-matched sparse model consistently improves over dense T5-Large across natural language understanding, question answering, and summarization tasks.The comparison uses 500k pre-training steps, a 1M batch size, and 524B C4 tokens.
  • 6.2 ST-MoE-32B: ST-MoE-32B contains 269B parameters while matching the computational cost of a dense 32B Transformer.Its reduced parameter count is intended to ease serving and fine-tuning.
  • 6.2 ST-MoE-32B: 91.2 SuperGLUE test-server average establishes state-of-the-art performance, while XSum and CNN-DM also reach state-of-the-art without training changes.The SuperGLUE validation accuracy is 93.2.
  • 6.2 ST-MoE-32B: 94.8 on ARC Easy versus 92.7 and 86.5 on ARC Challenge versus 81.4 improve the current test-server state of the art.The reported comparisons are test-server submissions.
  • 6.2 ST-MoE-32B: Closed-book WebQA reaches 47.4 accuracy versus 42.8 previously, while Natural Questions reaches 41.9 versus 41.5.The model improves over prior state of the art on two of three closed-book question-answering tasks.
  • 6.2 ST-MoE-32B: The sparse model scores 90.8 exact match on SQuAD, below T5-XXL's 91.3, and does not improve on the smallest tasks CB and WSC.Trivia QA also fails to improve over GPT-3 and GLAM.

7 TRACING TOKENS THROUGH THE MODEL

Token-routing analysis finds clear encoder expert specialization but much weaker decoder specialization. Multilingual encoders specialize by token patterns rather than by language, while routing behavior remains shaped by group-wise load balancing.

  • Routing Procedure: Tokens are routed in groups, with auxiliary load balancing encouraging even expert utilization and limiting dispatch and gathering costs.Experts compete for assignments within each group rather than across the entire batch.
  • 7.1 Encoder Experts: Encoder experts specialize in sentinel tokens and patterns including punctuation, verbs, proper names, and counting.The examples are summarized in Table 13.
  • 7.2 Decoder Experts: Decoder experts show little meaningful specialization, with sentinel tokens routed more uniformly than in encoder layers.Encoder routing entropy is low, whereas decoder routing entropy is high and nearly uniform across 32 experts.
  • 7.3 Multilingual Experts: Multilingual encoder experts specialize in sentinel tokens, numbers, conjunctions, articles, and proper names, sometimes across languages.Examples include related tokens such as “for” and “pour.”
  • 7.3 Multilingual Experts: The multilingual experts show no evidence of language specialization: English, Japanese, French, and Chinese tokens are routed indiscriminately.The authors connect this pattern to groups containing only a few languages and group-wise load balancing.
  • 7.4 Open Questions: The observed encoder specialization leaves open how architectures without learned routing achieve strong performance.The paper proposes studying learned versus random routing at larger scale.

8 RELATED WORK

Related work positions sparse expert models as an efficient route to scale across language and other modalities, while highlighting advances in routing, deployment, and multitask modularity.

  • Mixture-of-Experts History: Mixture-of-Experts evolved from whole-network ensemble-like models into components of deeper networks and later scaled to large machine-translation systems.The cited history spans Jacobs and Jordan’s early work, Eigen’s deep-network formulation, and Shazeer’s 137B model.
  • Scale in NLP: Sparse models target dense-model quality more efficiently, with prior work reporting a 4x pre-training speed-up and GPT-3-level quality using one-third the energy.These results are attributed to Fedus et al. and Du et al.
  • Routing Algorithms: BASE layers and related methods recast routing as assignment or optimal-transport problems, removing or changing auxiliary load-balancing losses.Prior work also examined the scaling properties of alternative routing algorithms and single expert layers.
  • Other Modalities: MoE methods have improved results in vision and speech, including a 15B V-MoE model matching ImageNet state of the art with fewer computational resources.Other work applies MoE layers across image patches and channels or to automatic speech recognition.
  • Deployment: Deployment research addresses token-level serving burdens through distillation and routing at the sentence or task level.The motivation is the memory required to store expert parameters.
  • Multitask Learning: Multitask MoE research explores task-specific routers and conditioning expert activation on domain or task labels.Some approaches instead infer the relevant label.

9 DISCUSSION

The discussion presents sparse models as robust across many design choices but identifies important boundaries: multilingual pre-training can be unstable, small-scale findings may not transfer, and some architectural changes require stabilization.

  • Unpredictable Multilingual Dynamics: Multilingual pre-training can produce smaller speed-ups and unstable routing because groups contain only a few languages despite training on more than 100.The authors hypothesize that variance across groups and batches drives chaotic routing.
  • Robustness: Sparse models remain robust across routing algorithms, token dropping rates, architectural changes, and many hyperparameters.Fine-tuning batch size and learning rate remain important exceptions.
  • Scaling Findings: Improvements found at small scale can vanish or reverse in larger or longer-trained models, including earlier top-n routing findings.The paper reports that 8x larger experiments favored small boosts from top-(n + 1) routing over top-n routing.
  • Lower-Precision Training: The router z-loss stabilizes training by encouraging smaller-magnitude router logits and can preserve or improve quality.The authors suggest this may inform lower-precision formats with more compressed exponential ranges.
  • Multiplicative Operations: Adding multiplicative interactions to expert layers speeds pre-training by 4% without changing step time, while router z-loss prevents additional instability.The paper links the efficiency to low communication overhead under model parallelism.
  • Scaling Dynamics: Encoder-decoder models and deeper models are more unstable under fixed FLOPs, motivating activation constraints and possible attention z-loss penalties.The paper attributes encoder-decoder instability partly to their higher ratio of attention layers.

10 CONCLUSION

The work argues that better-balanced computation makes sparse models more practical and usable while maintaining strong benchmark performance.

  • Sparse models achieve state-of-the-art performance across a wide range of competitive public benchmarks.
  • A sparse model with 1/5th the size and a better FLOPs-to-parameters balance improves usability through lower memory overhead.
  • The auxiliary load-balancing loss encourages uniform token routing across experts by minimizing when routing fractions and router probabilities approach 1/N.
  • The auxiliary-loss coefficient α = 10−2 balances load effectively without overwhelming the primary cross-entropy objective.

B ROUTER Z-LOSS TRAINING DYNAMICS

The router z-loss coefficient controls router-logit dynamics during pre-training, with the selected coefficient producing stable behavior in both encoder and decoder layers.

  • cz = 0.001 was selected as the best router z-loss coefficient for both encoder and decoder layers.
  • With cz values of 1e-2 or larger, router z-loss shrinks near zero during pre-training.
  • Pi denotes the fraction of router probability allocated to expert i across all tokens in batch B.

C IMPROVED ARCHITECTURAL MODIFICATIONS

The paper evaluates architectural, routing, data, and fine-tuning choices for improving sparse-model quality, convergence, and efficiency, while identifying setting-specific trade-offs.

  • IMPROVED ARCHITECTURAL MODIFICATIONS: A dense FFN immediately before or after each sparse layer improves quality more than placing the same added layers randomly elsewhere.
  • IMPROVED ARCHITECTURAL MODIFICATIONS: All models use GELU-Linear FFNs rather than ReLU FFNs, with additive or multiplicative learned biases applied after the first matrix multiplication.
  • IMPROVED ARCHITECTURAL MODIFICATIONS: Multiplicative expert biases improve performance and achieve a 4% convergence-time speedup over a strong sparse baseline.
  • IMPROVED ARCHITECTURAL MODIFICATIONS: Using word embeddings alone for routing hurts quality, while adding them to layer activations becomes neutral after extended pre-training.
  • IMPROVED ARCHITECTURAL MODIFICATIONS: Batch Prioritized Routing is restricted to encoder layers because decoder use would expose future-token information.
  • IMPROVED ARCHITECTURAL MODIFICATIONS: Batch Prioritized Routing improves quality at capacity factors ≤1 but loses its advantage at 1.25 and underperforms top-2 routing.
  • IMPROVED ARCHITECTURAL MODIFICATIONS: The training data combines C4 with the dataset introduced in GLaM, using mixture weights to sample dataset sources.
  • IMPROVED ARCHITECTURAL MODIFICATIONS: Top-3 routing benefits slightly from lower thresholds than top-2 routing, while tuned absolute thresholds can match performance.

J NEGATIVE RESULTS

Several proposed changes did not improve sparse-model performance, including routing features, additional balancing, pre-training noise, and mixed fine-tuning data.

  • Encoding whether tokens were previously routed or dropped made no difference in performance.
  • Adding explicit expert positional information did not improve performance or speed up early convergence.
  • The tested pre-training noise strategies were intended to reduce pre-training and fine-tuning discrepancies.
  • Adding load balancing for lower n−1 experts in top-n routing provided minimal benefits.
  • Mixing pre-training span-corruption data into fine-tuning did not improve performance and increased training loss.
Loading 2202.08906v2…