Z-Loss and Logit Regularisation

Softmax has a blind spot. softmax(z) and softmax(z + c) produce the identical output for any constant c added to every logit, because the constant appears in both the numerator and the shared denominator and cancels (see softmax-logits). Cross-entropy loss, computed purely from softmax's output, inherits that blind spot: nothing in the standard training objective penalises a model whose logits drift uniformly upward over training, even though the predictions never change. At small scale this is harmless. At the scale where activations are stored in bfloat16, and especially inside mixture-of-experts routers, it is a real source of training instability.

Why the drift matters in low precision

bfloat16 has roughly 8 bits of mantissa, so its relative precision is fixed but its representable range is enormous; large numbers lose absolute precision fast. A model whose logits have quietly drifted to the range of hundreds or thousands, with the actually-informative differences between logits still only a few units wide, is asking softmax to resolve small relative differences between large, coarsely-rounded numbers. exp() of those logits can also overflow before the max-subtraction stabiliser (see softmax-logits) even gets a chance to help, because the subtraction happens after the logits are already computed and stored at reduced precision. The result is loss spikes and occasional divergence late in long training runs, exactly when restarting from checkpoint is most expensive.

The auxiliary loss

Z-loss penalises the log of the softmax normaliser directly, pushing it toward zero:

Z = sum_j exp(z_j)
z_loss = coefficient * (log Z)^2
total_loss = cross_entropy_loss + z_loss

with coefficient typically small, on the order of 1e-4. Unlike cross-entropy, log Z is not shift-invariant: adding a constant c to every logit adds exactly c to log Z (log sum_j exp(z_j + c) = c + log sum_j exp(z_j)). This is precisely the direction cross-entropy cannot see, so z-loss fills the gap: it gives the optimiser a gradient signal that specifically discourages uniform logit drift, without changing what the softmax output (and therefore the model's predictions) actually is at any fixed point during training.

What it does and does not fix

Z-loss does not change the ranking or relative spacing of logits, only their absolute scale, so it has no direct effect on prediction accuracy in principle; the coefficient is kept small precisely to avoid it fighting the main objective. What it buys is numerical headroom: keeping log Z near zero keeps individual logits in a range where bfloat16 rounding error stays small relative to the differences that matter, and keeps exp(z_j) comfortably away from overflow. Fedus, Zoph and Shazeer, 2021 (Switch Transformer) discuss this stabilisation problem in the context of large sparse models trained in reduced precision; a router-specific z-loss, applied to the router's logits over experts rather than the vocabulary output, was analysed further in later sparse-model stabilisation work, since router logits are especially prone to drift because a handful of experts can dominate the router's output early in training and reinforce their own selection (see moe-load-balancing-loss).