Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Mixtral › Expert Layer Design

Michael BrenndoerferUpdated December 31, 202520 min read

Language AI HandbookAll chapters

Part 33 of 67Mixture of Experts

Chapter 248Mixtral › Expert Layer Design

Reading section 3 of 15 · Expert Layer Design · 20 minSection 3 of 15 · 20 min

Section 3 of 15 · 20 min read

Mixtral › Expert Layer DesignOne section · Sections in this chapter3 / 15 ↓

Sections in this chapter15

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

Mixtral's expert design follows the principles we established in the Expert Networks chapter, with each expert implemented as a standard SwiGLU feed-forward network. This architecture uses gating to selectively activate different parts of the feed-forward computation. The intuition is straightforward: rather than passing every feature through the same transformation, a gated architecture lets the network learn which features matter most for a given input, emphasize those, and suppress others. This selective activation is more expressive than a simple linear transformation followed by a nonlinearity.

Think of each expert as a specialist who has developed a personal toolkit for handling certain kinds of inputs. When you route a token about Python list comprehensions to a code-focused expert, that expert applies its learned transformations for processing programming concepts. When you route a token from a French sentence to a language-fluency expert, a different set of transformations applies. But critically, neither expert needs to know about the other's domain: each operates independently on its assigned inputs, creating an output vector in the same representational space so the results can be combined. The expertise is encoded in the weights of each expert's projection matrices, which were learned during training to excel at their respective specializations.

SwiGLU is well-suited for experts because its gating mechanism lets each expert make soft decisions about which input features are relevant to its domain. An expert specialized in mathematical reasoning might learn to gate on features that encode numerical context, suppressing linguistically relevant features that another expert handles better. This adaptive feature selection lets each expert process the same input differently, producing diverse transformations for the weighted output combination.

Building on our discussion of Gated Linear Units from Part XV, each expert computes a gated transformation of its input. The SwiGLU activation uses gating to control information flow. Each expert transforms the input through three learned projection matrices, combining a gating mechanism with element-wise multiplication to selectively activate features. Given an input vector $x$ , expert $i$ computes a two-branch structure. One branch determines which features to emphasize (the gate), while the other gives the features to be gated (the up projection). The element-wise product of these two branches creates a filtered representation that is then projected back to the original dimension.

In[4]:

Code

%pip install torch
import torch
import torch.nn.functional as F

torch.manual_seed(42)

d_model = 4096
d_ffn = 14336

W_gate = torch.randn(d_model, d_ffn) * 0.01
W_up = torch.randn(d_model, d_ffn) * 0.01
W_down = torch.randn(d_ffn, d_model) * 0.01

x = torch.randn(d_model)

gate_output = F.silu(x @ W_gate)
up_output = x @ W_up
expert_output = (gate_output * up_output) @ W_down

Out[5]:

Console

Input shape: torch.Size([4096])
Gate output shape: torch.Size([14336])
Up output shape: torch.Size([14336])
Expert output shape: torch.Size([4096])
Expert output magnitude (mean abs): 0.2167

The computation processes a single token's hidden state through the gating mechanism. The input shape (4096) expands to the FFN intermediate dimension (14336) for the gate and up projections, then contracts back to the original dimension (4096) in the output. This dimension preservation ensures compatibility with subsequent transformer layers while letting the expert to compute in a higher-dimensional space where features can be more expressively transformed. The expansion to a higher dimension is important. It gives the network with more degrees of freedom to learn complex feature interactions before compressing the representation back to the original size.

The formal mathematical expression for each expert's computation can now be stated precisely. The formula captures the three-step process we described above: project to a higher dimension with gating, apply element-wise multiplication to create the gated representation, and project back to the model dimension:

\text{Expert}_i(x) = (\text{SiLU}(xW_{\text{gate}}^{(i)}) \odot xW_{\text{up}}^{(i)})W_{\text{down}}^{(i)}

where:

$x \in \mathbb{R}^{d_{\text{model}}}$ : the input hidden state vector for a single token
$W_{\text{gate}}^{(i)} \in \mathbb{R}^{d_{\text{model}} \times d_{\text{ffn}}}$ : gate projection matrix for expert $i$ . Determines which features to emphasize.
$W_{\text{up}}^{(i)} \in \mathbb{R}^{d_{\text{model}} \times d_{\text{ffn}}}$ : up projection matrix for expert $i$ . Expands the representation to the intermediate dimension.
$W_{\text{down}}^{(i)} \in \mathbb{R}^{d_{\text{ffn}} \times d_{\text{model}}}$ : down projection matrix for expert $i$ . Projects back to the model dimension.
$\text{SiLU}(z) = z \cdot \sigma(z)$ : the Sigmoid Linear Unit activation function, where $\sigma(z) = \frac{1}{1 + e^{-z}}$ is the sigmoid function
$\sigma(z) = \frac{1}{1 + e^{-z}}$ : sigmoid function, squashing any real number to the range (0, 1) and giving smooth gating
$\odot$ : element-wise (Hadamard) multiplication
$d_{\text{model}} = 4096$ : the hidden dimension of the model
$d_{\text{ffn}} = 14336$ : the intermediate dimension of the expert feed-forward network

Out[6]:

Visualization

Line chart showing SiLU (blue solid) and ReLU (red dashed) activation functions over the range -5 to 5. SiLU curves smoothly below zero letting small negative values, while ReLU has a hard kink at zero. Annotations highlight the smooth transition at the origin and SiLU's negative value region. — Comparison of SiLU and ReLU activation functions across the input range [-5, 5]. SiLU has a smooth derivative through the origin and retains small negative outputs, while ReLU has a non-differentiable kink at zero and clips all negative inputs to zero.

The formula implements a gated activation pattern that deserves careful attention. The SiLU activation applied to the gate projection creates a soft gating mechanism, letting the network to learn which features to pass through based on the input. Unlike a hard gate that produces binary on/off decisions, SiLU produces smooth values that can range continuously, letting gradient-based learning to refine gating behavior during training. The element-wise multiplication combines the gated features with the up-projected representation, effectively scaling each dimension of the up projection by the corresponding gate value. Finally, the down projection returns to the original dimensionality, compressing the gated representation back to a form compatible with the rest of the transformer.

Having established how individual experts transform their inputs, we now turn to the necessary question of which experts should process each token. The routing network determines this assignment by computing a score for each expert and selecting the top two. This routing mechanism is central to what makes Mixtral a sparse model. Instead of processing every token through all parameters, the router dynamically selects a small subset of experts based on the token's representation.

The routing function maps hidden states to expert weights, selecting exactly two experts per token and assigning them normalized weights that sum to 1. The computation proceeds in three steps. First, a linear projection computes raw scores for all experts, measuring how suitable each expert is for the given input. Second, TopK masking sets all but the two highest scores to negative infinity. This keeps only two experts can be selected. Third, softmax normalizes these masked scores into probabilities that sum to one. The routing function computes:

g(x) = \text{Softmax}(\text{TopK}(xW_g, k=2))

where:

$x \in \mathbb{R}^{d_{\text{model}}}$ : the input hidden state vector for a single token
$W_g \in \mathbb{R}^{d_{\text{model}} \times n_{\text{experts}}}$ : the routing weight matrix. Projects hidden states to expert scores, with $n_{\text{experts}} = 8$ .
$xW_g \in \mathbb{R}^{n_{\text{experts}}}$ : the raw router logits. One score per expert.
$\text{TopK}(z, k=2)$ : operation that sets all but the two highest values in vector $z$ to $-\infty$ . Ensures only the top-2 experts can receive non-zero routing weights.
$\text{Softmax}(z)_i = \frac{e^{z_i}}{\sum_{j=1}^{n_{\text{experts}}} e^{z_j}}$ : normalization function that converts the masked logits into probabilities
- $e^{z_i}$ : exponential of the $i$ -th masked logit. Ensures all values become positive.
- $\sum_{j=1}^{n_{\text{experts}}} e^{z_j}$ : sum of exponentials across all $n_{\text{experts}}$ positions. Is a normalizing constant so the output probabilities sum to 1.
$g(x) \in \mathbb{R}^{n_{\text{experts}}}$ : the final routing weights, where exactly two elements are non-zero and sum to 1

Key Insight: TopK Before Softmax

The ordering of operations, TopK masking before softmax, is not arbitrary. If softmax were applied first, all experts would receive non-zero weights (softmax never produces exactly zero outputs), and the model would always activate all experts. Sparsity would be purely aspirational rather than guaranteed. By masking before softmax, we enforce hard sparsity: exactly $k$ experts receive non-zero weight and the rest receive exactly zero. The router still learns to produce higher scores for better-matched experts (because the top-2 selected experts compete for weight through softmax), but the architecture guarantees computational sparsity regardless of how the router's scores are distributed.

The ordering of operations is necessary to understanding this function. TopK masking happens before softmax. This keeps exactly two experts are selected regardless of the score distribution. Without this masking, softmax would distribute weight across all experts, defeating the sparsity goal entirely since every expert would receive some fraction of the computation. The exponential function in softmax amplifies differences between scores, so the two highest-scoring experts typically receive most of the weight even before masking. However, the explicit masking gives a hard constraint that guarantees exactly two non-zero weights. This makes the sparsity property architectural rather than emergent. This is a deliberate design choice because rather than hoping the softmax distribution concentrates on a few experts, the TopK operation enforces this concentration mechanically.

Out[7]:

Visualization

Bar chart of raw router logits for 8 experts. Experts 1 and 3 have the tallest bars at 1.5 and 1.2, highlighted in blue. A red dashed threshold line marks the top-2 cutoff. — Raw router logits across all experts. Experts 1 and 3 score highest (1.5 and 1.2), establishing candidates for top-2 selection.

Bar chart after TopK masking. Experts 1 and 3 retain their scores in blue; all inactive experts remain at zero, with one annotation noting that their masked logits are negative infinity. — TopK masking operation enforcing exactly two active experts. Scores for experts 1 and 3 are retained while all others are set to negative infinity.

Bar chart of final softmax routing weights. Only experts 1 and 3 have non-zero bars with weights 0.57 and 0.43 labelled, summing to 1.0. — Final routing weights after softmax normalization. Experts 1 and 3 receive weights of 0.57 and 0.43 (summing to 1.0), while others remain at zero due to TopK masking.

With routing weights determined, the final MoE layer produces its output by taking a weighted sum of the selected experts. Since the router selects exactly two experts per token, the computation aggregates their outputs in a straightforward manner. The formula computes a convex combination, where each expert processes the input independently, creating its own output vector, and the router weights (which sum to 1) determine how much each expert contributes to the final result. This blending allows the model to combine complementary knowledge from multiple experts rather than relying solely on one expert's output. For instance, a token related to "Python programming for data science" might route to both a code-focused expert and a math-focused expert, blending their perspectives.

The MoE output combines the selected experts using their routing weights:

\text{MoE}(x) = \sum_{i \in \text{TopK}(g(x))} g_i(x) \cdot \text{Expert}_i(x)

where:

$x \in \mathbb{R}^{d_{\text{model}}}$ : the input hidden state vector for a single token
$\text{TopK}(g(x)) \subset \{1, 2, ..., 8\}$ : the set of exactly two expert indices selected by the router
$g_i(x) \in [0, 1]$ : the normalized routing weight for expert $i$ from the softmax output. Represents the proportion of expert $i$ 's contribution.
$\text{Expert}_i(x) \in \mathbb{R}^{d_{\text{model}}}$ : the output vector produced by expert $i$ when applied to input $x$ .
$\sum_{i \in \text{TopK}(g(x))} g_i(x) = 1$ : routing weights sum to 1. Ensures the output magnitude is comparable to a single expert.
$\text{MoE}(x) \in \mathbb{R}^{d_{\text{model}}}$ : the final layer output, a weighted blend of two expert outputs

For Mixtral, this sum always contains exactly two terms because of the top-2 routing constraint. To make this concrete, consider a specific example: if experts 3 and 7 are selected with weights 0.6 and 0.4, the output becomes:

\text{MoE}(x) = 0.6 \cdot \text{Expert}_3(x) + 0.4 \cdot \text{Expert}_7(x)

where:

$0.6$ and $0.4$ : routing weights from $g(x)$ for experts 3 and 7, summing to 1
$\text{Expert}_3(x)$ and $\text{Expert}_7(x)$ : the output vectors from applying each expert's SwiGLU transformation to input $x$
The final output is a weighted average, so the magnitude stays consistent with single-expert outputs

The weighted average formulation has an important property. Because the weights sum to 1, the output magnitude remains stable regardless of which experts are selected. If weights did not sum to 1, the MoE layer's outputs would vary in scale depending on routing decisions, potentially destabilizing training and which makes it harder for subsequent layers to process the representations consistently.

This design differs from Switch Transformer, which we covered in the previous chapter, in one key aspect: Mixtral uses top-2 routing rather than top-1. The additional expert gives redundancy and smoother gradient flow during training, though it doubles the compute per token compared to Switch's more aggressive sparsity.

Why Top-2?Link Copied

The choice of k=2 is a deliberate tradeoff between efficiency and quality. This choice illuminates the basic design tensions in sparse architectures. Before examining the benefits, it is worth appreciating that the choice of $k$ affects training and inference, with further consequences for deployment that are not obvious from the formula alone.

Key Insight: Why k=2 Rather Than k=1

The main point behind choosing top-2 over top-1 is that the second expert gives qualitatively different benefits from the first. The first expert handles the primary transformation; the second acts as a correction or complementary perspective. With only one expert, every routing error is irreversible: if the router sends a token to the wrong expert, the output is simply wrong and there is no recovery within that layer. With two experts, the model can blend contributions from a primary specialist and a secondary specialist. This makes routing less sensitive to imprecision during early training when the router has not yet learned reliable patterns. Over the course of training, this redundancy accelerates convergence because every token gives learning signal to two experts rather than one, and every expert receives more diverse training examples because it shares responsibility with other experts rather than exclusively owning a domain.

Higher $k$ values give several benefits, as we discussed in the Top-K Routing chapter:

Gradient diversity: More experts receive gradients per token, which accelerates learning. When only one expert processes a token, only that expert's parameters are updated. With two experts, both receive learning signal from the same token, speeding up convergence.
Redundancy: If one expert produces poor output, the other compensates. This reduces sensitivity to individual expert failures or poorly matched routing decisions.
Smoother specialization: Experts can develop overlapping competencies rather than hard boundaries. This prevents the model from fragmenting knowledge into non-communicating silos.

The cost is straightforward: top-2 routing requires twice the FLOPs of top-1. For Mixtral, this means each token processes through two 14,336-dimensional FFNs rather than one. This is not a trivial increase, and the Mistral team's choice to accept this cost indicates they found the quality benefits substantial.

Let's quantify this compute cost. Each SwiGLU expert performs three matrix multiplications: gate projection, up projection, and down projection. Each projection maps between the model dimension and the FFN intermediate dimension. Since top-2 routing activates two experts, the per-expert cost is doubled.

To calculate FLOPs, we count the multiply-accumulate operations for each matrix multiplication. For a matrix multiplication mapping from dimension $d_{\text{model}}$ to $d_{\text{ffn}}$ , the cost is approximately $d_{\text{model}} \times d_{\text{ffn}}$ FLOPs (simplified from $2 \times d_{\text{model}} \times d_{\text{ffn}}$ by omitting constant factors for clarity). Activating two experts requires computing three projections (gate, up, down) for each of the two selected experts. The total computational cost is:

2 \times (3 \times d_{\text{model}} \times d_{\text{ffn}}) \text{ FLOPs per token}

where:

The factor of 2 accounts for the two active experts selected by top-2 routing
The factor of 3 is the three projection matrices in each SwiGLU expert (gate, up, and down projections)
$d_{\text{model}} = 4096$ : the model's hidden dimension
$d_{\text{ffn}} = 14336$ : the intermediate dimension of each expert's feed-forward network
The product $d_{\text{model}} \times d_{\text{ffn}}$ estimates FLOPs for one matrix multiplication, simplified from $2 \times d_{\text{model}} \times d_{\text{ffn}}$ by omitting the constant factor

To build intuition for what this compute cost means in practice, it helps to compare it to an equivalent dense model. A dense FFN that performs the same number of FLOPs would have intermediate dimension:

d_{\text{ffn,dense}} = 2 \times 14336 = 28672

where:

$d_{\text{ffn,dense}}$ : the intermediate dimension of an equivalent dense FFN, which would perform the same number of FLOPs as Mixtral's top-2 MoE routing
The factor of 2 comes from activating two experts, each with dimension 14,336
A dense FFN with this intermediate dimension would perform the same three projections (gate, up, down) but process all tokens through a single large FFN instead of routing to different experts

This comparison reveals a key insight: top-2 MoE has the same computational cost as a moderately-sized dense FFN, yet maintains access to 8 experts worth of capacity (8 × 14,336 = 114,688 total expert parameters). The sparsity allows the model to select from a much larger parameter space without proportionally increasing compute. Put another way, Mixtral pays the compute cost of a 28K-width FFN but gets to choose which 28K-worth of parameters to use from a pool of 114K-worth. This is far cheaper than the alternative of using all 8 experts. If all experts were active, the effective FFN dimension would be:

d_{\text{ffn,all}} = 8 \times 14336 = 114688

where:

$d_{\text{ffn,all}}$ : the intermediate dimension if all 8 experts processed every token
The factor of 8 accounts for all experts being active simultaneously
This dimension would result in approximately 4 times the computational cost of top-2 routing (since $114688 / 28672 \approx 4$ )

The sparsity from top-2 routing therefore gives a 4x compute savings compared to activating all experts, while still maintaining access to the full 8 experts' worth of learned knowledge. This is the basic efficiency proposition of MoE. Store more parameters than you compute with, selecting the relevant subset dynamically based on the input.

Out[8]:

Visualization

Bar chart comparing effective FFN dimensions for top-1 routing (14,336), top-2 routing (28,672), and all-8-experts (114,688). Bars labeled with dimension values and relative cost multipliers 0.5x, 1.0x, and 4.0x respectively. — Effective FFN dimensions across routing strategies. Top-2 routing (1.0x baseline cost at 28,672 dimensions) balances quality with efficiency, compared to top-1 routing (0.5x cost) and full 8-expert activation (4.0x cost).

Dual-axis bar chart with k values 1 through 8 on the x-axis. Blue bars show effective FFN dimension scaling linearly; pink bars show expert capacity percentage. A green shaded region and annotation mark Mixtral's k=2 choice. — Compute versus capacity tradeoff across top-k routing choices. Top-2 routing (k=2) reaches 25% capacity utilization at 28,672 effective dimensions, balancing compute cost against available expert capacity.

Empirically, Mistral AI found that top-2 routing materially outperformed top-1 at matched compute budgets, suggesting the quality benefits outweigh the efficiency loss for their target applications.