Load Balancing › Load Balancing

In the previous chapters, we explored how Mixture of Experts architectures use gating networks to route tokens to specialized expert networks. Top-k routing selects the most relevant experts for each token, allowing sparse computation while maintaining model capacity. By activating only a small fraction of the model's parameters for any given input, MoE architectures can scale model capacity beyond what would be computationally feasible with dense networks. A model with 64 experts and top-2 routing activates only about 3% of its expert parameters per token, yet retains access to all expert knowledge through the training process. This routing mechanism also introduces a challenge that can undermine the architecture's benefits.

The gating network can develop strong preferences for certain experts while ignoring others. This creates an imbalance that wastes model capacity, causes computational inefficiencies during distributed training, and in extreme cases leads to complete model collapse. This failure mode has derailed training runs, caused wasted compute, and motivated research on keeping experts busy. To train MoE models reliably, you need to understand why they break.

Load balancing addresses these problems by so all experts receive a roughly equal share of training tokens. Without explicit balancing mechanisms, MoE models consistently converge to pathological states where only a handful of experts process the large majority of tokens. The remaining experts sit idle, consuming memory and adding model parameters that contribute nothing to the output. This is one of the central paradoxes of MoE training: the architecture is designed to provide more capacity, yet without careful attention to routing dynamics, you often end up with less effective capacity than a simpler dense model.

Understanding why this happens and how to measure it is needed before exploring solutions like auxiliary losses in subsequent chapters. This chapter builds the diagnostic foundation: we examine the mechanics of imbalance, characterize failure modes, and develop quantitative tools to detect problems before they become catastrophic. Think of this chapter as learning to read an MoE model's diagnostic signals. Before you can prescribe treatment, you need to know what healthy behavior looks like and how to recognize failures.

In the next chapters, we'll cover auxiliary balancing losses and router z-loss, which are the primary tools practitioners use to actively correct these imbalances during training. For now, the goal is deep understanding of the problem itself.