Router Z-Loss › The Router Instability Problem

To understand why router instability occurs, we first need to revisit how routing decisions are made and why the underlying mathematics creates a hidden vulnerability. Recall from the discussion of gating networks that the router produces logits $z_i = W_r h_i$ for each input token $h_i$, where $W_r$ is the router weight matrix. These logits represent the raw, unnormalized scores that indicate how strongly the router prefers each expert for a given token. These logits pass through a softmax to produce routing probabilities:

$$p_j = \frac{\exp(z_j)}{\sum_{k=1}^{N} \exp(z_k)}$$

where:

• $p_j$: the probability assigned to expert $j$
• $z_j$: the router logit for expert $j$
• $N$: the total number of experts
• $\exp(z_j)$: the exponential of the logit, so a positive value
• $\sum_{k=1}^{N} \exp(z_k)$: the sum over all $N$ experts, serving as the normalization constant

The softmax function turns arbitrary real-valued logits into a proper probability distribution where all values are positive and sum to one. This transformation is elegant and widely used, but it introduces subtle numerical challenges when the input logits span extreme ranges.

During training, nothing explicitly constrains the magnitude of these logits. The neural network is free to produce logits of any size, and the gradient descent process will adjust parameters in whatever direction minimizes the loss function. The softmax function is scale-invariant in its output probabilities: multiplying all logits by a constant doesn't change the relative probabilities. If you double all the logits, you still get the same ranking and the same probability assignments. However, scale dramatically affects the sharpness of the distribution and the numerical stability of the computation. This simple property creates a hidden danger that only appears as training progresses.

It is worth pausing on this scale-invariance property because it has important consequences. The router's job, from a functional perspective, is to rank experts and select the best ones. Rankings are fully determined by the relative order of logits, not their absolute values. Yet the training loss is sensitive to the absolute values through the softmax computation and its gradients. This mismatch between what the router needs (relative ranking) and what the optimization process adjusts (absolute values) creates the conditions for unconstrained logit growth.

As training progresses, the router learns to route tokens to appropriate experts by increasing the logits for preferred experts relative to others. This is exactly what we want: the router should learn to distinguish which expert is best suited for each type of input. Without constraints, however, this creates a natural logit drift toward larger logit magnitudes. If expert 3 is consistently preferred for certain tokens, the router learns to make $z_3$ large and positive while other logits become large and negative. The optimization process finds that making these differences more extreme leads to more confident routing decisions.

This drift accelerates over time through a self-reinforcing feedback loop. Larger logits create sharper probability distributions, which produce stronger gradients for the "winning" experts. When the probability assigned to the chosen expert approaches 1.0, the gradient signal becomes very clear: this expert should handle this token. These stronger gradients further increase the logit gap, which creates even sharper distributions, which produce even stronger gradients. This positive feedback loop pushes logits toward extreme values without any natural stopping point.

To illustrate this phenomenon concretely, consider a router that starts with logits around $[1.0, 0.8, 0.5, 0.2]$. The corresponding probabilities might be roughly $[0.35, 0.29, 0.21, 0.15]$, which is a reasonably soft distribution. After many training steps, if expert 1 is consistently correct, the logits might evolve to $[10.0, 3.0, 2.0, 1.0]$, giving probabilities closer to $[0.99, 0.009, 0.0003, 0.0001]$. The training process has learned to be more confident, but the logit magnitudes have increased dramatically. Left unchecked, this drift continues until logits reach values that cause computational problems.

Think of logit drift as a ratchet mechanism: each training step tightens the routing distribution slightly, and there is no counter-pressure to loosen it. Every time the model correctly identifies an expert assignment, the gradient pushes the logits further apart, reinforcing the decision. The model is rewarded for confidence and responds by becoming ever more confident, regardless of whether the extra confidence improves task performance. The problem is unbounded confidence that eventually strains numerical computations.

Large logits cause several numerical problems that can derail training entirely. These issues show why z-loss targets logit magnitudes:

• Softmax overflow: Computing $\exp(z_j)$ for $z_j > 709$ (float64) or $z_j > 88$ (float32) produces infinity. When any exponential term overflows, the entire softmax computation becomes undefined. Even a single overflowed value corrupts the probability calculation for all experts.

• Softmax underflow: Computing $\exp(z_j)$ for very negative $z_j$ produces zero due to floating-point underflow. This might seem harmless since small probabilities should be near zero anyway. However, when combined with division operations, these zeros can cause division-by-zero errors or produce undefined gradients.

• Gradient instability: The softmax gradient involves terms like $p_j(1 - p_j)$, which vanishes as probabilities approach 0 or 1. When the router becomes extremely confident, assigning probability 0.9999 to one expert and 0.0001 to others, the gradients become tiny. This gradient saturation prevents the router from learning to change its routing decisions even when it should.

• Loss of precision: When logits span a huge range, floating-point arithmetic loses significant digits. A computation involving $\exp(50)$ and $\exp(-50)$ cannot represent both values accurately in the same calculation, even with double precision. This precision loss introduces numerical noise that accumulates over training steps.

While the softmax-with-max-subtraction trick (subtracting $\max_j z_j$ before exponentiation) prevents overflow in the forward pass, it does not address the instability that large logit magnitudes create for gradient flow. The backward pass still suffers from saturated gradients and precision issues. The max-subtraction trick only helps with the exponential computation; it does not prevent the underlying logit drift that causes these problems. Z-loss addresses this root cause rather than individual symptoms.

The basic issue is that the max-subtraction trick is a computational workaround, not a solution to the underlying problem. You can stably compute $\text{softmax}([1000, 999, 998])$ using max-subtraction: subtract 1000 from all values, then compute $\text{softmax}([0, -1, -2])$. But the gradient of this softmax with respect to the original logits is still nearly zero, because the probabilities are $[0.731, 0.269, 0.099]$, and the gradient involves these probabilities directly. The router has essentially made a near-deterministic decision, and near-deterministic decisions have tiny gradients. Max-subtraction makes the computation numerically stable without which makes it informationally useful.

Router instability typically manifests in several recognizable ways that you should monitor:

1. Training loss spikes: Sudden, unexpected increases in loss occur as numerical issues cascade through the network. The loss might be decreasing smoothly for thousands of steps, then suddenly jump to a much higher value or even become NaN. These spikes often indicate that router logits have grown large enough to cause computational problems.

2. Gradient explosions: NaN or inf values appearing in gradients signal that the numerical instability has propagated to the backward pass. Once gradients contain invalid values, parameter updates become meaningless, and the model parameters can become corrupted. Recovering from gradient explosions often requires rolling back to an earlier checkpoint.

3. Expert collapse: The router becoming stuck on a single expert due to saturated gradients represents a subtle but serious failure mode. When logits become extreme, the router cannot learn to route tokens to different experts even when the current routing is suboptimal. The gradients for non-selected experts become vanishingly small, freezing the routing pattern in place.

4. Poor generalization: Overconfident routing that doesn't transfer to new inputs indicates that the model has memorized routing patterns rather than learning generalizable features. When the router makes extremely sharp decisions based on subtle input features, small changes in input distribution can completely change routing behavior in unpredictable ways.

These issues become more severe as model scale increases. Larger models have more parameters in the router network, more capacity to push logits to extreme values, and longer training runs that provide more opportunity for drift to accumulate. The ST-MoE paper from Google observed that without z-loss, larger MoE models frequently experienced training instabilities that smaller models avoided. This scale-dependent behavior makes z-loss increasingly important as the field moves toward ever-larger MoE architectures.

These failures often appear suddenly after training seems to be going well. A common failure pattern is a long period of healthy, decreasing loss followed by an abrupt instability that wipes out gains from thousands of training steps. Logit drift is slow and cumulative: logits grow gradually over many batches, but the numerical consequences trigger only when they exceed specific thresholds. This makes preventive regularization, like z-loss, more effective than reactive measures such as checkpoint rollback.