--- title: "Router z-loss vs load balancing loss" description: "Compare router z-loss with MoE load balancing loss using the formulas and two counterexamples that isolate logit scale from expert utilization." canonical_url: "https://fanout.sh/blog/router-z-loss-vs-load-balancing-loss" md_url: "https://fanout.sh/blog/router-z-loss-vs-load-balancing-loss.md" last_updated: "2026-08-29" access: "public" --- # Router z-loss vs load balancing loss Compare router z-loss with MoE load balancing loss using the formulas and two counterexamples that isolate logit scale from expert utilization. - Author: Suraj Gaud - Published: 2026-08-29 - Track: AI research - Access: Public - Tags: mixture of experts, router z-loss, load balancing loss, MoE training, expert routing, ST-MoE, Switch Transformer Router z-loss vs load balancing loss is a comparison between two different controls on a mixture-of-experts router. Load balancing loss watches where tokens go. Router z-loss watches the absolute scale of the logits that send them there. One loss can be ideal while the other is bad. A router can split tokens evenly with logits near 100, or send every token to one expert while its z-loss is exactly zero. That separation is useful during training. Expert counts diagnose balance. Log-sum-exp statistics diagnose z-loss. Treating either metric as a substitute for the other hides the failure you need to fix. ## Router z-loss vs load balancing loss The[Switch Transformer paper](https://www.jmlr.org/papers/v23/21-0998.html)defines a common load balancing loss as N times the sum of f_i times P_i across N experts. For expert i, f_i is the fraction of tokens actually dispatched to it. P_i is the average router probability assigned to it before top-k selection. Uniform routing gives f_i = P_i = 1/N. The unweighted loss is then 1, not 0. Its useful zero point is the excess above that floor: L_balance minus 1. The[ST-MoE paper](https://arxiv.org/abs/2202.08906)defines router z-loss as the batch average of the squared log-sum-exp of each token's router logits. The complete training objective has separate coefficients for language-model cross-entropy, load balancing loss, and router z-loss. The terms are additive because they constrain different properties. Fanout's[MoE routing explainer](/blog/mixture-of-experts-routing-explained)covers the token-to-expert path. The comparison here stays on the two auxiliary objectives. ## Load balancing loss measures utilization The hard fraction f records what top-k routing did. If 75 of 100 top-1 assignments select expert 0, then f_0 is 0.75. The soft fraction P records the router's average probability mass. It remains differentiable even though the selected expert index does not. The dot product couples intent to dispatch. An expert with both a large f_i and a large P_i raises the loss. Gradient descent can reduce P_i, which lowers the chance that later tokens choose that expert. Under perfect balance, both vectors are uniform and the loss reaches its floor of 1. If every token chooses one expert with probability 1, the loss reaches N. The range therefore depends on expert count. Report the raw loss with N, or report its excess over 1. A value of 1.4 means something different operationally for two experts and 128 experts. Load balancing also depends on the group over which f and P are computed. A local microbatch can look skewed while a global batch is balanced, or hide a sequence-level hotspot inside a balanced aggregate. The loss says nothing direct about the absolute router logits. Softmax discards a common offset, so logits can drift upward together without changing f, P, or the chosen experts. ## Router z-loss measures absolute logit scale For one token with logits x_1 through x_N, compute z = log of the sum of exp(x_i). Router z-loss averages z squared across the batch. Log-sum-exp is close to the largest logit. Its value lies between max(x) and max(x) plus log N, which makes it a smooth scale monitor. ST-MoE introduced this penalty after finding that selective float32 routing was not enough to make its largest sparse models reliable. In its sweep, a z-loss coefficient of 0.001 stabilized all three reported runs. The penalty reacts to a common logit shift. Add 100 to every expert logit and softmax probabilities stay unchanged, but log-sum-exp also rises by 100. Load balancing cannot see this shift. Subtracting the maximum inside a stable softmax implementation does not make z-loss redundant. Numerically stable softmax prevents a forward-pass overflow. Z-loss changes the learned router so that the original logits stay smaller before low-precision operations and exponentials amplify their rounding error. The earlier[router z-loss guide](/blog/router-z-loss-moe-training)follows that stability mechanism. Here, the counterexamples show why balance metrics cannot certify it. ## Zero z-loss can coexist with bad balance Take two experts and give every token the logits log(0.9) and log(0.1). Their exponentials sum to 1, so log-sum-exp is 0. Router z-loss is also 0. Softmax returns probabilities 0.9 and 0.1. Top-1 routing sends every token to expert 0, so f = [1, 0] and P = [0.9, 0.1]. The unweighted load balancing loss is 2 times (1 times 0.9 plus 0 times 0.1), which equals 1.8. Its floor is 1, so the excess is 0.8. Nothing is numerically large in this router. It still wastes one expert and can overflow expert 0's capacity while expert 1 sits idle. Z-loss cannot fix this example by itself. It already sees its preferred partition value. The needed gradient must come from load balancing or another routing-control mechanism. The example also corrects a common shortcut. Small logits do not imply uniform routing. Relative gaps determine softmax probabilities, while z-loss is sensitive to the partition function's absolute position. ## Perfect balance can coexist with huge z-loss Now route half the tokens with logits [100, 99] and the other half with [99, 100]. Top-1 assignments split evenly, so f = [0.5, 0.5]. The average probabilities also split evenly because the second half mirrors the first. P = [0.5, 0.5], and the load balancing loss is exactly 1. Each token's log-sum-exp is about 100.313. Squaring it gives about 10,062.75, so the batch z-loss is enormous despite perfect utilization. Subtracting 100 from every logit preserves the same softmax probabilities and the same assignments. It changes log-sum-exp to about 0.313 and z-loss to about 0.098. This pair of routers is indistinguishable to the load balancing objective. Z-loss strongly prefers the shifted version because it preserves routing with much smaller absolute numbers. The case is not an argument for comparing raw loss magnitudes directly. Training uses coefficients, and the two formulas have different natural scales. It is an argument for logging both. A balanced expert histogram does not prove that the router's exponentials and low-precision arithmetic are healthy. ## The coefficients are not interchangeable Switch Transformer used a load balancing coefficient of 0.01 after sweeping values from 0.1 to 0.00001. ST-MoE selected 0.001 for router z-loss in its reported setup. Those numbers are results from different objectives and experiments, not a universal ten-to-one ratio. Copying the ratio without inspecting the unweighted losses can make either term dominate. Track each weighted contribution beside the language-model loss. Also track max and mean log-sum-exp, tokens per expert, dropped-token rate, and routing entropy. If only a few experts are overloaded while log-sum-exp is stable, adjust the balance mechanism. Raising z-loss may soften probabilities, but it does not directly target the dispatch counts. If expert counts are even while log-sum-exp drifts upward, adjust z-loss or the source of logit growth. Raising the balance coefficient adds pressure to a metric that is already at its floor. If both fail, tune them as separate controls. First identify whether the imbalance is local, global, or sequence-specific. Then check whether the z-loss spike starts before or after the routing skew. ## A practical comparison for training runs Use load balancing loss to answer whether paid-for expert capacity is receiving work. Pair it with per-expert token counts because one scalar can hide which experts are hot. Use router z-loss to answer whether logits are drifting along a direction that softmax probabilities ignore. Pair it with the log-sum-exp distribution rather than only the batch mean. Keep the coefficients in configuration and the unweighted losses in logs. A weighted term can look quiet simply because its coefficient is small. Test common-offset invariance as a unit check. Adding the same constant to all logits must leave routing probabilities and load balancing loss unchanged, while z-loss must change. Test the two counterexamples as another check. They catch implementations that accidentally compute z-loss after softmax or assume a balanced batch guarantees small logits. The[AI research course](/ai-research)connects these loss functions to router optimization and training diagnostics. The operational rule is narrower: utilization and logit scale need separate measurements. --- This representation contains public Fanout content only. Protected Pro lessons, account data, billing, checkout, and pricing are not included. Browse the public content map: https://fanout.sh/sitemap.md