Gating Networks › Router Training

Training the router presents a unique challenge that distinguishes it from standard neural network components: the router's output involves a discrete selection (which experts to use), but backpropagation requires continuous, differentiable operations. Discrete choices create discontinuities in the loss surface, making gradients undefined at decision boundaries. How, then, do gradients flow through routing decisions?

The key insight that lets router training is that although expert selection is discrete, the routing weights are continuous. The model cannot backpropagate through the binary decision of "use expert 3 instead of expert 5," but it can backpropagate through the continuous weights applied to the selected experts. Gradients flow through the weights applied to selected experts:

$$y = \sum_{i \in \text{TopK}} \tilde{p}(e_i|x) \cdot E_i(x)$$

where:

• $y$: the output of the MoE layer
• $\tilde{p}(e_i|x)$: the renormalized routing weight for selected expert $i$
• $E_i(x)$: the output computed by expert $i$
• $\text{TopK}$: the set of indices for the selected experts

During backpropagation, the loss gradient with respect to router logits comes from two sources:

1. Through the weights: How changing $\tilde{p}(e_i|x)$ affects the weighted combination of expert outputs
2. Through the softmax: How the raw logits affect the routing probabilities

For a selected expert $i$, the gradient through routing weights follows the chain rule:

$$\frac{\partial \mathcal{L}}{\partial G(x)_i} = \frac{\partial \mathcal{L}}{\partial y} \cdot \frac{\partial y}{\partial \tilde{p}(e_i|x)} \cdot \frac{\partial \tilde{p}(e_i|x)}{\partial G(x)_i}$$

where:

• $\frac{\partial \mathcal{L}}{\partial G(x)_i}$: the gradient of the loss with respect to the router logits
• $\frac{\partial \mathcal{L}}{\partial y}$: the gradient flowing back from the network output
• $\frac{\partial y}{\partial \tilde{p}(e_i|x)}$: the sensitivity of the output to the expert's weight
• $\frac{\partial \tilde{p}(e_i|x)}{\partial G(x)_i}$: the gradient of the softmax (or Top-K) function

This gradient chain has an intuitive interpretation. The first term captures how much the overall loss depends on this MoE layer's output. The second term captures how much changing the weight of a particular expert would change that output (in effect the expert's output vector). The third term captures how sensitive the routing weight is to changes in the raw logit.

This gradient updates the router to assign higher weights to experts that produce better outputs for each input. If an expert produces an output that reduces the loss, the gradient will adjust the router to give that expert a higher weight for similar inputs in the future.

One subtlety deserves careful attention: gradients only flow through selected experts. If expert $j$ wasn't in the Top-K, there's no direct gradient signal about whether selecting it would have been better. The model never computes $E_j(x)$, so it cannot know how good that output would have been.

This creates a basic limitation in router learning. The router learns to improve its choices among selected experts based on their actual outputs, but has limited ability to discover that an unselected expert might be better. Imagine a scenario where expert 7 would be perfect for certain mathematical tokens, but the router never selects it for those tokens during early training. Without ever seeing expert 7's output on mathematical content, the gradient signal cannot tell the router to adjust toward selecting expert 7.

MoE systems can get stuck in suboptimal routing patterns. An expert might become permanently underutilized simply because random initialization didn't give it enough early opportunities.

Several techniques address this limitation:

• Load balancing losses penalize uneven expert usage, encouraging exploration of all experts by adding a penalty when some experts receive many more tokens than others
• Random routing occasionally selects experts not in the Top-K for gradient purposes, giving gradient signal about alternatives
• Auxiliary losses give additional training signals beyond the main task loss, guiding the router toward desirable behaviors

We'll explore these techniques in detail in the upcoming chapters on load balancing and auxiliary losses.

The routing weights are typically small values (especially when split across multiple experts), which can lead to gradient scaling issues. If two experts share a token with weights 0.5 and 0.5, the gradient flowing to each expert is scaled by 0.5. With more experts sharing, these weights become even smaller.

Beyond gradient scaling, another concern is numerical stability. If routing logits grow very large, the softmax can produce extreme probabilities (nearly 0 or 1), leading to vanishing gradients and unstable training. Some implementations use router z-loss to prevent routing logits from growing too large, maintaining stable training dynamics. The z-loss adds a penalty proportional to the log-sum-exp of routing logits:

$$\mathcal{L}_z = \frac{1}{B}\sum_{x} \left(\log \sum_{i=1}^{n} \exp(G(x)_i)\right)^2$$

where:

• $\mathcal{L}_z$: the router z-loss
• $B$: the batch size
• $x$: an input token representation in the batch
• $G(x)_i$: the raw routing score (logit) for expert $i$
• $\sum_{i=1}^{n} \exp(G(x)_i)$: the partition function (denominator of softmax) for input $x$
• $\log \sum_{i=1}^{n} \exp(G(x)_i)$: the Log-Sum-Exp of the logits, serving as a smooth approximation of the maximum logit

The log-sum-exp quantity is closely related to the maximum logit: as one logit grows much larger than the others, the log-sum-exp approaches that maximum value. By penalizing the squared log-sum-exp, the z-loss discourages any single expert from receiving extremely high routing scores. This encourages the router to produce moderate logits rather than extreme values, improving numerical stability throughout training.