Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer

Reading guide

Problem and motivation

Deep learning models heavily benefit from parameter capacity when trained on large datasets. However, standard architectures activate the full network for every input token or example. This leads to an impractical increase in computational cost as both network size and dataset volume scale.

While conditional computation (activating only subsets of a model per example) was proposed theoretically, earlier implementations struggled on modern accelerators due to:

• GPUs favoring uniform arithmetic over branching.
• The "shrinking batch problem," where routing subsets of examples leaves individual sub-networks with tiny, inefficient batch sizes.
• Inter-device network bandwidth bottlenecks.
• Imbalanced routing, where a few sub-networks dominate while others starve.

Key ideas

• Sparsely-Gated MoE Layer: Feed-forward sub-networks ("experts") are placed between recurrent layers (such as stacked LSTMs). Only a sparse top-$k$ selection of experts is evaluated per position.
• Noisy Top-K Gating: Tunable Gaussian noise is added to gating logits before the top-$k$ selection and softmax to provide continuous, backpropagatable gradients and encourage balanced exploration across experts.
• Combined Data and Model Parallelism: Standard layers and gating run data-parallel across devices, whereas experts are partitioned model-parallel across GPUs. Batches across all devices are combined so experts receive sufficiently large aggregated mini-batches.
• Balancing Constraints: Explicit auxiliary losses on expert importance and expert load prevent the gating network from collapsing to degenerate solutions that overutilize a handful of experts.

Method (with the key equations)

An MoE layer consists of $n$ feed-forward expert networks $E_1,\dots ,E_n$ and a gating network $G(x)$. The layer output is: $$y=\sum _{i=1}^{n}G(x{)}_i{E}_i(x)$$ Computation is avoided whenever $G(x{)}_i=0$.

Noisy Top-K Gating

To combine sparsity with differentiable gating, noise is injected into linear projections: $$H(x{)}_i=(x\cdot W_g{)}_i+\text{StandardNormal}()\cdot \text{Softplus}((x\cdot W_{\text{noise}}{)}_i)$$ $$\text{KeepTopK}(v,k{)}_i=\{\begin{array}{ll}{v}_i& \text{if }{v}_i\text{ is in the top }k\text{ elements of }v\\ -\infty & \text{otherwise}\end{array}$$ $$G(x)=\text{Softmax}(\text{KeepTopK}(H(x),k))$$

Load-Balancing Losses

To enforce equal total gating weights, an importance loss is defined over a batch $X$ using the coefficient of variation ($CV$): $$\text{Importance}(X)=\sum _{x\in X}G(x)$$ $$L_{\text{importance}}(X)=w_{\text{importance}}\cdot CV(\text{Importance}(X){)}^2$$ To ensure equal numbers of examples per expert, the smooth load estimator uses the normal CDF $\Phi$: $$P(x,i)=\Phi \left(\frac{(x\cdot W_g{)}_i-\text{kth\_excluding}(H(x),k,i)}{\text{Softplus}((x\cdot W_{\text{noise}}{)}_i)}\right)$$ $$\text{Load}(X{)}_i=\sum _{x\in X}P(x,i)$$ $$L_{\text{load}}(X)=w_{\text{load}}\cdot CV(\text{Load}(X){)}^2$$

Results (with the most important numbers)

• 1 Billion Word Benchmark: A computationally matched MoE-4096-h model achieved a test perplexity of 34.1 (a 24% reduction compared to computationally matched baseline LSTMs). A high-budget MoE model reached 28.0 test perplexity in 47 hours on 32 K40 GPUs, outperforming the prior state-of-the-art result of 34.7 (which took 59 hours on 32 K40 GPUs).
• 100 Billion Word Corpus: Scaling up to 65,536 experts (68 billion parameters) dropped test perplexity by 39% compared to the dense baseline (28.9 vs. 47.0), maintaining 0.72 TFLOPS/GPU.
• Machine Translation: MoE-augmented GNMT achieved 40.56 BLEU on WMT'14 En$\to$Fr (+1.34 BLEU over GNMT without RL) and 26.03 BLEU on WMT'14 En$\to$De (+1.12 BLEU) while requiring less than half the ops per timestep.
• Multilingual Translation: A single multilingual MoE model outperformed a multilingual GNMT baseline on 11 of 12 language pairs (up to +5.84 BLEU on Korean$\to$English) and reduced dev perplexity by 19%.

Limitations and open questions

• Extreme sparsity can lead to diminishing returns or performance regressions; perplexity degraded when scaling from 65,536 to 131,072 experts on the 100B-word dataset.
• Distributed communication overhead remains dependent on the hidden-to-input/output ratio, making small expert capacities less bandwidth-efficient.
• Applying MoE recurrently inside RNN recurrence steps breaks the batch-unrolling optimization across sequence positions, requiring activation recomputation schemes.