Mixtral › Worked Example: Tracing a Token Through Mixtral

To make the abstract architecture concrete, let us trace a single token through one complete Mixtral MoE layer step by step. Suppose we are processing the token "gradient" appearing in the sentence "The gradient of the loss function guides training." After the attention layer has integrated context from surrounding tokens, we receive a hidden state vector representing this token in context.

We begin with a hidden state vector $\mathbf{x} \in \mathbb{R}^{4096}$. For our worked example, we will use a small toy version with $d_{\text{model}} = 4$ and $n_{\text{experts}} = 4$ with top-2 routing, to keep the numbers tractable by hand.

Step 1: Router computes logits.

The router weight matrix $W_g \in \mathbb{R}^{4 \times 4}$ maps our hidden state to a score for each of 4 experts. Suppose:

$$\mathbf{x} = [1.2, -0.3, 0.8, 0.5]$$

and the router produces logits:

$$\mathbf{s} = \mathbf{x} W_g = [2.1, 0.4, 1.7, 0.3]$$

where:

• $\mathbf{s} \in \mathbb{R}^{n_{\text{experts}}}$: the raw router logits, one per expert
• $s_0 = 2.1$: Expert 0 scores highest, likely a math or technical expert given our context
• $s_2 = 1.7$: Expert 2 scores second, perhaps a science/reasoning expert
• $s_1 = 0.4, s_3 = 0.3$: Experts 1 and 3 score much lower

Step 2: TopK masking selects exactly 2 experts.

We keep the top-2 scores and set the rest to $-\infty$:

$$\mathbf{s}_{\text{masked}} = [2.1, -\infty, 1.7, -\infty]$$

Only experts 0 and 2 are eligible to receive this token.

Step 3: Softmax converts masked scores to routing weights.

$$\begin{aligned} g_0 &= \frac{e^{2.1}}{e^{2.1} + e^{1.7}} = \frac{8.166}{8.166 + 5.473} = \frac{8.166}{13.639} \approx 0.599 \\ g_2 &= \frac{e^{1.7}}{e^{2.1} + e^{1.7}} = \frac{5.473}{13.639} \approx 0.401 \end{aligned}$$

where:

• $g_0 \approx 0.599$: Expert 0 receives 60% of the weight
• $g_2 \approx 0.401$: Expert 2 receives 40% of the weight
• $g_0 + g_2 = 1.000$: Weights sum to exactly 1 by construction of softmax over the selected pair

Step 4: Selected experts process the input.

Expert 0 applies its SwiGLU transformation to $\mathbf{x}$, creating output $\mathbf{e}_0 \in \mathbb{R}^{4}$. Expert 2 independently applies its SwiGLU transformation to the same $\mathbf{x}$, creating $\mathbf{e}_2 \in \mathbb{R}^{4}$. These experts have different learned weight matrices, so they produce different outputs even from the same input.

For our toy example, suppose:

$$\mathbf{e}_0 = [0.8, 1.2, -0.3, 0.6] \quad \text{(Expert 0 output)}$$ $$\mathbf{e}_2 = [0.3, 0.7, 0.9, -0.2] \quad \text{(Expert 2 output)}$$

Step 5: Weighted combination produces the MoE output.

$$\begin{aligned} \text{MoE}(\mathbf{x}) &= g_0 \cdot \mathbf{e}_0 + g_2 \cdot \mathbf{e}_2 \\ &= 0.599 \cdot [0.8, 1.2, -0.3, 0.6] + 0.401 \cdot [0.3, 0.7, 0.9, -0.2] \\ &= [0.479, 0.719, -0.180, 0.359] + [0.120, 0.281, 0.361, -0.080] \\ &= [0.599, 1.000, 0.181, 0.279] \end{aligned}$$

where the final output is a convex combination of the two expert outputs, with Expert 0's contribution weighted more heavily because it received the higher router score.

Step 6: Residual connection adds the MoE output back.

The MoE output is added to the pre-normalization residual:

$$\mathbf{h}_{\text{out}} = \mathbf{x}_{\text{residual}} + \text{MoE}(\text{LayerNorm}(\mathbf{x}_{\text{residual}}))$$

This residual connection ensures the MoE layer learns incremental refinements rather than complete replacements, which is why pre-training deep networks with residual connections is so much more stable than without them.

This six-step process repeats independently for every token in the sequence and for every layer in the 32-layer network, with each layer having its own independent set of 8 experts and its own router. The full forward pass through Mixtral involves 32 such routing decisions per token, with each decision potentially selecting different pairs of experts from the available 8.