Mixture of Experts: Outrageously Large Neural Networks

A small gating network sends each word to a few out of thousands of expert networks, so the parameter count can grow while the compute per word stays flat.

The layer holds up to 131,072 feed-forward networks and runs only the handful with the highest gate scores on each input. Making that train on 2017 GPUs took noisy top-k gating, two load-balancing losses, and a way of batching that keeps every expert's matrix multiplies large.

Google Brain, January 2017. A language model with 4.3 billion parameters in one layer beat the best published perplexity on the 1 Billion Word benchmark, after the same 10 passes over the training data, while doing 6% of the previous best model's multiply-adds per word. The same layer, with one or two experts per token, is the core of Switch Transformers and Mixtral.

Noam Shazeer and six co-authors, including Geoffrey Hinton and Jeff Dean, wrote this paper to make an old idea pay off at scale. The idea is conditional computation: a network whose parts switch on or off depending on the input, so it can hold far more parameters than it uses on any one example. The paper builds one such part, the sparsely-gated mixture-of-experts (MoE) layer, and gets it to train at up to 137 billion parameters. Five pieces explain it:

1. the layer: many expert networks and a gate that weights them, with zero-weight experts skipped;
2. noisy top-k gating, which makes the gate sparse and keeps it trainable;
3. batching tricks that keep each expert's share of a batch large enough for a GPU to run efficiently;
4. two balancing losses that stop the gate from sending everything to a few experts;
5. the measurements: perplexity and BLEU against compute-matched baselines, from 4 to 131,072 experts.

Parameters versus compute

In an ordinary network every weight takes part in processing every example. A model with twice the parameters does about twice the compute per example, and since bigger models also want more training data, the paper notes that training cost grows roughly with the square of model size: double the parameters and the data together and you do about four times the work.

Language is a domain where more parameters keep helping. The previous best result on the 1 Billion Word benchmark (Jozefowicz et al., 2016) used LSTM language models whose LSTM layers ranged from 2 million to 151 million parameters, and the larger ones were better. An LSTM is a recurrent network that reads text one word at a time while carrying a state vector; the sequence-to-sequence translation models of the period were stacks of them. A language model is scored by perplexity, the exponential of its average negative log-probability per word. A perplexity of 34 means the model is, on average, as uncertain as if it were choosing uniformly among 34 words; lower is better.

Conditional computation had been proposed several times before 2017 (Bengio et al., 2013 and 2015; Davis & Arel, 2013; Cho & Bengio, 2014, among others), and none of it had produced a large gain in model size or quality. The paper lists five obstacles:

• GPUs are much faster at arithmetic than at branching, so gating has to switch large blocks on and off, not single units;
• GPUs need large batches, and a gate that splits a batch among many blocks leaves each block a small one;
• the network between GPUs is thousands of times slower than their arithmetic, so moving data to the right block can dominate;
• earlier schemes needed extra loss terms to control sparsity and balance (Bengio et al., 2015 used three);
• earlier work used image datasets of up to 600,000 images, too little signal to train billions of parameters.

The rest of the paper addresses all five: the first and fourth with the layer's design, the second and third with how the layer is distributed across GPUs, the fifth by training on text corpora of 1 billion to 100 billion words.

How a mixture-of-experts layer works

A mixture-of-experts layer has $n$ expert networks $E_1, \dots, E_n$ and a gating network $G$. For an input vector $x$, the gate outputs $n$ weights, and the layer's output is the weighted sum of the experts' outputs:

The saving comes from making $G(x)$ sparse. Wherever $G(x)_i = 0$, the term vanishes, so $E_i(x)$ is never computed. The experts in this paper all have the same shape and separate weights: an input layer, one hidden layer of ReLU units, and an output layer. In the 1 Billion Word models the input and output size is 512 and the hidden size is 1024, so one expert holds

weights (the paper rounds to 1M) and costs the same number of multiply-adds per input. A layer with 4096 such experts holds about 4.3 billion weights. With 4 experts selected per input, each input costs about 4.2 million multiply-adds in the MoE layer no matter how many experts exist.

The mixture-of-experts structure goes back to Jacobs, Jordan, Nowlan and Hinton (1991), where a softmax gating network weighted several expert networks and the whole model was the mixture, with every expert evaluated. Eigen, Ranzato and Sutskever (2013) stacked two MoE layers inside a deep network. This paper keeps the MoE as one layer of a larger model and makes the gate sparse, which turns the mixture into a way to add parameters without adding compute. A hospital triage desk is a fair picture: each patient is sent to a few of the hospital's specialists, and the hospital can hire more specialists without making any single visit longer. The picture fails in one place: nobody decides in advance what each expert specializes in. The gate and the experts are trained together, so the division of labour comes out of training.

In the language models the MoE layer sits between two stacked LSTM layers and is applied separately at every position of the text, so consecutive words can go to different experts. The paper calls this applying the MoE "convolutionally", since the same layer with the same weights runs at each time step. A residual connection adds the layer's input to its output, as in ResNet.

Noisy top-k gating

The standard gate for a mixture of experts multiplies the input by a trainable matrix $W_g$ of shape $d \times n$, giving one unnormalized score (a logit) per expert, and takes a softmax, which exponentiates the logits and divides by their sum so the weights are positive and add to 1:

A softmax output is never exactly zero, so (2) runs every expert. The paper makes the gate sparse by keeping only the $k$ largest logits and setting the rest to $-\infty$ before the softmax, which makes their gates exactly 0. It also adds Gaussian noise to the logits first:

Reading (4) term by term: $(x \cdot W_g)_i$ is the clean logit for expert $i$. A second trainable matrix $W_{\text{noise}}$, the same shape as $W_g$, sets how much noise each expert's logit gets for this input. Softplus, $\log(1 + e^z)$, turns that into a positive standard deviation, and $\epsilon_i$ is a fresh standard normal draw per expert. The noise is applied only during training; the official implementation multiplies the standard deviation by zero at inference and also adds a floor of 0.01 to it.

A worked example with $n = 8$ and $k = 2$: suppose the noisy logits for one input come out as

The two largest are expert 4 (1.7) and expert 5 (1.5). The other six become $-\infty$, and the softmax over the survivors gives $e^{1.7}/(e^{1.7}+e^{1.5}) = 0.55$ and 0.45. The layer computes

and the other six experts do no work for this input. Expert 1, with logit 1.4, missed by 0.1; on another draw of the noise it might have been selected instead of expert 5.

Figure 1 runs this gate on eight experts. Drag $k$ from 1 to 8 and watch the kept logits (teal) and the gate values below them, and the readout of how sensitive the gates are to the logits. Switch between training and inference to remove the noise, and redraw the noise to see near-ties change places.

The readout at $k = 1$ explains a choice the paper makes. The gate is trained by ordinary back-propagation: the task loss depends on $G(x)_i$ because (1) multiplies each expert's output by it. With $k \ge 2$, each kept gate is a softmax over several logits, and its derivative with respect to its own logit is $G_i(1 - G_i)$, 0.55 × 0.45 = 0.25 in the example above, so a gradient reaches $W_g$. With $k = 1$, the softmax has a single finite input and returns exactly 1 whatever that input is, so the derivative is zero and the task loss sends nothing to $W_g$. The paper puts it as "if we choose $k > 1$, the gate values for the top k experts have nonzero derivatives with respect to the weights of the gating network", and every model in it uses $k = 2$ or $k = 4$.

Which experts are selected is still a discontinuous function of the logits: as a logit crosses the $k$-th largest, an expert's gate jumps from 0 to a finite value. The paper calls these "theoretically scary discontinuities" and reports that it had not seen them cause trouble in practice. The alternative it rejects is the approach of Bengio et al. (2015), which treated the gates as boolean decisions and trained them with REINFORCE, a policy-gradient estimator from reinforcement learning that has much higher variance than a direct gradient. Keeping the gate values continuous and back-propagating through them, as here, also passes gradients through the gate to the layer's input.

The noise has a second job, covered in the load-balancing section below: it turns the number of inputs each expert receives into a smooth function of $W_g$. Both matrices start at zero (Appendix A), so at the first step every clean logit is 0, every noise standard deviation is $\mathrm{Softplus}(0) = \ln 2 \approx 0.69$, and the top-k choice is uniformly random, so every expert starts with an equal expected share of the batch.

Keeping each expert's batch large

A GPU is efficient on a large matrix multiply and wastes most of its time on a small one, because the cost of loading a weight matrix from memory is paid once per batch. A dense layer processing a batch of $b$ inputs multiplies one $b$-row matrix by its weights. An MoE layer splits that batch: if the gate picks $k$ of $n$ experts per input, each expert receives on average

inputs. With $k = 4$ and $n = 4096$, a batch of 18,750 words (one GPU's share of the 1 Billion Word batch) gives each expert about 18 words, which leaves its weights loaded for almost no computation. The paper calls this the shrinking batch problem and enlarges the batch in two ways.

Mixing data parallelism and model parallelism. In ordinary data-parallel training, each of $d$ GPUs holds a full copy of the model and runs its own batch. Here the LSTM layers and the gating network are replicated that way, but each expert exists once, on one GPU. All $d$ GPUs run their batches in lockstep, and at the MoE layer every GPU sends each input to the GPUs that hold its selected experts, so an expert receives the inputs routed to it from all $d$ batches at once:

examples per expert, a factor of $d$ more. The 1 Billion Word models trained on 16 GPUs with about 300,000 words per step in total, so the 4096-expert model gives each expert about 300,000 × 4 / 4096 ≈ 293 words per step instead of 18. Adding experts and GPUs in proportion keeps the per-expert batch, the memory per GPU and the step time constant, which is how the paper scales to 131,072 experts on 128 GPUs.

Applying the MoE to all time steps at once. An LSTM processes a sentence one word at a time, but the MoE layer's input at each position depends only on the first LSTM's output there. The implementation waits for the first LSTM to finish the whole sequence and sends every position of every sentence through the MoE as one batch, multiplying the batch by the number of time steps. This would not work if the MoE were inside the recurrence (for example, replacing an LSTM's weight matrices), since each step's input would depend on the previous step's output; the paper leaves that case open.

The network between GPUs is the other limit, since each input and output of an expert crosses it. For the transfer to stay cheap relative to the computation, an expert must do many multiply-adds per number sent. An expert with input size $d_{in}$, hidden size $h$ and output size $d_{out}$ does $h(d_{in} + d_{out})$ multiply-adds for $d_{in} + d_{out}$ numbers moved, a ratio of exactly $h$. The 512-1024-512 expert does 1024 multiply-adds per number transferred, and the multilingual translation model's experts, with 8192 hidden units, do 8192. The paper estimates that a GPU needs a ratio in the thousands, and a wider hidden layer raises it.

Why experts collapse, and the balancing losses

Left alone, the gating network tends to settle on a few experts and give them large weights for almost every input. An expert that is selected more often gets more gradient updates, so it improves faster; because it now produces better outputs, the task gradient raises its gate logits; with higher logits it is selected more often, which closes a feedback loop. Eigen et al. (2013) saw the same thing. The other experts receive few inputs, stop improving, and their parameters are wasted, and on a cluster the GPU holding a popular expert can run out of memory.

The paper measures balance with the importance of each expert over a batch $X$, the sum of the gate values it received:

This is a vector with one entry per expert. It adds a loss that is small when every entry is the same:

where the coefficient of variation is the standard deviation of the entries divided by their mean,

and $w_{\text{importance}}$ is a hand-tuned weight. Dividing by the mean makes the loss independent of the batch size: doubling every importance leaves it unchanged. Squaring it gives a gradient proportional to each expert's deviation from the mean.

A worked example with four experts and a batch of 8 inputs. Each input's gates sum to 1, so the importances always sum to 8 and their mean is fixed at 2. Suppose they come out as (3.6, 2.4, 1.2, 0.8). The deviations from 2 are 1.6, 0.4, −0.8 and −1.2, their squares average to 4.8 / 4 = 1.2, so $\mathrm{CV}^2 = 1.2 / 2^2 = 0.30$, and with $w_{\text{importance}} = 0.1$ the loss is 0.03. Because the mean is fixed, the gradient of $\mathrm{CV}^2$ with respect to expert $i$'s importance is

which is +0.2 for the first expert and −0.15 for the last. Back-propagated through the softmax, it lowers the first expert's gate values and raises the last one's wherever the two compete for an input.

Figure 2 is a toy version of the feedback loop, built for this page. Eight experts with noisy top-2 gating see 64 inputs per step; each expert's error falls as it sees more data, which is the part that makes the imbalance reinforce itself. Press Play at $w = 0$, then move $w$ to any positive value and play it again.

At $w = 0$ the busiest expert ends at 64 of 64 inputs, four times the mean, the most any expert can have with $k = 2$. The balanced runs end with the busiest expert at 1.13 to 1.25 times the mean.

Equal importance is not the same as equal work. Importance sums gate values, and the number of inputs an expert processes is a count. One expert could receive 2 inputs with gate 0.9 each and another 6 inputs with gate 0.3 each: both have importance 1.8, and the second does three times the computation and holds three times the activations in memory. On a GPU cluster the busiest expert sets the step time and the memory peak, so the paper adds a second loss on the count itself, the load.

A differentiable count of examples

The number of inputs an expert receives is an integer, and an integer has no useful gradient: a small change to $W_g$ either changes nobody's top k (gradient zero) or flips a selection (a jump). Appendix A replaces the count with its expected value under the gating noise, which is smooth.

Fix an input $x$ and an expert $i$, and keep the noise already drawn for every other expert. Expert $i$ makes the top $k$ exactly when its noisy logit beats the $k$-th largest noisy logit among the other experts, a threshold the paper writes $\mathrm{kth\_excluding}(H(x), k, i)$. Its own noisy logit is its clean logit plus a Gaussian with standard deviation $\mathrm{Softplus}((x\cdot W_{\text{noise}})_i)$, so the probability of a fresh draw beating the threshold is a normal CDF (the paper's equation (8) is this probability; (9) evaluates it):

$\Phi$ is the standard normal cumulative distribution function: the probability that a standard normal draw falls below its argument. Summing over the batch gives a smooth estimate of each expert's count, and the load loss has the same form as (7):

Take a clean logit of 1.0, a noise standard deviation of 0.5, and a threshold of 1.3 from the other experts. Then $P = \Phi(-0.6) = 0.27$: the expert is selected on about a quarter of noise draws. The derivative with respect to the clean logit is $\varphi(-0.6)/0.5 = 0.67$, where $\varphi$ is the normal density, so the load loss can push this expert's logit up or down even on inputs where it was not selected. The importance loss cannot do that: an unselected expert has a gate of exactly 0 and receives no gradient through (6).

In the official code the threshold depends on whether expert $i$ is currently in the top $k$. If it is, the $k$-th largest of the others is the $(k{+}1)$-th largest overall; if it is not, it is the $k$-th largest overall. That is why the forward pass takes the top $k + 1$ logits, one more than it uses.

In Figure 3, slide the clean logit across the threshold and shrink $\sigma$. The shaded area in the top panel is $P$; the bottom panel plots $P$ against the clean logit next to the 0-or-1 count with no noise.

The noise level is itself trained, through $W_{\text{noise}}$. A larger $\sigma$ widens the region of logits where the load gradient is nonzero and makes routing more random; a smaller one makes routing sharper and the load estimate closer to the hard count.

Appendix A tests the two losses on the 256-expert 1 Billion Word model, trained for 10 epochs at each setting:

With neither loss, perplexity is 39.8 and the busiest expert receives 17.8 times the average load. Any setting with at least one loss lands between 35.6 and 35.7, so the balancing losses cost nothing measurable in quality. They differ in the last column: more weight on the load loss gives a lower peak, down to 1.07 times the mean, which is what decides memory and step time on distributed hardware. The 1 Billion Word models use 0.1 for both weights; the translation models use 0.01.

Hierarchical MoE

With tens of thousands of experts, a single gate has to score all of them for every input. A two-level hierarchy cuts that down. A primary gate picks among $a$ groups, and each group is itself an MoE with its own gate over $b$ experts:

The 4096-expert 1 Billion Word model has 16 groups of 256, one group per GPU, with $k = 2$ at each level: 2 groups, then 2 experts in each, so every input still runs exactly 4 experts. Its gating networks score 16 groups and 2 × 256 experts per input instead of 4096. In the paper's accounting the two LSTM layers cost 2 million multiply-adds per word each and the 4 experts 4 million, for the 8 million "ops/timestep" budget that all the comparisons below hold fixed. The primary gate runs data-parallel like the LSTMs; each secondary MoE lives on one GPU.

What the paper measured

The experiments hold compute fixed and vary the number of experts. Every model in the first set does about 8 million multiply-adds per word in the forward pass, excluding the output softmax. The baselines with the same budget have no sparsity: two extra LSTM layers (4xLSTM-512), a single wide expert of 4096 units (MoE-1-Wide), four layers of 1024 (MoE-1-Deep), and an LSTM-2048-512 model identical to one of Jozefowicz et al. They score between 44.7 and 46.1 after 10 epochs.

Figure 4 plots the results from Tables 7 and 8 of the paper. Tap or hover over a point for its name, value and parameter count; the tabs switch between the two corpora and the measured GPU efficiency.

On the 1 Billion Word benchmark, the 4-expert model (all experts always on, so no sparsity) matches the baselines at 45.0. Perplexity then falls with every increase in experts: 39.7 at 32, 35.7 at 256, 34.1 at 4096. The last is 24% below the LSTM-2048-512's 44.7 at the same number of multiply-adds per word, a drop of 0.27 nats in per-word cross-entropy (ln 44.7 = 3.80, ln 34.1 = 3.53). The gains shrink beyond about 1 billion parameters: 1024 to 4096 experts buys only 0.5.

Two more models raise the compute alongside a 4-billion-parameter MoE layer, with larger LSTMs and fewer, larger experts. The paper compares all three with the previous best published model, 2xLSTM-8192-1024 at 151 million multiply-adds per word:

The low-budget MoE beats the published model with 5.9% of its multiply-adds per word. The high-budget MoE, with about the same compute, reaches 28.0, 19% lower. (The published model reaches 30.6 after 100 epochs; the MoE models were trained for 10.)

The paper hypothesized that a larger training set would let even more experts help, and repeated the experiment on about 100 billion words of Google News. After one pass over 100 billion words, perplexity falls from 47.0 for the compute-matched 4xLSTM-512 to 28.9 with 65,536 experts (69 billion parameters), 39% lower, and then rises slightly to 29.2 with 131,072. After only 10 billion words the curve flattens from 16,384 experts on, and the gap between the two curves widens with the expert count: the extra parameters help more when there is more data to put in them. With 4 of 65,536 experts active, the layer is 99.994% sparse per input.

The efficiency tab shows what the extra experts cost in GPU efficiency. The MoE models ran at 0.72 to 1.14 TFLOPS per GPU up to 65,536 experts, against 1.23 for the baseline and a K40's theoretical peak of 4.29. The 131,072-expert model fell to 0.30. Each of its experts saw roughly 2.5 million × 4 / 131,072 ≈ 76 words per batch, and the paper attributes the drop to not growing the batch with the GPU count (128 GPUs instead of 32).

In machine translation, the paper modifies Google's GNMT system (Wu et al., 2016), an LSTM encoder-decoder with attention, cutting it from 9 encoder and 8 decoder LSTM layers to 3 and 2 and inserting an MoE layer of up to 2048 experts in each. The result has 8.7 billion parameters and does 85 million multiply-adds per time step, against GNMT's 278 million parameters and 214 million. Translation quality is scored with BLEU, the overlap of word n-grams between the model's output and human reference translations, on a 0 to 100 scale. Its BLEU scores on WMT'14 are 40.56 for English to French and 26.03 for English to German, 1.34 and 1.12 above GNMT before its reinforcement-learning refinement step, which the MoE models did not use. On Google's production English-French data it beats GNMT by 1.01 BLEU after one day of training on 64 GPUs, against six days on 96 K80s for GNMT. These single-pair models use a different gate, described in Appendix F: during training each expert takes exactly the top $m = k|X|/n$ inputs of the batch, so every expert gets the same batch size, and at inference a learned per-expert threshold replaces the batchwise choice. The paper says the infrastructure ran faster with equal batch sizes at the time and has since been fixed.

The multilingual experiment trains one model on twelve language pairs at once, as Johnson et al. (2016) did with a single GNMT. This MoE model uses noisy top-k gating with 512 experts per layer and $k = 2$. Its development-set perplexity is 3.35 against the multilingual GNMT's 4.14, 19% lower; it beats the multilingual GNMT on 11 of 12 pairs, by up to 5.84 BLEU (Korean to English), and beats the twelve separately trained single-pair GNMT models on 8 of 12. The one pair where it loses to the multilingual GNMT is English to Korean, 16.62 against 18.41, which the paper attributes to overtraining: rare pairs were heavily oversampled, so the model saw the same few real examples many times.

Table 9 of the paper shows the experts specializing. It lists contexts that select three experts in the 2048-expert translation model: expert 381 fires around "researchers", "innovation" and "technological"; expert 752 on the "a" in "plays a critical", "assume a leadership", "played a vital"; expert 2004 on words about speed, such as "rapidly", "swiftly" and "fastest".

Implementation notes

The forward pass of one MoE layer, following the paper and the tensor2tensor implementation:

# Noisy top-k MoE layer, forward pass on a batch X of shape [B, d].
# W_g, W_noise: [d, n], zero-initialised. expert i maps d -> h -> d (ReLU).
clean = X @ W_g                                  # [B, n] clean logits
if training:
    std = softplus(X @ W_noise) + 1e-2           # [B, n] per-expert noise std
    H = clean + randn(B, n) * std                # eq (4)
else:
    H = clean                                    # no noise at inference
top_val, top_idx = top_k(H, k + 1)               # k+1: the extra one is a threshold
gates_k = softmax(top_val[:, :k])                # [B, k], each row sums to 1
G = scatter(gates_k, top_idx[:, :k], n)          # [B, n], zero outside the top k
y = zeros(B, d)
for i in range(n):
    rows = nonzero(G[:, i])                      # inputs routed to expert i
    y[rows] += G[rows, i:i+1] * expert[i](X[rows])   # eq (1)
importance = G.sum(axis=0)                       # eq (6), shape [n]
load = prob_in_top_k(clean, H, std, top_val, k).sum(axis=0)   # eqs (9)-(10)
aux_loss = w * (cv_squared(importance) + cv_squared(load))

Three groups of parameters receive gradients from this code. Each expert gets gradient only from the rows routed to it, scaled by their gate values. $W_g$ gets gradient from the task loss through the softmax over the kept logits (zero when $k = 1$), from the importance loss by the same path, and from the load loss through the numerator of (9), including for experts that were not selected. $W_{\text{noise}}$ gets gradient through the noise term inside the kept logits and through the denominator of (9). The selection itself (top_k, scatter) passes no gradient.

The 1 Billion Word models pass the MoE output through a sigmoid before dropout and the residual connection, and train with Adam under a learning rate that rises linearly for 1000 steps and then decays as the inverse square root of the step. The 100-billion-word models fit up to 1 billion parameters per GPU with two memory savings. The experts' hidden activations are recomputed in the backward pass instead of stored, the trade gradient checkpointing makes. And the optimizer state for the expert weights is shrunk: Adam normally keeps two extra numbers per parameter, so the paper sets $\beta_1 = 0$ to drop the first-moment estimate, and replaces the matrix of second-moment estimates with a vector of row averages and a vector of column averages, reconstructing each entry as their outer product divided by the mean. For one 512 × 1024 expert matrix that is 1,536 numbers instead of 524,288. Shazeer and Stern published this factored estimate as the Adafactor optimizer in 2018.

From LSTMs to Switch and Mixtral

The next generation moved the MoE into the Transformer, replacing feed-forward blocks with sets of experts. GShard (Lepikhin et al., 2020) used top-2 gating with a fixed capacity per expert, so inputs beyond an expert's capacity skip it and pass through the residual connection. This paper has no such limit: an expert processes every input routed to it, which is why imbalance shows up here as memory pressure on one GPU rather than as dropped inputs. Switch Transformers (Fedus, Zoph and Shazeer, 2021) went to $k = 1$. They avoided the zero gradient described above by changing the order of operations: the softmax runs over all $n$ logits first, and the chosen expert's output is multiplied by its probability, a number below 1 that still depends on $W_g$. Switch also replaced the two squared-CV losses with a single loss on the product of each expert's token fraction and mean gate probability.

Mixtral 8x7B (Jiang et al., 2024) uses 8 experts per layer and selects 2 per token: 47 billion parameters in total, 13 billion used per token. The gate is the paper's structure, a linear layer, top-k, and a softmax over the kept logits, without the noise.