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).

Router z-loss versus output z-loss

The same penalty can be applied at two different points in a model: on the final vocabulary logits (stabilising the language-modelling head) or on a mixture-of-experts router's logits over experts (stabilising routing decisions). Both use the identical formula; they differ only in which logit vector z refers to. Router z-loss is the more commonly discussed variant in practice, because router instability compounds: a router whose logits have drifted to extreme values produces a near-deterministic, saturated top-k selection (see softmax-logits on softmax saturation), which starves the load-balancing loss of the soft gradient signal it needs to correct imbalance (see moe-load-balancing-loss).

When it falls down

• It is a stability tool, not a quality lever. Tuning the z-loss coefficient higher does not improve the model; past a small threshold it starts fighting the primary loss and can measurably hurt task performance.
• Mostly invisible in fp32 debugging. Small-scale experiments run in full precision rarely exhibit the drift severely enough to need z-loss, which can make it easy to skip during prototyping and only discover the need for it once a run is scaled up into bf16 and starts spiking.
• Does not replace gradient clipping or careful learning-rate scheduling. Z-loss addresses one specific failure mode (logit-scale drift); most large training runs still need standard optimisation stabilisers alongside it.
• Coefficient is corpus- and scale-dependent. A value tuned on one model size or vocabulary is a starting point, not a transferable constant, since both Z's typical magnitude and bfloat16's rounding behaviour interact with vocabulary size and hidden dimension.

Further reading

• Fedus, Zoph and Shazeer, 2021, Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity - large sparse model training and the precision instabilities that motivate logit regularisation.
• Milakov and Gimelshein, 2018, Online normalizer calculation for softmax - the log-sum-exp mechanics underlying Z and why it is numerically delicate.