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:

$$z = W_r h,\qquad p(e\mid h)=\mathrm{softmax}(z)_e.$$

Top-k gating (sparse activation)

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

$$\mathrm{MoE}(h)=\sum_{e\in S}\tilde p_e\,f_e(h),\qquad \tilde p_e = \frac{p(e\mid h)}{\sum_{j\in S} p(j\mid h)}.$$

Load balancing (avoid expert collapse)

In the cited MoE setups, routing can become imbalanced, so auxiliary losses and routing noise encourage balanced expert usage. A classic form uses:

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

and:

$$\mathcal L_{\text{LB}} = N_E\sum_{i=1}^{N_E} f_i\,P_i.$$

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. If expert $e$ can accept $C_e$ assignments, then an assignment $(t,e)$ is served only while:

$$\mathrm{load}(e) < C_e.$$

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:

$$\mathrm{TopK}(p(\cdot\mid h_t), k)\quad\text{chooses candidates; capacity chooses which candidates run.}$$

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.