Continuous Function

Foundation reference

Sparse Mixture of Experts: Routing, Load Balancing & Expert Parallelism

A sparse mixture-of-experts layer routes each token to a few expert networks and combines their outputs. Mixtral 8x7B has 47B parameters but uses about 13B for each token, because every layer routes each token to 2 of its 8 experts (Jiang et al., 2024).

A newer notebook covers this concept. This older page keeps its summary, math and sources.

Concept 22 of 100EfficiencyStudy phase 6
Open the demoOpen the newer notebookAll concepts
Key equation
MoE(h)=∑e∈Sp~e⋅fe(h)\mathrm{MoE}(h) = \sum_{e \in S} \tilde{p}_e \cdot f_e(h)
Reading map and next steps

Study guide

Where to start

Concept 22 of 100MoEEfficiency · Study phase 6: Modern efficiency & inference
Question

Why does an MoE layer save arithmetic but not memory, given that every expert's weights must stay available (or sharded across devices) although each token uses only k of them?

Open the demoGo to the mathOpen the newer notebookNext concept
How to study this pagePredict · Check · Generalize · Continue
PredictSay what you expect.

Before you move a control in the demo, say what you expect to change.

CheckCheck against the math and the papers.

Compare what the demo shows with the key equation and with the papers cited below.

GeneralizeName what stays the same.

Identify the part of the idea that survives a change of notation or model, so you can recognize it in another paper or system.

ContinueContinue in the newer notebook.

The newer notebook covers the same concept and has an interactive demo.

Why It Matters for Modern Models

  • A sparse mixture-of-experts layer routes each token to a few expert networks and combines their outputs. Mixtral 8x7B has 47B parameters but uses about 13B for each token, because every layer routes each token to 2 of its 8 experts (Jiang et al., 2024).
  • DeepSeek-V2 has 236B parameters with 21B active per token and a 128K-token context, and pairs MoE with a compressed KV cache (multi-head latent attention) to keep serving costs down (DeepSeek-AI, 2024).
  • MoE saves arithmetic but not memory: all experts must be stored, and routing tokens dynamically splits each batch into uneven groups per expert, which makes serving harder.
  • Where LLM Serving changes how a fixed model is run, MoE changes the model itself so that each token costs less computation.

Questions to Think About

  • Why does a router trained only by backpropagation tend to send most tokens to a few experts unless a balancing loss is added?
  • At what granularity should load be balanced (per micro-batch, per global batch, per sequence), and how can balancing too finely work against expert specialization?
  • How does MoE differ from an ensemble, given that each token passes through a different subnetwork?
  • Why does distributed MoE need two all-to-all exchanges per layer, one to send tokens to their experts and one to bring the results back?

Interactive Demo

Core Math

If you prefer intuition first, try the demo above, then come back for the math.

Key Equation
MoE(h)=∑e∈Sp~e⋅fe(h)\mathrm{MoE}(h) = \sum_{e \in S} \tilde{p}_e \cdot f_e(h)

A sparse mixture-of-experts (MoE) layer lets the total parameter count grow while the compute per token stays roughly fixed, because each token runs only kk of the experts (Shazeer et al., 2017).

Router probabilities (per token, per layer):

Given a token's hidden state h∈Rdh \in \mathbb{R}^d, a linear router produces expert probabilities:

z=Wrh,p(e∣h)=softmax(z)ez = W_r h, \qquad p(e \mid h) = \mathrm{softmax}(z)_e

Top-k gating (sparse activation):

Let S=TopK(p(⋅∣h),k)S = \mathrm{TopK}(p(\cdot \mid h), k). Only those experts run, and their outputs are mixed:

MoE(h)=∑e∈Sp~e⋅fe(h),p~e=p(e∣h)∑j∈Sp(j∣h)\mathrm{MoE}(h) = \sum_{e \in S} \tilde{p}_e \cdot f_e(h), \qquad \tilde{p}_e = \frac{p(e \mid h)}{\sum_{j \in S} p(j \mid h)}

Load-balancing loss (Switch Transformers form, Fedus et al., 2022):

With NEN_E experts, let fif_i be the fraction of tokens in the batch sent to expert ii and PiP_i the average router probability given to expert ii:

LLB=NE∑i=1NEfi⋅Pi\mathcal{L}_{\text{LB}} = N_E \sum_{i=1}^{N_E} f_i \cdot P_i

With uniform loads, fi=1/NEf_i = 1/N_E for every expert, the loss equals 1 whatever the PiP_i. Fedus et al. describe it as minimized under uniform routing. That holds when ff and PP coincide, but the hard counts fif_i need not match the average probabilities PiP_i, and an uneven split can score lower. With two experts and four tokens whose router probabilities are (0.51,0.49)(0.51, 0.49) three times and (0.01,0.99)(0.01, 0.99) once, top-1 routing gives f=(0.75,0.25)f = (0.75, 0.25) and P=(0.385,0.615)P = (0.385, 0.615), so

LLB=2 (0.75⋅0.385+0.25⋅0.615)=0.885<1.\mathcal{L}_{\text{LB}} = 2\,(0.75 \cdot 0.385 + 0.25 \cdot 0.615) = 0.885 < 1.

What the loss reliably supplies is a training signal, added with a small weight (0.01 in Fedus et al.). The counts fif_i are not differentiable; the gradient flows through PiP_i, with ∂LLB/∂Pi=NEfi\partial \mathcal{L}_{\text{LB}} / \partial P_i = N_E f_i. Measured in PP, that gradient is largest for the experts that received the most tokens, so a descent step shifts router probability, on balance, toward less-loaded experts: it lowers the load-weighted sum ∑ifiPi\sum_i f_i P_i. The router reaches PP only through each token's softmax, so in a single step one busy expert's average probability can still rise.

Key Papers

Mixtral of Experts

Jiang et al.2024arXiv
Read paper →

DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models

Dai et al.2024ACL
Read paper →

DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model

DeepSeek-AI2024arXiv
Read paper →

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

Shazeer et al.2017ICLR
Read paper →

Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity

Fedus, Zoph & Shazeer2022JMLR
Read paper →

Connections

Prerequisites

AttentionMaximum LikelihoodEfficiencyLLM ServingEfficient Attention

Leads to

MoE ServingState Space ModelsMixture-of-Depths

Where to go next

A newer notebook covers MoE, with an interactive demo.

Continue in the newer notebook
Other connections (9)

Which limitations does it address, or leave to a later idea?

  • MoE scheduling: a limitation of MoE that MoE Serving addresses. Open MoE Serving

    MoE makes compute sparse but introduces dispatch/combine all-to-all traffic and severe load skew. MoE serving systems were invented to schedule routing, reduce stragglers, and manage expert pools efficiently.

  • More params, same FLOPs: a limitation of Scaling Laws that MoE addresses. Open Scaling Laws

    In a dense model every parameter is used for every token, so parameter count and compute per token grow together. A mixture-of-experts model adds parameters while keeping the compute per token roughly fixed, by sending each token to only a few experts (Shazeer et al., 2017; Fedus et al., 2022).

Where does it break down?

  • Routing stragglers: where MoE breaks down, taken up in MoE Serving. Open MoE Serving

    Mixture-of-experts routing gives experts uneven loads, and the busiest expert sets the latency of the whole layer. Without careful scheduling and load balancing, the saving from sparse computation is lost to waiting on the slowest expert.

Where else does the same idea appear?

  • Sparse selection: an analogy between MoE and Attention. Open Attention

    A mixture-of-experts layer sends each token to its top-k experts and ignores the rest. Softmax attention is dense, giving every position some weight, although trained heads often concentrate on a few positions; sparse-attention variants make that selection explicit.

  • Top-k then renormalize: the same technique appears in Decoding & Sampling. Open Decoding & Sampling

    MoE keeps top-k experts and renormalizes gate. Top-k decoding applies same truncation-and-renormalization to vocabulary logits before sampling.

  • Softmax routing: the same technique appears in Attention. Open Attention

    Both use a softmax to turn scores into weights: attention over the positions in a sequence, a mixture-of-experts router over a fixed set of experts. The router is a learned linear layer, not an attention head, and it keeps only the top-scoring experts instead of averaging over all of them.

  • Sparse features: an analogy between MoE and Superposition. Open Superposition

    Superposition lets a layer represent more features than it has dimensions, as overlapping directions, which works because only a few features are active at once (Elhage et al., 2022). A mixture-of-experts layer likewise relies on each token needing only a few experts. Both exploit sparsity.

  • Top-k gating: the same technique appears in Sparse Autoencoders. Open Sparse Autoencoders

    k-sparse SAEs enforce TopK activations; MoE routing enforces TopK experts per token. Both implement sparse selection as "keep top-k, drop the rest, then compute through a smaller active set."

  • Draft-verify routing: an analogy between MoE and Speculative Decoding. Open Speculative Decoding

    Both put a cheap component in front of an expensive one. An MoE router, a small linear layer, decides which experts run for each token; a small draft model proposes tokens that the large target model then checks. The analogy is loose: routing saves computation by skipping experts, while speculative decoding still runs the target model and saves time by checking several tokens in one pass.