Router z-loss for stable MoE training
Router z-loss in MoE training is a small auxiliary term with a narrow job. It keeps the router's log-partition value near zero so its logits do not drift to a numerically awkward common offset.
It neither balances experts nor makes routing probabilities uniform. A router can have zero z-loss and still send almost every token to one expert.
That distinction is easier to see by calculating the loss from four logits, then shifting all four without changing a single routing probability.
Router z-loss in MoE training
For one token with N expert logits, call the logits x1 through xN. First compute log Z = log of the sum of exp(xj) across all experts.
Square log Z, then average that square across the B tokens in the batch. In plain notation, Lz = average over tokens of log-sum-exp(router logits) squared.
The ST-MoE paper introduced this router loss after selective float32 routing was not enough to stabilize its largest sparse models.
Its study used sparse models with 32 experts. Four of six baseline runs stayed stable, while all three runs with router z-loss stayed stable and kept model quality.
The training objective adds a coefficient beta: total loss = language-model loss + load-balancing term + beta times router z-loss.
ST-MoE chose beta = 0.001 after a sweep. OLMoE used the same value for OLMoE-1B-7B.
The coefficient matters because an unweighted z-loss can be much larger than the model's main loss. Beta sets its gradient contribution; it does not change where the z-loss itself reaches zero.
Calculate it from four logits
Take one token and four expert logits: 4, 2, 0, and -2.
Their exponentials are about 54.60, 7.39, 1, and 0.14. The sum is 63.13, so log Z is 4.145.
Square that value. The token's router z-loss is 17.182. With beta = 0.001, its contribution to the total training loss is 0.0172.
The softmax routing probabilities are about 0.865, 0.117, 0.0158, and 0.0021. Expert one wins by a wide margin.
Now subtract 4.145 from every logit. The shifted logits are about -0.145, -2.145, -4.145, and -6.145.
Softmax is invariant to a common additive shift, so those four routing probabilities stay exactly the same. The selected expert also stays the same.
But the shifted exponentials now sum to one. Log Z becomes zero, and the z-loss becomes zero.
Router z-loss chooses a stable offset for the logits while leaving the probability distribution free to be sharp, flat, balanced, or imbalanced.
Fanout's MoE routing explainer covers how the probabilities become top-k assignments after this calculation.
Why log-sum-exp is the target
Softmax has a redundant degree of freedom. Adding the same constant c to every logit changes log Z by c but leaves every probability unchanged.
The main model objective has little reason to care which common offset the router uses. Optimizer updates can therefore let that offset wander even when the router's decisions are already useful.
Router z-loss removes that freedom by preferring log Z = 0. At that point, the exponentials already sum to one and each logit can be read as the log of its softmax probability.
This is more precise than saying the loss merely makes every logit small. It controls the log-partition value, not an L2 norm of the logit vector.
For example, logits can remain far apart after the optimal common shift. The four-logit example still spans six units and still assigns 86.5 percent probability to one expert.
The square matters too. A positive or negative log Z is penalized, and the minimum sits at zero rather than at negative infinity.
Follow the gradient once
For one token, the derivative with respect to logit j is 2 times log Z times softmax(xj). Multiplying by beta gives the term that reaches the total gradient.
In the example, beta = 0.001 produces gradients of about 0.00717, 0.00097, 0.00013, and 0.000018 before the batch average.
Gradient descent subtracts those positive values because log Z is positive. It pushes most on the highest-probability expert, which contributes most to the partition sum.
If log Z were negative, the gradient signs would reverse and training would raise the logits. Both directions move the partition value toward zero.
At log Z = 0, every derivative from this auxiliary term is zero. The model's main and load-balancing losses can still change relative expert scores.
This gradient also explains why the reduction must be explicit. Averaging over tokens, summing over tokens, or averaging again across data-parallel ranks changes the effective coefficient.
Z-loss is not load balancing
Load-balancing loss asks whether tokens are distributed across experts. Router z-loss asks whether the logits use a numerically sensible offset.
The two losses can disagree without either being broken.
- Zero z-loss can coexist with one expert receiving nearly all probability.
- Perfectly balanced probabilities can have a large z-loss if every logit shares a large positive offset.
- A low load-balancing loss does not prove that exponentials and router matmuls are numerically well behaved.
- A low z-loss does not prove that expert capacity, communication, or token dropping is balanced.
OLMoE trained with a load-balancing coefficient of 0.01 and a z-loss coefficient of 0.001 because each term covers a different failure mode.
The token choice vs expert choice comparison explains which routing direction guarantees balance and which needs an auxiliary objective.
What the papers measured
ST-MoE compared stability techniques on models that were unstable often enough to study without paying for full-scale failures.
Tighter update clipping stabilized all three tested runs but damaged quality. Router z-loss stabilized all three and slightly improved the paper's quality measure relative to its baseline mean.
The authors connected the problem to roundoff error. Sparse expert models add many exponentials through routers, and small probability errors can change thresholded second-expert routing or the weight applied to an expert output.
OLMoE later ran a direct z-loss ablation. It reported fewer loss spikes, lower training and validation loss, and higher downstream performance with beta = 0.001.
That benefit cost about 2 percent training throughput. This is measured overhead, not a universal constant; it depends on implementation, hardware, token count, and how auxiliary losses are reduced.
Z-loss is a training expense. Inference does not optimize an objective, so an inference-only router can skip the calculation.
NVIDIA's Megatron Core router documentation describes its inference router as omitting z-loss and other training-only work.
Implementation checks that prevent silent mistakes
Compute log-sum-exp on the full router logits before top-k discards experts. The partition sum needs all N experts.
Use a stable log-sum-exp operation instead of exponentiating raw logits and taking a logarithm by hand. ST-MoE also casts tensors that feed exponentials to float32.
Mask padding tokens before the batch reduction. Otherwise easy zeroed positions can change both the reported z-loss and its effective coefficient.
Record the unweighted loss and beta separately. A dashboard that only shows beta times Lz makes it hard to distinguish healthier logits from a smaller coefficient.
Check how the framework averages across tokens, microbatches, tensor-parallel ranks, and data-parallel ranks. The same literal beta can create a different gradient scale under a different reduction.
Monitor router log Z alongside expert loads and the main loss. Those three traces separate offset drift, routing imbalance, and model-quality problems.
How to decide whether it is working
Start with the paper's 0.001 only as a reference point. The useful coefficient is the smallest one that controls log-partition drift and loss spikes without dominating the main objective.
Compare matched runs with and without the term. Track stability across seeds rather than celebrating one run that did not diverge.
Inspect the distribution of log Z, not just its mean. A small batch mean can hide a long tail of tokens with large positive or negative values.
Measure throughput because the auxiliary reduction and synchronization are real work. OLMoE accepted a roughly 2 percent cost for better stability and quality in its setup.
Keep load metrics separate. If z-loss falls while one expert remains overloaded, tune the routing or balancing mechanism instead of increasing beta.
The expert-parallelism guide shows why the busiest expert shard can set distributed step time even when the other GPUs are underused.
Router z-loss has one clean success condition: log-sum-exp stays controlled, stability improves across comparable runs, and the model pays an acceptable training cost. Expert balance is a different test.