Switch Transformer › Worked Example: Step-by-Step Routing Trace

To fully understand how a single token moves through a Switch layer, let's trace through the complete computation manually with concrete numbers. This numerical walkthrough makes the abstract formulas concrete and reveals exactly what happens at each step.

Consider a simple setting: we have 4 experts ($E = 4$), a model dimension of $d_{\text{model}} = 4$ (tiny, for illustration), and a batch with 8 tokens total. The capacity factor is $C = 1.5$.

Step 1: Compute capacity. We apply the capacity formula:

$$\text{capacity} = \left\lfloor \frac{8 \times 1.5}{4} \right\rfloor = \left\lfloor 3.0 \right\rfloor = 3$$

Each expert can accept at most 3 tokens before it is considered full and subsequent tokens are dropped.

Step 2: Router produces logits. For one specific token with representation $x = [1.2, -0.3, 0.8, 0.5]$, the router's linear layer produces logits $z = W_r \cdot x$. Suppose this gives $z = [2.1, 0.5, 1.8, -0.2]$.

Step 3: Softmax converts logits to probabilities. We compute:

$$\begin{aligned} e^{2.1} &= 8.17 \\ e^{0.5} &= 1.65 \\ e^{1.8} &= 6.05 \\ e^{-0.2} &= 0.82 \end{aligned}$$

The sum of exponentials is $8.17 + 1.65 + 6.05 + 0.82 = 16.69$. Dividing each by the sum:

$$p(x) = \left[ \frac{8.17}{16.69}, \frac{1.65}{16.69}, \frac{6.05}{16.69}, \frac{0.82}{16.69} \right] = [0.49, 0.10, 0.36, 0.05]$$

Step 4: Select the top expert. The highest probability is $p_0 = 0.49$ for expert 0, so $e^* = 0$. This token will be processed by expert 0.

Step 5: Check capacity and route. We check whether expert 0 has room. If fewer than 3 tokens have already been assigned to expert 0, we process this token. The expert runs its FFN on the token representation and returns output $E_0(x)$.

Step 6: Scale by router probability. The Switch layer output for this token is:

$$y = 0.49 \times E_0(x)$$

The factor 0.49 scales the output. This shows the router's moderate confidence. If the router were completely certain (probability 1.0), the expert output would pass through unchanged. If the router were uncertain (probability near 0.25 for each expert), the output would be substantially attenuated.

Step 7: Add residual. The final output adds the original input: $y_{\text{final}} = 0.49 \times E_0(x) + x$. This residual connection ensures the token retains its full original representation even if the expert's contribution is small.

This trace reveals why the router probability matters as a multiplicative weight: it gives a differentiable path for gradients from the output loss back to the router's linear projection weights, letting end-to-end training. Without this multiplicative weighting, the routing decision would be a non-differentiable argmax with no gradient signal for the router.