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.

What are f_i, P_i, α, and N in the load-balancing loss?

Why does imbalance occur — what is expert/routing collapse?

What is the capacity factor — effects of it being too large or small?

How are overflow tokens handled?

Trade-offs of the aux-loss weight α being too large or small?