Source-linked AI summary
Mixture-of-Experts with Expert Choice Routing
Yanqi Zhou, Tao Lei, Hanxiao Liu, Nan Du, Yanping Huang, Vincent Zhao, Andrew Dai, Zhifeng Chen, Quoc Le, James Laudon
TL;DR
Sparsely activated MoE models increase parameter capacity at fixed per-token computation, but routing can imbalance expert loads and under-train experts. This paper proposes expert choice routing, where experts select top-k tokens, enabling variable experts per token and fixed expert bucket sizes. The method achieves more than 2x faster convergence and stronger downstream results, including outperforming T5 11B dense on 7 of 11 tasks at smaller activation cost.
Problem
Conventional MoE routing can imbalance expert loads and under-train experts, while fixed top-k allocation gives every token the same number of experts regardless of importance.
Method
Expert choice routing has each expert select a fixed-capacity top-k set of tokens, allowing each token to receive a variable number of experts and guaranteeing load balance by design.
Results
The method achieves more than 2x training efficiency improvements over GShard and Switch Transformer and outperforms T5 11B dense on 7 of 11 downstream tasks.
Takeaways & Limitations
Expert choice provides strong scaling and downstream performance across the evaluated GLUE and SuperGLUE tasks while flexibly allocating computation across tokens.
Takeaways & Limitations
The current implementation may not immediately apply to autoregressive generation or very small serving batches, so the paper leaves routing improvements for future work.
Abstract
from arXiv · showhide
Sparsely-activated Mixture-of-experts (MoE) models allow the number of parameters to greatly increase while keeping the amount of computation for a given token or a given sample unchanged. However, a poor expert routing strategy (e.g. one resulting in load imbalance) can cause certain experts to be under-trained, leading to an expert being under or over-specialized. Prior work allocates a fixed number of experts to each token using a top-k function regardless of the relative importance of different tokens. To address this, we propose a heterogeneous mixture-of-experts employing an expert choice method. Instead of letting tokens select the top-k experts, we have experts selecting the top-k tokens. As a result, each token can be routed to a variable number of experts and each expert can have a fixed bucket size. We systematically study pre-training speedups using the same computational resources of the Switch Transformer top-1 and GShard top-2 gating of prior work and find that our method improves training convergence time by more than 2x. For the same computational cost, our method demonstrates higher performance in fine-tuning 11 selected tasks in the GLUE and SuperGLUE benchmarks. For a smaller activation cost, our method outperforms the T5 dense model in 7 out of the 11 tasks.
1 Introduction
Scaling model capacity improves quality, but conventional MoE routing can imbalance expert loads and waste capacity. Expert choice routing assigns fixed token capacity to each expert while allowing tokens to receive different numbers of experts, yielding faster convergence and strong downstream performance.
- Motivation: Scaling model capacity, dataset size, and training time improves computer vision and language-model performance, motivating greater attention to training efficiency.Training efficiency is defined as the computation required to achieve superior model quality than the state-of-the-art system.
- Problem: Conventional token-choice routing can imbalance expert loads, under-utilize some experts, and leave some tokens unprocessed.Auxiliary balancing losses have limited effectiveness, while over-utilized experts are constrained by a maximum capacity.
- Method: Expert choice lets each expert select its top-k tokens instead of having each token select top-scoring experts.The method provides fluid parameter allocation based on learned token-to-expert importance and guarantees load balance without an auxiliary loss.
- Results: Over 2× faster training convergence is achieved by the 8B/64E model versus Switch Transformer top-1 and GShard top-2 gating.The comparison uses the corresponding top-1 and top-2 gating counterparts.
- Results: The method shows strong downstream performance across evaluated GLUE and SuperGLUE scales, with the 8B/64E model outperforming T5 11B dense in 7 of 11 tasks.The reported downstream evaluation covers selected tasks from both benchmarks.
2 Related Work
Related work scales neural networks through model and conditional computation, while MoE methods increase capacity by selectively activating experts. Prior routing approaches address balancing or alternative assignments, but often require auxiliary losses or remain limited to pre-training evaluation.
- Scaling: Model parallelism splits weights and tensors across cores, while pipeline parallelism splits layers across devices using pipelined micro-batches.These approaches support continued scaling while motivating improvements in training and serving efficiency.
- Conditional Computation: Conditional computation increases network capacity without increasing computation by activating parameters on demand per example or token.Gating decisions may be binary or sparse and continuous, stochastic or deterministic.
- Mixture of Experts: Switch Transformer selects one top expert per token, while prior MoE methods require auxiliary losses that do not guarantee balanced loads.Hard capacity limits can still leave tokens unprocessed, and several alternative routing methods are evaluated only on pre-training perplexity.
3 Method
The method replaces token-choice routing with expert choice, addressing load imbalance and fixed per-token computation by assigning each expert a fixed token capacity while allowing variable expert counts per token. It integrates this routing into sparse MoE feed-forward layers and supports an optional constraint on expert selections per token.
- Routing pitfalls: Token-choice routing can create load imbalance, under-specialized experts, and wasted capacity while assigning every token the same number of experts.Over-utilized experts may reject tokens because of capacity limits, whereas under-utilized experts leave model capacity unused.
- Expert choice routing: Expert choice lets each expert independently select its top-k tokens using a fixed expert capacity, guaranteeing load balance by design.The selected capacity is the number of tokens each expert can process.
- Expert choice routing: Variable expert assignments allow compute to vary across tokens rather than forcing every token to receive exactly k experts.The method is motivated by allocating more computation according to input complexity.
- Routing computation: The routing function computes affinity scores and selects the k largest token scores for each expert, producing index, gating, and permutation matrices.The permutation matrix gathers selected tokens for each expert, while the gating matrix supplies expert weights.
- MoE layer: The selected tokens are processed by independent feed-forward experts and then returned to their original batch order using permutation and gating matrices.The expert outputs and final outputs can be computed efficiently with Einstein summation operations.
- Additional constraint: An optional regularized routing variant limits the maximum number of experts that can select each token using entropy-regularized linear programming and Dykstra’s algorithm.The constraint is used to study how variable expert counts affect pre-training and fine-tuning performance.
4 Experiments
The experiments evaluate expert choice across training efficiency, scaling, downstream fine-tuning, routing heterogeneity, and capacity-factor ablations. Expert choice improves convergence and downstream performance while preserving benefits from scaling experts and allowing variable routing.
- Training Efficiency: More than 2x faster convergence than GShard top-2 is achieved by EC-CF2 at matched activated model size and computational cost.EC-CF2 reaches GShard top-2’s perplexity in less than half the steps, while GShard steps are 20% slower.
- Fine-tuning on GLUE and SuperGLUE: EC-CF2 consistently outperforms Switch top-1 and GShard top-2 across 11 GLUE and SuperGLUE tasks, increasing average accuracy by more than 2%.Compared with the dense 8B counterpart, EC-CF2 increases the average fine-tuning score by 3.4 points.
- Ablation: Limiting tokens to two experts decreases fine-tuning accuracy by 0.8 points on average, whereas allowing up to three experts matches vanilla expert choice.The ablation supports variable expert counts per token as useful for downstream performance.
- Heterogeneity Matters: Most tokens route to one or two experts, while 23% route to three or four and about 3% route to more than four.The routing distribution demonstrates that expert choice learns heterogeneous allocations across tokens.
- Ablation: Reducing the capacity factor from two to one degrades perplexity but still outperforms top-1 gating, while expert choice outperforms dense models in perplexity and convergence time.The dense-model comparison also reports larger sparse-gating benefits for 100M-parameter experts.
5 Conclusion
The paper proposes a routing method for sparsely activated MoE models that addresses load imbalance and under-utilization while allowing different numbers of experts per token. It reports over 2x training-efficiency improvements and strong gains on 11 GLUE and SuperGLUE datasets.
- Expert choice routing addresses load imbalance and under-utilization in conventional MoE models.
- Expert choice routing enables different numbers of experts to be selected for each token.
- More than 2x training-efficiency improvements are reported against GShard and Switch Transformer models.
- The method achieves strong gains when fine-tuning on 11 GLUE and SuperGLUE datasets.
6 Limitations
The paper identifies deployment settings where expert choice routing does not immediately apply, including autoregressive generation and very small serving or inference batches. It also notes that MoE models retain a large memory footprint because parameters scale with the number of experts.
- Autoregressive text generation is not immediately supported because routing uses past and future tokens for top-k selection.
- Very small serving or inference batches pose another setting where expert choice routing does not immediately apply.
- Possible remedies include grouping tokens by sequence, selecting a global top-k, and capping expert or token selections.
- MoE models have a large memory footprint because the total number of parameters increases with the number of experts.
- Increasing the number of experts requires reserving many hardware devices, so dynamic power savings do not eliminate static reserved power.
7 Checklist
The checklist reports that the paper addresses ethics, reproducibility, training details, compute reporting, error bars, and asset citation. The dataset is not released, and the paper states that it contains no personal data, personally identifiable information, or offensive content.
- The authors state that the paper describes its limitations and has no potential negative societal impacts to discuss.
- The paper reports including code, data, instructions, training details, error bars, and compute-resource information needed to reproduce the main results.
- The paper reports citing the creators of existing assets.
- The dataset is not released yet.
- The paper reports no use of people’s data and no personally identifiable information or offensive content in the dataset.
A Comparison on Fine-tuning with a Dense Model
The expert choice MoE model is compared with a dense model on fine-tuning performance across 11 GLUE and SuperGLUE tasks. The reported comparison states that the expert choice model consistently and significantly outperforms the dense model.
- The expert choice MoE model consistently outperforms the dense model across 11 GLUE and SuperGLUE tasks.
- The comparison uses an 8B dense model and an Expert Choice EC-CF2 8B/64E model.
- The table caption reports that the expert choice method significantly outperforms the dense model in downstream tasks.
B Capacity Factor
Capacity factors determine how many experts each token can receive on average, enabling comparisons between expert choice configurations and prior gating methods.
- A capacity factor of n indicates how many experts each token can receive on average.
- EC-CF2 matches the computational footprint of GShard top-2 gating.
- EC-CF1 matches the computational footprint of Switch Transformer top-1 gating.
- EC-CF0.5 tests whether substantially lowering the capacity factor can still provide strong performance.
- Table 2 reports average fine-tuning results from three independent runs for each routing method.
C Capped Expert Choice
Capping expert choice limits how many experts can receive each token, while retaining load-balanced routing advantages over top-2 gating.
- Figure 1 compares validation perplexity for EC-BASE, EC-CAP2, EC-CAP3, and GShard top-2 gating on 100M/64E models.
- Restricting tokens to two experts worsens perplexity relative to base expert choice, while EC-CAP2 and EC-CAP3 still outperform top-2 gating.
- EC-CAP3 achieves perplexity comparable to EC-BASE because only a small fraction of tokens use more than three experts.
- Hashing-based routing performs substantially worse than expert choice in average fine-tuning scores and variance.
E Fine-tuning Details
Fine-tuning uses distinct hyperparameter settings for the dense baseline and 8B/64E MoE models, with selected layers frozen in the dense model.
- The 8B dense model uses a constant learning rate of 0.0001 and dropout rate of 0.1 during fine-tuning.
- For the 8B dense model, attention and feed-forward layers are frozen while embeddings and layer normalization remain trainable.
- The dense-model fine-tuning configuration was selected as optimal based on the authors’ hyperparameter search.
- The authors also performed hyperparameter searches for both baseline models and expert choice.