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

Router probabilities

Let a token hidden state be $h\in\mathbb R^d$. A linear router produces logits and probabilities:

Top-k gating (sparse activation)

Let $S=\mathrm{TopK}(p(\cdot\mid h), k)$. Only experts in $S$ run:

Load balancing (avoid expert collapse)

Left alone, a router tends to favour a few experts, which then get trained more and favoured further. Shazeer et al. (2017) add routing noise and auxiliary losses to keep usage balanced. The Switch Transformer form (Fedus et al., 2022) uses, over a batch:

• $f_i$: fraction of tokens assigned to expert $i$
• $P_i$: average gating probability for expert $i$

and:

Here $N_E$ is the number of experts. The loss equals 1 when routing is perfectly uniform and grows when tokens and router probability pile onto the same experts; it is added to the training loss with a small coefficient (0.01 in Fedus et al.).

Uniform loads do not necessarily give its smallest value, because the hard counts $f_i$ need not match the average probabilities $P_i$. With two experts and top-1 routing, four tokens whose router probabilities are $(0.51, 0.49)$ three times and $(0.01, 0.99)$ once give $f=(0.75, 0.25)$, $P=(0.385, 0.615)$ and $\mathcal L_{\text{LB}}=0.885$. What the loss supplies is a gradient: $\partial\mathcal L_{\text{LB}}/\partial P_i = N_E f_i$. Measured in $P$, it 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 $\sum_i f_i P_i$. The router reaches $P$ only through each token's softmax, so in a single step one busy expert's average probability can still rise.

Expert capacity (selected is not always served)

Top-k routing only proposes token-expert assignments. A serving or training system may also cap each expert at a finite number of token slots for the batch. In the Switch Transformer, every expert gets the same capacity, set by a capacity factor $c$ (Fedus et al., 2022):

A capacity factor above 1 leaves room for uneven routing, at the cost of memory and computation spent on empty slots. If expert $e$ can accept $C_e$ assignments, then an assignment $(t,e)$ is served only while:

Once the expert is full, later assignments to that same expert overflow unless the implementation has an explicit rerouting or fallback rule. The important distinction is:

The interactive demo uses a deliberately simple dispatch rule: process tokens in batch order, fill each expert's slots, and mark later assignments to full experts as overflowed. Real implementations can add different overflow, padding or rerouting policies, but they still have to account for finite expert capacity.

Under this first-come rule, each expert ends with $\min(D_e,C_e)$ assignments, where $D_e$ counts the assignments that name expert $e$. That number does not depend on the order of the batch; which assignments overflow does. The practice after the demo reorders a small batch: the loads and the number of overflows stay the same, but one token loses both of its experts.

Expert parallelism

Large MoE models place different experts on different devices. Each MoE layer then sends every token's hidden state to the devices that hold its selected experts and sends the results back, usually as two all-to-all exchanges between devices (Lepikhin et al., 2021). This communication grows with the number of tokens and with $k$, and it is one reason a sparse model can be slower in wall-clock time than its computation count suggests. The MoE serving page looks at this cost in detail.