Gating Networks › Routing Score Computation

Raw routing logits from the linear projection must be converted into routing weights that determine how much each expert contributes to the output. The logits can be any real number, positive or negative, and they don't have any inherent probabilistic meaning. This transformation from arbitrary scores to real weights happens through a softmax function, creating a probability distribution over experts that we can interpret and use.

The softmax function is the standard tool for converting a vector of arbitrary real numbers into a probability distribution. In the context of routing, it turns our expert logits into normalized routing weights:

$$p(e_i|x) = \frac{\exp(G(x)_i)}{\sum_{j=1}^{n} \exp(G(x)_j)}$$

where:

• $p(e_i|x)$: the probability of selecting expert $i$ for input $x$
• $G(x)_i$: the raw routing score (logit) for expert $i$
• $\exp(G(x)_i)$: the exponential function applied to the logit, so the value is positive
• $\sum_{j=1}^{n} \exp(G(x)_j)$: the sum of exponentials across all $n$ experts, serving as the normalizing constant

Let's see what this formula does. The exponential function serves two purposes: first, it ensures all values are positive (since probabilities cannot be negative); second, it amplifies differences between logits. If one expert has a logit of 3 and another has 1, the exponential makes the first about $e^2 \approx 7.4$ times larger, not just 2 points higher. This amplification makes the router's preferences more decisive.

The denominator is a normalizing constant. This keeps all probabilities sum to exactly 1. This normalization is needed because we want the routing weights to stand for a proper probability distribution.

These routing probabilities $p(e_i|x)$ stand for the router's "confidence" that expert $i$ is appropriate for input $x$. Higher probabilities indicate stronger preference, and we can compare probabilities across experts to understand the router's relative confidence in each option.

In a soft MoE (which we introduced conceptually in the sparse models chapter), the final output would be a weighted combination of all expert outputs:

$$y = \sum_{i=1}^{n} p(e_i|x) \cdot E_i(x)$$

where:

• $y$: the combined output of the MoE layer
• $p(e_i|x)$: the routing probability for expert $i$
• $E_i(x)$: the output computed by expert $i$ for input $x$

However, this soft combination requires computing all expert outputs, which defeats the purpose of sparsity. If we must run every expert for every token, we gain no computational savings from having multiple specialized experts rather than one large one.

Practical MoE systems use hard routing to obtain the computational benefits of sparsity: only the top-$K$ experts (by routing probability) are computed. The next chapter covers Top-K routing in detail, but the basic idea is needed to understand how the router functions in practice.

The process works in four steps:

1. Compute all routing probabilities via softmax
2. Select the $K$ highest-scoring experts
3. Renormalize selected probabilities to sum to 1
4. Compute only those $K$ experts

For example, with 8 experts and $K=2$, only 2 experts run per token. This gives roughly 75% computational savings in the expert portion of the layer. The routing weights for selected experts become:

$$\tilde{p}(e_i|x) = \frac{p(e_i|x)}{\sum_{j \in \text{TopK}} p(e_j|x)}$$

where:

• $\tilde{p}(e_i|x)$: the renormalized weight for selected expert $i$
• $p(e_i|x)$: the original softmax probability for expert $i$
• $\text{TopK}$: the set of indices corresponding to the $K$ experts with highest scores
• $\sum_{j \in \text{TopK}} p(e_j|x)$: the sum of probabilities for the selected experts, used as a normalizing constant

This renormalization ensures the weighted combination of selected experts produces properly scaled outputs. Without renormalization, if the two selected experts had probabilities 0.3 and 0.2, they would sum to only 0.5, effectively scaling down the output. Renormalization adjusts them to approximately 0.6 and 0.4. This keeps the output magnitude remains appropriate.

Some implementations add temperature scaling to the softmax to give additional control over routing behavior:

$$p(e_i|x) = \frac{\exp(G(x)_i / \tau)}{\sum_{j=1}^{n} \exp(G(x)_j / \tau)}$$

where:

• $p(e_i|x)$: the temperature-scaled probability for expert $i$
• $G(x)_i$: the raw routing logit
• $\tau$: the temperature parameter that controls the sharpness of the distribution
• $\exp(G(x)_i / \tau)$: the exponentiated scaled logit; using $\tau < 1$ amplifies differences, making the distribution sharper

This concept relates to the decoding temperature we covered in Part XXVIII, where temperature controls the randomness of text generation. The principle is the same: temperature modulates how peaked or flat the probability distribution is.

Lower temperatures sharpen the distribution, pushing probability mass toward the highest-scoring expert and approaching hard assignment to a single expert. Higher temperatures soften the distribution, spreading probability more evenly and approaching uniform routing where every expert has roughly equal chance of selection.