Z-loss

Summary

Z-loss adds a term proportional to $(lo g Z)^{2}$ to the cross-entropy objective, where $Z = \sum_{v} exp (ℓ_{v})$ is the softmax partition function over the vocabulary and $ℓ_{v}$ are the pre-softmax logits.

L = L_{CE} + λ_{z} \cdot (lo g Z)^{2}

A typical coefficient is $λ_{z} = 1 0^{- 4}$ . The term encourages $lo g Z$ to stay near zero, anchoring the common offset that cross-entropy leaves unconstrained. It does not bound every logit difference.

Two reasons, in order of how often they come up:

Stability. Without z-loss, large LLM training runs occasionally see the lm_head logits drift to very large magnitudes, which can cause numerical issues in the softmax (especially in BF16 or FP8). PaLM and DeepSeek both used z-loss for this.
A controlled cooldown intervention. The Marin 8B retrospective reports a successful intervention in its recipe. Excluding normalization parameters from weight decay is an implementation choice, not proof that no other control exists. Z-loss anchors the common-offset direction; it does not bound every parameter norm or logit difference.

The mechanism

The cross-entropy loss for a single token with target $y$ is:

L_{CE} = - ℓ_{y} + lo g Z

Note that $L_{CE}$ is invariant to adding a constant to all logits (because $lo g Z$ shifts by the same amount that $ℓ_{y}$ does). That invariance means cross-entropy alone provides no gradient pressure on the common offset of the logits. The model can drift the entire logit vector up or down without changing the loss.

Learning objective

A shared logit shift leaves probabilities unchanged but changes z-loss

Derivation

SAME PREDICTION, FREE OFFSET

Logits: [0, 1]

Probabilities: [0.269, 0.731]

Target CE: 0.313

log Z: log(1 + e) = 1.313

SHIFT BOTH BY -1.313

Logits: [-1.313, -0.313]

Probabilities: [0.269, 0.731]

Target CE: 0.313

log Z: 0, so z-loss is 0

Read it this way: compare the rows across panels. Subtracting the same constant from every logit preserves their differences, so softmax and cross-entropy do not notice. The partition value does notice: z-loss prefers the equivalent representative whose log Z is near zero. It anchors a shared offset; it does not force all logits to be equal or erase the model's probability distribution. Original calculation checked against the PaLM training objective.

Check: Why can z-loss change while softmax probabilities and target cross-entropy stay exactly the same under a common shift?

Read it this way: Adding 2 to all logits multiplies the softmax denominator by exp(2). Softmax probabilities are unchanged, but log-normalizer-based z-loss changes.
Which quantity changes even when every softmax probability stays fixed?

Adding $λ_{z} (lo g Z)^{2}$ breaks the invariance. The gradient with respect to a logit $ℓ_{v}$ becomes:

\frac{\partial L}{\partial ℓ _{v}} = (p_{v} - 1 [v = y]) + 2 λ_{z} (lo g Z) \cdot p_{v}

The second term pulls $lo g Z$ toward zero, anchoring the otherwise free common offset. The pressure is small per-step but persistent.

When to use it

Consider a measured ablation when diagnostics indicate a relevant failure mode:

Track log-partition values, logit differences, finite-value failures, and held-out task behavior rather than choosing from model size alone.
Treat Marin’s cooldown result and coefficient as a source-specific example to reproduce or test, not a universal prescription.
Compare the actual low-precision implementation and a recoverable control run before attributing an instability to the common-offset direction.

Short fine-tuning and large pretraining both need a workload-specific decision; neither label settles whether the extra objective helps.

How to set the coefficient

A coefficient such as $1 0^{- 4}$ is a starting point for an ablation, not a scale-free guarantee of a small penalty. For that coefficient, $lo g Z = 2$ contributes 0.0004, while $lo g Z = 100$ contributes one. State token/batch reduction and log base: natural-log cross-entropy is measured in nats, not bits. The cross-entropy lesson separates the probability objective from added regularizers.

If you set $λ_{z}$ much larger, the model spends capacity matching $lo g Z = 0$ instead of fitting the data and downstream perplexity worsens. If you set it much smaller, you lose the regularization effect.

Subtle effects worth knowing

Subsequent Marin experiments showed that z-loss doesn’t just prevent logit blowup. In the steady state, z-loss increases the weight norm of the lm_head and decreases the scale of the final layer norm. This makes sense in retrospect: the final layer norm has a disproportionate impact on logit scale (one parameter affecting all logits), and it’s the only player not constrained by weight decay. Z-loss redistributes the pressure.

What an interviewer expects you to say

If asked about z-loss:

State that it’s a penalty on $(lo g Z)^{2}$ added to cross-entropy.
Explain that cross-entropy is invariant to adding a constant to all logits, so it provides no gradient on absolute logit scale. Z-loss breaks that invariance.
Mention both use cases: stability against logit blowup and regularization on logit scale during deep cooldowns.
Inspect the actual decay groups and precision policy; explain why a source-specific intervention is not a bound on every logit or parameter.