Router Z-Loss › Combined Auxiliary Losses

Real MoE systems use both load balancing loss (from the previous chapter) and z-loss together. These losses address complementary problems and work synergistically to create stable, effective MoE training. Neither loss alone is sufficient: load balancing without z-loss can still experience numerical instability, while z-loss without load balancing allows expert collapse and utilization imbalance.

These two problems are separate. Load balancing addresses a resource-allocation question: are all experts being used effectively? Z-loss addresses a computational-stability question: are the router's computations numerically well behaved? A model could theoretically have perfectly balanced expert load while still experiencing numerical instability in the router computation. Conversely, a model could have numerically stable logits while all routing probability concentrates on a single expert. Solving one problem does not solve the other, which is why both losses are included in the training objective.

Think of load balancing and z-loss as two different maintenance tasks for the same machinery. Load balancing is like making sure all the machinery components are getting equal use, preventing any single component from wearing out while others sit idle. Z-loss is like making sure the machinery is operating within safe temperature and pressure ranges. You need both: balanced use prevents premature failure of individual components, while operating within safe ranges prevents the entire system from breaking down.

The complete MoE training objective combines the task loss with both auxiliary losses:

$$\mathcal{L}_{\text{total}} = \mathcal{L}_{\text{task}} + \alpha_{\text{aux}} \mathcal{L}_{\text{aux}} + \alpha_z \mathcal{L}_z$$

where:

• $\mathcal{L}_{\text{total}}$: the complete joint objective
• $\mathcal{L}_{\text{task}}$: the primary task loss
• $\mathcal{L}_{\text{aux}}$: the auxiliary balancing loss (encourages uniform expert utilization)
• $\mathcal{L}_z$: the z-loss (prevents logit instability)
• $\alpha_{\text{aux}}, \alpha_z$: weighting coefficients for the auxiliary terms

This combined objective guides the model toward three goals simultaneously: achieving good task performance, maintaining balanced expert utilization, and preserving numerical stability. The relative importance of these goals is controlled by the two auxiliary coefficients.

Why does this formula make sense? Notice that the task loss operates at the macro scale of model performance, typically measured in nats or bits. The auxiliary losses operate at the micro scale of routing behavior. The coefficients $\alpha_{\text{aux}}$ and $\alpha_z$ serve as unit conversion factors, adjusting the scale of each auxiliary loss to make it numerically comparable to the task loss gradient. Without appropriate coefficients, one term could completely dominate the gradient, preventing the other terms from having any effect.

The two auxiliary losses serve distinct purposes that together create a reliable training framework:

Load balancing loss ensures that all experts are used and that no expert becomes overloaded or underutilized. Z-loss ensures that the routing computation itself remains numerically stable. These concerns are orthogonal: a model could have perfectly balanced load with numerically unstable logits, or stable logits with highly imbalanced load. Both auxiliary losses are needed to achieve both goals.

Understanding this orthogonality also clarifies when to increase each coefficient. If you observe expert imbalance (some experts receiving too many or too few tokens), increase $\alpha_{\text{aux}}$. If you observe training instability (loss spikes, NaN values, or LSE growth), increase $\alpha_z$. The two problems have distinct diagnostic signatures and distinct remedies, which makes them tractable to address independently.

The two losses can interact in subtle ways that you should understand:

• Synergy: Z-loss keeps routing soft, which makes it easier for load balancing to spread tokens across experts. Hard routing with very peaked probabilities resists load balancing pressure because the gradient from load balancing loss is proportional to softmax probabilities. When probabilities are nearly binary (0.99 for one expert, 0.01 for others), the load balancing gradient can only adjust the dominant expert's logit. Softer probabilities distribute the gradient across all experts, making load balancing more effective.

• Tension: Very strong z-loss ($\alpha_z$ too high) can make routing nearly uniform, reducing the importance of load balancing. When z-loss forces all logits to be similar, the resulting routing probabilities are already close to uniform regardless of load balancing. In this regime, further increasing load balancing coefficient has little effect. Conversely, aggressive load balancing might push some expert logits up to attract more tokens, which increases z-loss.

• Balance: Finding the right coefficient combination requires considering both effects. A common starting point is $\alpha_{\text{aux}} = 10^{-2}$ and $\alpha_z = 10^{-4}$, making load balancing the stronger constraint with z-loss giving background stabilization. This ratio reflects the fact that expert imbalance is typically a more pressing problem than numerical instability in early training, though both must be controlled.

The interaction between these losses also has an indirect effect on the experts themselves. When load balancing maintains reasonable utilization and z-loss maintains soft routing, the experts receive diverse training signals from many different token types. This diversity is what allows each expert to develop a distinct specialization over the course of training. Without these regularization terms working together, experts either collapse to handling everything (with no specialization) or become so specialized that they are useless for the full distribution of inputs.