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

Neural networks absorb more knowledge as they get more parameters, but in an ordinary network every parameter does work on every input. Doubling the model doubles the cost of every example, and with growing datasets, training cost rises roughly with the square of scale. Conditional computation, running only part of the network per example, was the obvious answer, and earlier attempts had not shown large gains in practice. GPUs are slow at branching, conditional parts see small batches, network bandwidth limits what can be split across machines, and gates tend to favor a few paths.

Shazeer and colleagues built a layer that made it work. A mixture-of-experts (MoE) layer holds n experts, each a small feed-forward network, and a gating network that, for each token, computes a score per expert, adds Gaussian noise scaled by a learned amount, keeps the top k scores and softmaxes them. Only those k experts run, and their outputs are summed with the gate weights. With k = 4 and 4,096 experts, 99.9% of the layer is idle for any given token.

Three engineering pieces made it efficient. The experts live on separate devices while the rest of the model runs data-parallel, so each expert receives tokens from every device's batch at once. The MoE is applied to all time steps of the previous LSTM layer together, enlarging its batch again. And a loss term penalizes uneven expert use, because otherwise the gate settles on a few experts that then train faster and get chosen even more.

The payoff was large. A 4-billion-parameter MoE language model using 6% of the compute of the best LSTM of the day beat it, a 68-billion-parameter one cut perplexity 39%, and an 8.7B-parameter MoE translation model beat Google's production GNMT system.

Most frontier and many open models today are mixtures of experts. This paper explains the trade they make, and that trade shows up in every deployment conversation.

Compute follows active parameters, memory follows total parameters. An MoE model's cost per token depends on the k experts it runs, but every expert's weights must sit in GPU memory. A model can be cheap per token and still need several GPUs just to hold it. Size hardware by total parameters and estimate throughput by active ones.

Batching works differently. Each expert only sees the tokens routed to it. At low traffic, experts receive tiny batches and GPUs run inefficiently; at high traffic, an unlucky expert can become a hot spot. Serving MoE well needs enough concurrent load and expert placement across devices.

Routing is learned, and it can be lopsided. The paper's own ablation shows an unbalanced gate sending nearly 18 times the average load to one expert. When a customer asks why an MoE model's latency varies with input type, uneven expert load is a likely place to look.