Load Balancing › Limitations and Impact

Load balancing in MoE models involves basic tensions that do not have clean resolutions. Understanding them supports principled tradeoffs rather than blindly optimizing metrics.

The first tension is between enforced balance and natural specialization. Perfect balance, where every expert receives exactly the same number of tokens, may not be optimal for model quality. Different token types require different processing, and some semantic categories may be more prevalent in the data than others. Forcing exact uniformity could degrade performance by preventing natural specialization patterns. Consider a model trained on a mix of English and code: code tokens might cluster toward certain experts that develop numerical or syntactic processing capabilities. Forcing those experts to process equal amounts of natural-language text and code could prevent them from fully specializing. The goal should be approximate balance that prevents pathological collapse rather than strict uniformity.

The second tension involves measurement stability. The metrics we've explored have limitations. Token fractions measured per batch can fluctuate substantially; an expert might appear underutilized in one batch but be needed for a specific topic that appears in the next. Longer-term averaging provides more stable estimates but responds slowly to shifts in expert utility. Entropy-based metrics capture the router's confidence distribution but do not distinguish between an expert appropriately confident on relevant tokens and one that is inappropriately dominant. No single metric tells the complete story, and over-optimizing any single metric can introduce distortions.

The third tension is between proactive and reactive balancing. Capacity constraints (reactive) prevent severe imbalance from causing computational problems but introduce token dropping, which degrades model quality. Auxiliary losses (proactive, covered in the next chapter) prevent imbalance from developing but add a training signal that may conflict with the primary task objective. Getting the balance right requires careful calibration: too little auxiliary loss and you still get collapse; too much and the auxiliary signal overwhelms the task signal, degrading the model's primary task performance.

Despite these limitations, load balancing is necessary for practical MoE deployment. Without it, the promised efficiency gains of sparse computation evaporate as models collapse to effectively dense computation through a small subset of experts. The training instabilities from severe imbalance, including gradient variance, capacity oscillation, and convergence failures, make unbalanced MoE models difficult to train at scale. Every major MoE system in production (Switch Transformer, GLaM, Mixtral, DeepSeek-MoE) uses some form of load balancing, which is the strongest possible endorsement for its necessity. These systems collectively represent trillions of tokens of training experience and billions of dollars of compute, and their builders independently concluded that balancing mechanisms are non-negotiable.

The failure to balance has real consequences beyond academic embarrassment. Training runs that collapse mid-way through are a significant waste of resources. For large models trained on expensive clusters, a single failed run due to expert collapse might cost hundreds of thousands of dollars and weeks of time. The investment in understanding and implementing proper load balancing is repaid many times over.

The next two chapters address load balancing through loss function design. Auxiliary balancing losses add a penalty term that encourages the router to spread tokens across experts. Router z-loss specifically targets the softmax's tendency to sharpen into collapse-prone distributions. Together with capacity constraints, these mechanisms enable stable training of MoE models with balanced expert utilization.