Router Z-Loss › Z-Loss Coefficient

The z-loss coefficient $\alpha_z$ controls the strength of logit regularization and represents the primary hyperparameter that you must tune when using this technique. Before choosing a value, it helps to understand the mathematical relationship between the coefficient and the equilibrium logit magnitude. In equilibrium, the z-loss gradient equals the gradient from other terms pushing logits up, which means that larger coefficients correspond to smaller equilibrium logit values.

The z-loss coefficient $\alpha_z$ controls the strength of logit regularization and represents the primary hyperparameter that you must tune when using this technique. The total training objective with z-loss takes the form:

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

where:

• $\mathcal{L}_{\text{total}}$: the combined training objective that guides all parameter updates
• $\mathcal{L}_{\text{task}}$: the primary task loss (cross-entropy for language modeling)
• $\alpha_z$: the z-loss weighting coefficient, determining how strongly logit magnitude is penalized
• $\mathcal{L}_z$: the auxiliary z-loss term defined in the previous section

Choosing an appropriate coefficient requires balancing stability against routing effectiveness. Too small a coefficient provides insufficient regularization and allows instability. Too large a coefficient over-regularizes the router, preventing it from learning sharp routing patterns that enable expert specialization. The optimal coefficient lies in a middle ground that must be found empirically for each setting, though some general guidelines narrow the search considerably.

The task loss and z-loss operate on different scales. A typical cross-entropy task loss might be in the range of 3-8 during language model training (corresponding to perplexity values of 20-3000). A typical z-loss value might be 4-16 for well-behaved logits (corresponding to LSE values of 2-4). Setting $\alpha_z = 10^{-3}$ means the z-loss contributes roughly 0.004-0.016 to the total objective, a small fraction of the task loss. This relative scaling is intentional: z-loss should be a minor influence on the training objective, nudging the router toward stability without dominating the optimization.

Typical z-loss coefficients fall in the range $10^{-4}$ to $10^{-2}$, with several factors influencing the optimal choice:

• Model scale: Larger models often need stronger z-loss (higher $\alpha_z$) because they have more capacity for logit growth. The router network in a large model has more parameters and processes higher-dimensional representations, giving it more "room" to develop extreme logit patterns. The longer training runs typical of larger models also provide more time for drift to accumulate.

• Number of experts: More experts slightly increase the LSE for a given max logit because of the $\log N$ term in the LSE bounds, so the coefficient may need adjustment. However, this effect is typically small compared to other factors, since $\log N$ grows slowly with the number of experts.

• Training stability requirements: If training instabilities occur despite using z-loss, increasing $\alpha_z$ by 2-5x often helps stabilize training. This is a common first response to observed instability. However, if instabilities persist even with high z-loss coefficients, other factors like learning rate or batch size may need adjustment.

• Routing sharpness: Very high $\alpha_z$ can make routing too soft, reducing the benefit of expert specialization. When routing becomes nearly uniform, different experts receive similar token mixtures, preventing them from developing distinct specializations. Monitoring expert utilization and specialization can reveal if the z-loss coefficient is too high.

The ST-MoE paper used $\alpha_z = 10^{-4}$ as a starting point, finding it sufficient to stabilize training while preserving routing quality. This value has become a common default in subsequent MoE implementations, though you should treat it as a starting point rather than a universal optimum.

A practical approach to setting the z-loss coefficient involves the following steps, which balance systematic exploration with responsive adjustment:

1. Start conservative: Begin with $\alpha_z = 10^{-4}$ and monitor training stability. This conservative starting point provides some regularization without being so strong that it prevents useful routing patterns from developing. Most training runs will be stable at this coefficient.

2. Track LSE statistics: Log the mean and max LSE values across training batches. Healthy values typically stay below 10, though this depends on the number of experts. If LSE values grow steadily throughout training, the z-loss coefficient may need to be increased. If LSE values drop quickly to very small values, the coefficient may be too high.

3. Watch for instability: If loss spikes or NaN gradients occur, increase $\alpha_z$ by 2-5x. These events indicate that the current regularization strength is insufficient to prevent numerical problems. After increasing the coefficient, training may need to restart from an earlier checkpoint before the instability occurred.

4. Check routing quality: If routing becomes too uniform (all experts receiving similar load), reduce $\alpha_z$. Uniform routing defeats the purpose of having multiple experts, since all experts would receive the same training signal and develop similar capabilities. Some variation in expert load is healthy and indicates that the router is learning useful distinctions.

5. Validate on held-out data: Ensure z-loss isn't hurting generalization through over-regularization. While z-loss typically improves generalization by preventing overconfident routing, excessively strong regularization could prevent the model from learning useful routing patterns altogether.

In practice, many production MoE training setups monitor both LSE statistics and routing entropy simultaneously, using them as complementary diagnostics. LSE statistics reveal whether logit magnitudes are creeping up, while routing entropy reveals whether the routing distribution is becoming too uniform or too peaked. Both metrics should stabilize at healthy values during a well-tuned training run.

Some implementations use dynamic z-loss coefficients that change during training, adapting the regularization strength to the needs of different training phases:

• Warmup: Start with higher $\alpha_z$ during early training when instabilities are most common, then reduce as training stabilizes. Early training often sees rapid parameter changes that can push logits to extreme values, making strong regularization especially valuable. As training converges and parameters change more slowly, weaker regularization allows sharper routing.

• Curriculum: Increase $\alpha_z$ as model scale grows during progressive training. Some training approaches start with a small model and gradually increase its size. As the model grows, the capacity for logit drift increases, motivating stronger regularization.

• Adaptive: Adjust $\alpha_z$ based on observed LSE statistics, increasing when values exceed a threshold and decreasing when values are safely bounded. This approach automatically responds to the actual training dynamics rather than following a predetermined schedule. The risk is that adaptive schedules can be difficult to debug if they behave unexpectedly.

Dynamic schedules add implementation complexity but can improve training efficiency by applying strong regularization only when and where it is needed. For most practical purposes, a fixed coefficient tuned carefully at the start of training works well. Dynamic schedules become more valuable in very long training runs where the optimal regularization strength may change over time.