Z-Loss and Training Stability

Softmax converts logits into probabilities by exponentiating them and dividing by their sum. Call that sum Z. A common z-loss term is the square of ln Z, multiplied by a coefficient and averaged over the relevant positions. Google’s T5X implementation adds such a term to cross-entropy. The name describes the normalization quantity being controlled, not a new replacement for the prediction task. The motivation becomes clearer from softmax’s shift property. Adding the same constant to every logit leaves its probabilities unchanged in exact arithmetic. Cross-entropy therefore does not identify a unique common offset for the logits. Z-loss responds to that offset because it changes ln Z. Encouraging ln Z toward zero can help control this otherwise unconstrained direction, while numerical implementation and precision still matter. For a constructed two-class example, logits [0, 0] produce equal probabilities and Z = 2. The unweighted penalty is (ln 2)², about 0.48045. Subtract ln 2 from both logits and each exponential becomes 0.5. Now Z = 1, the penalty is zero, and the probabilities remain equal. Zero auxiliary loss has not made the prediction correct; it has changed the logit normalization. ST-MoE adapts this idea to router logits and reports improved stability in its tested sparse-model configurations. Its router z-loss is distinct from the load-balancing auxiliary loss that addresses expert usage. Do not treat either result as a universal guarantee. Select the coefficient and target logits deliberately, track task loss and auxiliary loss separately, and inspect held-out quality. A run can become unstable for other reasons, including optimization settings, data problems, or numerical errors elsewhere in the computation.