Router Z-Loss › Summary

Router z-loss addresses a necessary but often overlooked problem in MoE training: the tendency for router logits to grow unboundedly, causing numerical instabilities that can derail training. The key concepts are:

• The instability problem: Router logits naturally drift toward larger values during training because nothing explicitly constrains them. Large logits cause softmax overflow, gradient instability, and poor generalization.

• Z-loss formulation: The z-loss penalizes the squared log-sum-exp of router logits, $\mathcal{L}_z = \frac{1}{B} \sum_i (\log \sum_j \exp(z_{i,j}))^2$. This provides smooth gradients that push large logits back toward reasonable values.

• Stabilization mechanism: Z-loss creates a restoring force that counteracts logit drift, maintaining bounded logits throughout training without preventing the router from learning discriminative patterns.

• Coefficient selection: Typical values range from $10^{-4}$ to $10^{-2}$, with larger models often needing higher coefficients. Start conservative and adjust based on observed LSE statistics and training stability.

• Combined auxiliary losses: Z-loss works alongside load balancing loss to address complementary problems. Load balancing encourages uniform expert utilization while z-loss ensures numerical stability. The combined objective is $\mathcal{L}_{\text{total}} = \mathcal{L}_{\text{task}} + \alpha_{\text{aux}} \mathcal{L}_{\text{aux}} + \alpha_z \mathcal{L}_z$.

With both auxiliary losses in place, MoE models train stably while maintaining balanced expert utilization. The next chapter explores expert parallelism: how to distribute MoE computation across multiple devices efficiently.