Mixture-of-Experts with Expert Choice Routing

By Yan-Quan Zhou, Yanqi Zhou, Tao Lei, Han-Chu Liu, Nan Du, Yanping Huang, Vincent Y. Zhao, Andrew M. Dai, Zhifeng Chen, Quoc V. Le and 1 more. 2022

810 citations recorded by the source corpus.

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.

Subfields: Computer Science Applications, Economics and Econometrics, Management Science and Operations Research, Model Architectures

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Computer Science Applications

In Economics and Econometrics

In Management Science and Operations Research

In Model Architectures

Citations in the corpus

Mixtral of Experts (2024, methodology): “There are multiple alternative ways of implementing G ( x ) [6, 15, 35], but a simple and performant one is implemented by taking the softmax over the Top-K logits of a linear layer [28].”
MegaBlocks: Efficient Sparse Training with Mixture-of-Experts (2022, background): “More recently, Zhou et al. (2022) proposed to reverse the routing problem such that each expert selects its top k scoring tokens.”
A Review of Sparse Expert Models in Deep Learning (2022, background): “Center: “Choose Top-k ” along the Tokens axis are routing algorithms such as Zhou et al. (2022).”
From Sparse to Soft Mixtures of Experts (2023, methodology): “Following the usual design for Sparse MoEs (Lepikhin et al., 2020; Fedus et al., 2022; Riquelme et al., 2021; Zhou et al., 2022), we replace a subset of the Transformer’s MLP blocks with Soft MoE blocks.”
A Survey on Mixture of Experts in Large Language Models (2024, positive): “Expert-choice gating demonstrates stronger empirical performance than token-choice top-2 gating in Gshard and GLaM models, as validated by Zhou et al. [85] and their subsequent Brainformers study [86].”
QMoE: Practical Sub-1-Bit Compression of Trillion-Parameter Models (2023, background): “Those range from static assignment based on input tokens IDs (Roller et al., 2021), over dynamic token-to-expert matching (Zhou et al., 2022), to “soft” routing of linear input combinations (Puigcerver et al., 2023).”
Confident Adaptive Language Modeling (2022, background): “Another line of work involves conditional computation to train larger models that only use a sparser subset of the full network during inference, for example by routing over mixture-ofexperts [9; 22; 39; 91], recurring modules [18; 29; 35], or accessing external memory [82].”
Sparse Mixers: Combining MoE and Mixing to build a more efficient BERT (2022, methodology): “Router design is an active research area (Lewis et al., 2021; Roller et al., 2021; Zhou et al., 2022; Clark et al., 2022).”
Sparse Upcycling: Training Mixture-of-Experts from Dense Checkpoints (2022, methodology): “, 2021), and Expert Choice routing (Zhou et al., 2022).”
Deep Class-Incremental Learning: A Survey (2023, background): “On the other hand, pretrained language-vision models [147], [237] have shown powerful generalizability in recent years, and exploring the ensemble of various pre-trained models [238], [239] is also an exciting topic.”
On the Representation Collapse of Sparse Mixture of Experts (2022, background): “Zhou et al. (2022) let each expert select top-k tokens rather than distribute tokens to experts.”
Mixture-of-Experts in Remote Sensing: A Survey (2026, unknown): “MoE with expert choice routing reverses the usual perspective and lets experts choose a fixed number of tokens rather than tokens choosing experts; this achieves better load balancing and allows the number of experts per token to vary [157].”