Cross-Entropy with Z-Loss

The problem

Softmax does not change if you add the same number to every logit (section 0). So nothing in the cross-entropy stops all the logits from drifting to huge values, and huge values hurt the precision of BF16 (brain floating point, section 8). The z-loss is a small extra term that penalizes the drift.

For each position $n$, call $\log Z_n = \mathrm{LSE}(z_n)$ the log-sum-exp of its logits (section 0): it is the "size" of the logits that softmax removes. The z-loss is

where:

• $N$ is the number of positions;
• $w$ is a small weight: OLMo 3 trains with $w = 10^{-5}$ (PaLM (Pathways Language Model) used $10^{-4}$);
• the result is added to the usual mean cross-entropy.

Example. One position with logits $(2, 1, 0)$ and right class 0, $w = 10^{-5}$.

• $\log Z = \log(e^2 + e^1 + 1) = 2.408$ and the cross-entropy is $\log Z - 2 = 0.408$.
• $\mathcal{L}_z = 10^{-5} \cdot 2.408^2 \approx 5.8 \times 10^{-5}$: negligible.
• Add 100 to every logit: the cross-entropy is still $0.408$, but $\log Z = 102.408$ and $\mathcal{L}_z \approx 0.105$, about 1,800 times larger. The z-loss grows with the square of the drift, so it pulls the logits back.

The gradient

By the chain rule, with $\partial \log Z / \partial z_c = \mathrm{softmax}(z)_c$ (section 2):

where $z_{n,c}$ is the logit of class $c$ at position $n$. It is a gradient that always pushes the logits down, in proportion to the probabilities, when $\log Z$ is positive.

Example. Two positions, logits [[2, 1, 0], [12, 10, 11]] and $w = 0.01$: the gradient is about [[0.016, 0.006, 0.002], [0.083, 0.011, 0.030]]. The second position has larger logits ($\log Z \approx 12.4$), so its gradient is larger.

What to implement

For logits of shape $(N, V)$ and integer targets of shape $(N,)$:

• cross_entropy_and_z_loss(logits, targets, z_weight=1e-5): return (ce, z): the mean cross-entropy and $\mathcal{L}_z$, as floats, both computed stably.
• z_loss_grad(logits, z_weight=1e-5): $\partial \mathcal{L}_z / \partial z$, shape $(N, V)$.

Requirements

• Stable for large logits: log-sum-exp with the maximum subtracted (section 0).
• No Python loops.

Hint

Compute $\log Z$ once, a vector of length $N$, and use it for the cross-entropy, for $\mathcal{L}_z$ and for the gradient. The softmax is np.exp(logits - log_z[:, None]).