MoE Load Balancing

without constraints, the gate makes tokens pile onto a few experts, causing load imbalance and routing collapse (a few experts absorb nearly all signal). Mitigations: the auxiliary balance loss $\mathcal{L}_{\text{aux}} = \alpha \cdot N \sum_{i=1}^{N} f_i P_i$, where $f_i$ is expert $i$'s fraction of tokens (actual utilization) and $P_i$ its mean routing probability — minimizing the product encourages uniform routing; plus a per-expert capacity cap (capacity factor: max tokens per expert = capacity factor × average load), where overflow tokens are dropped or routed through a residual path (recompute/residual connection).

a standard component of every top-k MoE training run (Switch Transformer, Mixtral, DeepSeek-V3); the capacity cap is also needed at inference to bound per-batch compute and prevent a hot expert from bloating latency.

imbalance idles some experts and overloads others, wasting memory and compute; expert collapse reduces MoE to a few effective experts, no better than a Dense model of the same size; the balance loss gives every expert a training signal, and the capacity cap makes training/inference deterministic.

Test Your Knowledge: Practice Questions for "MoE Load Balancing"

Single choice pitfall questions with instant feedback and mistake tracking.