Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Top-K Routing › Worked Example

Michael BrenndoerferNovember 15, 20256 min read

Language AI HandbookAll chapters

Part 33 of 67Mixture of Experts

Chapter 242Top-K Routing › Worked Example

Reading section 7 of 12 · Worked Example · 6 minSection 7 of 12 · 6 min

Section 7 of 12 · 6 min read

Top-K Routing › Worked ExampleOne section · Sections in this chapter7 / 12 ↓

Sections in this chapter12

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

To solidify understanding, let's trace through a concrete numerical example of top-2 routing from beginning to end. This example uses a small but fully realistic setup: 4 experts, a hidden dimension of 4, and a single token. Every number is computed explicitly so you can verify each step by hand.

The token has a hidden representation $\mathbf{x}$ with 4 dimensions. The router weight matrix $W_r$ has shape $4 \times 4$ (one row per expert, one column per hidden dimension). We set these to specific values so we can compute routing logits $\mathbf{h} = W_r \mathbf{x}$ by hand:

In[17]:

Code

# Token representation
x = torch.tensor([[0.5, -0.3, 0.8, 0.1]])

# Router weights (4 experts, hidden dim 4)
W_r = torch.tensor(
    [
        [0.2, -0.1, 0.4, 0.3],  # Expert 0 weights
        [-0.3, 0.5, 0.1, -0.2],  # Expert 1 weights
        [0.1, 0.2, -0.3, 0.6],  # Expert 2 weights
        [0.4, -0.4, 0.2, 0.1],  # Expert 3 weights
    ]
)

# Compute router logits: h = x @ W_r.T
router_logits = x @ W_r.T

Out[18]:

Console

Token representation x: [0.5, -0.30000001192092896, 0.800000011920929, 0.10000000149011612]

Router logits h = x @ W_r.T:
  Expert 0: 0.4800
  Expert 1: -0.2400
  Expert 2: -0.1900
  Expert 3: 0.4900

Let's verify the first logit manually. The dot product of the token $[0.5, -0.3, 0.8, 0.1]$ with expert 0's weights $[0.2, -0.1, 0.4, 0.3]$ is:

h_0 = (0.5)(0.2) + (-0.3)(-0.1) + (0.8)(0.4) + (0.1)(0.3)

\begin{aligned} h_0 &= 0.10 + 0.03 + 0.32 + 0.03 \\ &= 0.48 \end{aligned}

This confirms that expert 0's representation is moderately compatible with this token's content. The router logits for the other experts similarly measure how well each expert's learned "query" aligns with the token's representation.

Now we select the top-2 experts:

In[19]:

Code

# Select top-2
top_k_vals, top_k_idx = torch.topk(router_logits, k=2, dim=-1)

Out[20]:

Console

Top-2 expert indices: [3, 0]
Top-2 logit values: [0.49000000953674316, 0.48000001907348633]

We can see that experts 3 (with logit approximately 0.49) and expert 0 (with logit approximately 0.48) have the highest routing scores. These two experts "won" the routing competition for this token. Notice that the margin between the winners and the third-place expert is relatively small in this case, meaning slightly different token representations could flip the routing decision. This sensitivity to input content is exactly what the router is supposed to learn: distinguishing which expert best handles the current input.

Now we compute the normalized gating weights by applying softmax only to the two selected logits:

g_1 = \frac{\exp(h_{k_1})}{\exp(h_{k_1}) + \exp(h_{k_2})}, \quad g_2 = \frac{\exp(h_{k_2})}{\exp(h_{k_1}) + \exp(h_{k_2})}

In[21]:

Code

# Normalize weights over selected experts only
normalized_weights = F.softmax(top_k_vals, dim=-1)

Out[22]:

Console

Normalized weights: [0.5024999976158142, 0.4975000321865082]
Sum of weights: 1.0000

Interpretation:
  Expert 3 gets 50.2% weight
  Expert 0 gets 49.8% weight
  Experts 1 and 2 are not used (0% weight)

Out[23]:

Visualization

Bar chart of router logits for 4 experts, with Expert 0 and Expert 3 highlighted in green as the top-2 selected experts. — Router logits for the example token across 4 experts. Experts 0 and 3 (green) score highest and are selected, while experts 1 and 2 (gray) are bypassed.

Bar chart showing binary selection mask with Experts 0 and 3 selected (height 1) and Experts 1 and 2 unselected (height 0). — Binary selection mask for the top-2 routing decision. Experts 0 and 3 receive a selection indicator of 1, while the remaining experts receive 0.

Bar chart of final routing weights with Experts 0 and 3 each receiving approximately 0.50 weight and the other experts receiving zero. — Final normalized gating weights. Because experts 0 and 3 have nearly equal logits, their weights are approximately 0.50 each, showing the token benefits almost equally from both experts.

The routing has decided that expert 3 should contribute roughly 50.2% and expert 0 should contribute roughly 49.8% to this token's output. Experts 1 and 2 are completely bypassed: they compute nothing for this token and receive no gradient signal from it during training.

Because the two selected logits were very close in value (approximately 0.49 and 0.48), the softmax produces nearly equal weights. This near-equal split means both experts contribute roughly equally to the final representation. If we denote expert 3's output as $\mathbf{e}_3$ and expert 0's output as $\mathbf{e}_0$ , the final combined output is:

\mathbf{y} \approx 0.502 \cdot \mathbf{e}_3 + 0.498 \cdot \mathbf{e}_0

where:

$\mathbf{y}$ : the final combined output vector for this token
$\mathbf{e}_3, \mathbf{e}_0$ : the output vectors computed by expert 3 and expert 0 respectively when given $\mathbf{x}$ as input

This nearly equal weighting means both experts contribute almost equally to the final representation. As a contrast, if the logits had been 0.9 for expert 3 and 0.2 for expert 0, the softmax would produce weights of approximately $\frac{e^{0.9}}{e^{0.9} + e^{0.2}} \approx 0.67$ and $0.33$ , giving expert 3 twice as much influence. The further apart the logits, the more the combination resembles a single-expert decision. The closer together, the more it approximates an equal ensemble.

Why does this formula make sense? When two experts have similar routing scores, the model is essentially uncertain which expert is more appropriate. Giving them equal weight is a reasonable hedging strategy: rather than committing fully to one expert despite the uncertainty, the model takes a blend. When one expert strongly dominates, the model correctly places most weight on the clearly superior choice. The softmax over the selected pair implements this hedging automatically, without requiring any explicit uncertainty estimation.