Router Z-Loss › Worked Example: Z-Loss Step by Step

To make the mathematics concrete, let's trace through a single z-loss computation with specific numbers. Suppose we have a batch of $B = 3$ tokens and $N = 4$ experts. The router produces the following logit matrix:

$$Z = \begin{pmatrix} 3.0 & 1.0 & 0.5 & -0.5 \\ 1.5 & 1.2 & 1.0 & 0.8 \\ 8.0 & 2.0 & 1.0 & 0.0 \end{pmatrix}$$

where row $i$ contains the logits for token $i$ across all four experts. Notice that token 1 has moderate, nearly uniform logits, while token 3 has a very large logit for expert 1.

Step 1: Compute LSE for each token.

For token 1, with logits $[3.0, 1.0, 0.5, -0.5]$:

$$\text{LSE}(z_1) = \log(\exp(3.0) + \exp(1.0) + \exp(0.5) + \exp(-0.5)) = \log(20.09 + 2.72 + 1.65 + 0.61) = \log(25.07) \approx 3.22$$

For token 2, with logits $[1.5, 1.2, 1.0, 0.8]$:

$$\text{LSE}(z_2) = \log(\exp(1.5) + \exp(1.2) + \exp(1.0) + \exp(0.8)) = \log(4.48 + 3.32 + 2.72 + 2.23) = \log(12.75) \approx 2.55$$

For token 3, with logits $[8.0, 2.0, 1.0, 0.0]$:

$$\text{LSE}(z_3) = \log(\exp(8.0) + \exp(2.0) + \exp(1.0) + \exp(0.0)) \approx \log(2981 + 7.39 + 2.72 + 1.0) = \log(2992.1) \approx 8.00$$

Notice that token 3's LSE is almost exactly equal to its maximum logit (8.0), confirming the approximation is tight when one logit dominates strongly. Token 2's LSE is 2.55, compared to its maximum logit of 1.5, a gap of 1.05, close to $\log(4) \approx 1.39$ since the logits are nearly equal.

Step 2: Square each LSE value.

$$\text{LSE}(z_1)^2 \approx 3.22^2 = 10.37$$ $$\text{LSE}(z_2)^2 \approx 2.55^2 = 6.50$$ $$\text{LSE}(z_3)^2 \approx 8.00^2 = 64.00$$

Step 3: Average across the batch.

$$\mathcal{L}_z = \frac{1}{3}(10.37 + 6.50 + 64.00) = \frac{80.87}{3} \approx 26.96$$

Step 4: Understand what this tells us.

Token 3 contributes 64.00 to the sum, while tokens 1 and 2 together contribute only 16.87. Token 3's extreme logit of 8.0 for expert 1 is responsible for nearly 80% of the total z-loss, even though it is only one of three tokens. This is the quadratic amplification working as intended: logit values that are dangerously large receive a disproportionate share of the penalty, creating strong corrective pressure exactly where it is most needed.

Step 5: Compute the gradient for the dominant logit.

For token 3, the softmax probability for expert 1 is approximately:

$$p_{3,1} = \frac{\exp(8.0)}{\exp(8.0) + \exp(2.0) + \exp(1.0) + \exp(0.0)} \approx \frac{2981}{2992} \approx 0.996$$

The z-loss gradient for this logit is:

$$\frac{\partial \mathcal{L}_z}{\partial z_{3,1}} = \frac{2}{B} \cdot \text{LSE}(z_3) \cdot p_{3,1} \approx \frac{2}{3} \cdot 8.00 \cdot 0.996 \approx 5.31$$

This is a substantial gradient that will push $z_{3,1}$ downward in the next parameter update, counteracting the logit drift that produced this large value. The other logits for token 3 receive much smaller gradients, proportional to their much smaller softmax probabilities. This numerical trace confirms the theoretical properties we derived: the corrective pressure is both scale-proportional and routing-probability-weighted.