Load Balancing › Summary

Load balancing ensures that all experts in a Mixture of Experts model receive meaningful training signal and contribute to model predictions. Without explicit balancing mechanisms, MoE training exhibits rich-get-richer dynamics where initially preferred experts attract more tokens, receive more gradients, and become even more preferred. This positive feedback loop leads to expert collapse, a state where only one or two experts handle nearly all tokens and the architecture's capacity advantage is entirely lost.

Key load balance metrics include the load imbalance factor (measuring deviation from uniform allocation as a multiple of the ideal), coefficient of variation (quantifying spread in token fractions across all experts), and router entropy (capturing the concentration of routing probability mass. This provides the earliest warning signal). These metrics are complementary: LIF catches severe single-expert dominance, CV detects diffuse multi-expert imbalance, and entropy monitors the router's internal preferences before they manifest as hard routing changes. Monitoring all three gives the most complete picture of routing health.

The consequences of imbalance extend beyond wasted model capacity. In distributed training with expert parallelism, load imbalance creates stragglers that destroy parallel efficiency. A load imbalance factor of 2 halves throughput, and collapsed models may run at only $\frac{1}{E}$ efficiency despite having $E$ experts. This makes load imbalance a model quality problem and a basic obstacle to economical training at scale.

Capacity constraints provide one mechanism for limiting imbalance by rejecting tokens that would overflow an expert's quota. The capacity limit $K_{\text{max}} = \lfloor C \cdot N/E \rfloor$ sets a soft ceiling that scales naturally with batch size and the number of experts. However, this approach drops tokens rather than preventing the router from developing imbalanced preferences in the first place. The auxiliary balancing loss and router z-loss covered in the next chapters address this by shaping the training objective itself, rewarding uniform routing and penalizing the sharp softmax outputs that accelerate collapse.