Mixture-of-Experts Layers

A router sends each token to a small subset of expert FFNs. Increasing the number of experts adds stored weights while keeping the selected expert count fixed.

The feed-forward projections hold most of the weights in a standard transformer layer, as we saw in Post 4.5. Widening the FFN increases its capacity, but it also increases the work done for every token. Mixture-of-experts layers separate the number of available FFNs from the number each token uses.

Doubling a dense FFN’s hidden width doubles its projection work for every token. Sparse routing offers another way to add parameters: retain multiple FFNs, but evaluate only a selected subset for each token.

With expert width and the number of selected experts fixed, adding experts increases stored parameters without increasing the main expert-matrix work per token. Router scoring, dispatch, and communication can still become more expensive.

Prerequisites

You should be comfortable with the FFN block as covered in Post 4.5, particularly why it accounts for most of a transformer's parameters and why widening it is the natural lever for adding capacity. Softmax (Post 1.3) shows up in the router. The residual stream picture from the self-attention post matters because dropped tokens still pass through via the residual connection.

The basic idea

Instead of one big FFN, you have N separate FFNs, called experts. They all have the same architecture. They're all initialized independently. Sitting in front of them is a lightweight router: a single linear layer followed by a softmax. For each token, the router looks at the hidden state and produces a probability distribution over all N experts. Then you pick the top-k winners and only run those.

The rest of the experts are skipped entirely. They perform no expert FFN arithmetic for this token, even though their weights are still sitting in GPU memory.

The output is a weighted sum of the winning experts' outputs, weighted by their routing probabilities:

where $W_{\text{router}}$ has shape $(\text{d\_model}, \text{num\_experts})$. Select the top-k entries from $g(x)$ and renormalize the selected probabilities to weights $\tilde g_i=g_i/\sum_{j\in\text{top-k}}g_j$. For the top-2 convention used below, compute:

Each $\text{Expert}_i$ is a standard FFN (or SwiGLU FFN, as we saw in Post 4.5). Training adjusts the router’s scores along with the expert weights.

Mixtral 8x7B

Mixtral 8x7B replaces each layer's dense FFN with eight expert FFNs. The '7B' name refers to the underlying model scale, not the size of each individual expert. The model uses top-2 routing, so each token activates exactly 2 of the 8 experts.

Total parameters land around 46.7 billion. Active parameters per token are roughly 12.9 billion. Its total and active parameter counts separate storage from per-token arithmetic; they do not guarantee the quality of a dense 47B model or the latency of a dense 13B model. The Mixtral report compares it favorably with Llama 2 70B on most of its reported benchmarks and reports faster inference in its setup; that is not a hardware-independent serving-cost ratio.

The load-balancing problem

A router can concentrate traffic on a small number of experts, a failure mode called expert collapse. Initial routing preferences and unequal training opportunities can reinforce one another. Balancing methods try to prevent that concentration without forcing all experts to learn the same function.

That's the mechanism at toy scale. Here's what it looks like with a real 100M-parameter MoE I trained from scratch for my own interpretability experiments:

Auxiliary loss

Shazeer et al. used auxiliary balancing objectives; the expression below is the later Switch-style form. Add it to the language-model loss to discourage concentrated routing:

where $f_i$ is the fraction of routing assignments sent to expert $i$ in the current batch (normalizing by tokens times $k$ for top-$k$), and $P_i$ is the average routing probability assigned to expert $i$. Uniform fractions and probabilities give a reference value of $\alpha$. The dot-product surrogate encourages balance, but this is not a general proof of its global minimum. The coefficient $\alpha$ controls how aggressively you enforce balance.

This connects back to the loss functions we discussed in Post 1.4. You're adding a regularization term that trades off against the main language modeling objective. Too little pressure and experts collapse. Too much and you homogenize the routing. The router loses the freedom to develop meaningful preferences because it's too busy satisfying the balance constraint.

In my experiments, I found this tension directly. An auxiliary weight of 0.1 let experts collapse. A weight of 1.0 changed the load distribution and routing preferences. These small runs illustrate sensitivity to the coefficient; they do not establish a generally optimal value.

Worked example

For a single token entering a 4-expert MoE with top-2 routing, suppose our token's hidden state is $x$ and the router produces raw logits $[2.1, 0.4, 1.8, -0.3]$.

After softmax, these become probabilities: $[0.4965, 0.0907, 0.3678, 0.0450]$ (rounding to four places). The top-2 experts are Expert 0 and Expert 2. Their probabilities sum to about 0.8643. Renormalizing gives weights $0.5744$ and $0.4256$.

Now we run the token through both experts and combine:

Experts 1 and 3 do not run for this token. Their weights exist in memory but contribute zero FLOPs for this token. If each expert has a $d_{\text{model}} \times d_{\text{ff}}$ up-projection and a $d_{\text{ff}} \times d_{\text{model}}$ down-projection, we just paid for 2 of those instead of all 4.

For the auxiliary loss on this micro-batch, suppose across 100 tokens the fraction of routing assignments sent to each expert is $f = [0.45, 0.15, 0.30, 0.10]$ and the average softmax probabilities are $P = [0.40, 0.12, 0.35, 0.13]$. Then $L_{\text{aux}} = \alpha \cdot 4 \cdot (0.45 \times 0.40 + 0.15 \times 0.12 + 0.30 \times 0.35 + 0.10 \times 0.13) = \alpha \cdot 4 \cdot 0.316 = 1.264\alpha$. With perfectly uniform routing, $f_i = P_i = 0.25$ for all $i$, and the penalty drops to $\alpha \cdot 4 \cdot 4 \times 0.0625 = \alpha$.

Capacity factor and token dropping

GShard and Switch Transformer use per-expert capacity limits. Overflow can drop an assignment to an expert. With top-2 routing, a token may still receive its other expert’s output; if all its assignments are dropped, the residual path carries it without an expert update.

The capacity factor $C$ determines how much slack each expert gets. With $C = 1.25$, each expert can handle at most 25% more tokens than its perfectly balanced share. Anything beyond that gets dropped.

Dropping an expert assignment does not itself provide a differentiable signal to the router. Auxiliary balancing losses provide a separate training signal to reduce overload.

A capacity limit can discard an otherwise useful expert update. Overloaded experts therefore affect both the amount of computation performed and the representations produced, particularly while routing is still changing during training.

Switch Transformer: top-1 routing

Fedus et al. (2022) pushed the simplicity even further. The Switch Transformer uses top-1 routing: each token goes to exactly one expert. No second expert or combination of two outputs; the selected expert output is still multiplied by its router probability in Switch.

With top-1, you avoid the overhead of combining two experts' outputs. The Switch paper reports speedups for particular matched-compute training comparisons; the end-to-end gain depends on dispatch, hardware, and expert utilization.

Top-2 combines two expert outputs and requires more expert work than top-1 at the same expert width. It does not guarantee protection from routing mistakes; both routing schemes need to be evaluated with their training and balancing recipes.

Token-choice vs. expert-choice

Token-choice routing selects experts separately for each token. Zhou et al. (2022) instead select tokens separately for each expert.

The router still produces a score matrix of shape (number of routed tokens, num_experts). But instead of running top-k across the expert dimension (token picks experts), you run top-c across the token dimension (expert picks tokens). Ignoring integer rounding and capacity limits, each expert selects $c$ tokens:

This equalizes expert token counts within the routing batch. It does not guarantee full GPU utilization or eliminate communication costs.

But the token's perspective changes. Some tokens get claimed by three or four experts, receiving extra computation. Other tokens may be selected by none and pass through only via the residual connection. Selection count is not a measured difficulty score.

Selection across a routing batch can also make a token depend on other tokens in that batch, so causal generation needs a routing design that avoids access to future tokens. It is not a drop-in replacement for token-wise top-k in every autoregressive setting.

Do experts actually specialize?

There's some evidence that experts develop domain-specific preferences. In my own experiments, I saw routing patterns emerge where certain experts showed affinity for math tokens vs. story tokens vs. code tokens. The recorded expert-choice assignments also show domain preferences, though the two-slot logging format does not reveal every possible expert assignment.

But the specialization is messy. Many experts learn overlapping, general-purpose features. You rarely get a clean "math expert" and a clean "language expert." The routing patterns look more like soft preferences than hard partitions.

DeepSeek-V3 uses fine-grained experts (256 routed experts, top-8 routing) plus one always-active shared expert in each MoE layer. Its first three layers use dense FFNs. The shared expert processes every token; the router selects a subset of the other experts. The architecture encourages a separation between shared and routed computation, but does not assign fixed semantic jobs to individual experts.

Implementation

Here's a minimal MoE layer with top-2 routing and auxiliary loss. Each expert is a standard two-layer FFN. The router is a single linear projection.

import torch
import torch.nn as nn
import torch.nn.functional as F
 
class Expert(nn.Module):
    def __init__(self, d_model, d_ff):
        super().__init__()
        self.w1 = nn.Linear(d_model, d_ff)
        self.w2 = nn.Linear(d_ff, d_model)
 
    def forward(self, x):
        return self.w2(F.silu(self.w1(x)))
 
class MoELayer(nn.Module):
    def __init__(self, d_model, d_ff, num_experts=8, top_k=2, aux_weight=0.01):
        super().__init__()
        if not 1 <= top_k <= num_experts:
            raise ValueError("top_k must be between 1 and num_experts")
        self.router = nn.Linear(d_model, num_experts, bias=False)
        self.experts = nn.ModuleList(
            [Expert(d_model, d_ff) for _ in range(num_experts)]
        )
        self.top_k = top_k
        self.num_experts = num_experts
        self.aux_weight = aux_weight
 
    def forward(self, x):
        B, S, D = x.shape
        x_flat = x.reshape(-1, D)          # (B*S, D)
 
        # Router: produce scores over experts
        logits = self.router(x_flat)     # (B*S, num_experts)
        probs = F.softmax(logits, dim=-1)
 
        # Select top-k experts per token
        top_probs, top_idx = torch.topk(probs, self.top_k, dim=-1)
        # Renormalize the top-k weights
        top_probs = top_probs / top_probs.sum(dim=-1, keepdim=True)
 
        # Dispatch and combine
        out = torch.zeros_like(x_flat)
        for i, expert in enumerate(self.experts):
            mask = (top_idx == i).any(dim=-1)   # which tokens go here
            if not mask.any():
                continue
            # Gather weight for this expert
            weight = (top_probs * (top_idx == i).float()).sum(dim=-1)
            expert_out = expert(x_flat[mask])
            out[mask] += weight[mask].unsqueeze(-1) * expert_out
 
        # Auxiliary loss: encourage balanced routing
        # f_i = fraction of top-k assignments sent to expert i
        # P_i = mean routing probability for expert i
        tokens_per_expert = torch.zeros(self.num_experts, device=x.device)
        for i in range(self.num_experts):
            tokens_per_expert[i] = (top_idx == i).any(dim=-1).float().sum()
        f = tokens_per_expert / (x_flat.shape[0] * self.top_k)
        P = probs.mean(dim=0)
        aux_loss = self.aux_weight * self.num_experts * (f * P).sum()
 
        return out.view(B, S, D), aux_loss

This is a pedagogical implementation. Production systems like Megablocks use block-sparse matrix operations and custom CUDA kernels to avoid the Python for-loop over experts.

Serving challenges

All experts' weights need to be in memory (or at least swappable), even though only k of them run per token. For Mixtral 8x7B, that means loading around 47 billion parameters even though you only activate 13 billion per token.

In batched inference, different tokens can collectively activate all experts. The arithmetic remains sparse: 256 tokens routed to two experts require 512 token-expert evaluations, rather than 2,048 if every token ran through all eight. What changes is weight traffic, dispatch overhead, and the amount of useful work each expert can batch.

This is where expert parallelism comes in (which we'll cover in depth in Post 6.6). You distribute experts across GPUs. Each GPU holds a subset of experts. When a token needs an expert on a different GPU, you do an all-to-all communication to send the token there and get the result back. This adds latency, but it lets you fit much larger MoE models than any single GPU could hold.

Misconceptions

"MoE means more capacity for free." At fixed expert width and top-k, adding experts keeps expert arithmetic per token roughly flat. It still adds stored weights and can increase router, dispatch, and communication costs. Batch size, placement, and load balance determine how much of the arithmetic saving becomes a wall-clock or monetary saving.

"You only run the chosen experts, so serving must be cheap." Sparse arithmetic is only part of serving cost. A batch may activate every expert collectively, while each token still uses its selected subset. Weight storage, bandwidth, dispatch, communication, and batching determine whether those arithmetic savings translate into lower latency or cost.

"Each expert learns a distinct topic, like a math expert or a code expert." The specialization is messier than that framing suggests. Some routing preferences develop, especially with expert-choice routing, but the patterns are mostly low-level token-class features inside a single layer (whitespace tokens here, math symbols there) and they shuffle between layers. Cross-layer "an expert that handles all the math" rarely emerges. The domain-affinity charts earlier in this post are about as sharp as the specialization ever got in my own training runs.

What's next

Arc 5 shifts from architecture to runtime. The first post is Training View vs. Inference View, which compares the matrix shapes and memory traffic of training, prompt processing, and token-by-token generation.

Additional reading (and watching)

• Shazeer, N., et al. (2017). Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer. ICLR 2017. The paper that brought MoE into modern deep learning as a practical sparsely-gated layer, with the auxiliary load-balancing loss that most later work builds on.
• Fedus, W., Zoph, B., & Shazeer, N. (2022). Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity. JMLR 2022. Top-1 routing plus capacity factors; the first MoE scaled to a trillion parameters with a clean engineering story.
• Lepikhin, D., et al. (2021). GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding. ICLR 2021. Introduces per-expert capacity limits and automatic sharding — the serving-side plumbing that made large MoE trainable.
• Zhou, Y., et al. (2022). Mixture-of-Experts with Expert Choice Routing. NeurIPS 2022. The expert-choice flip: each expert picks its tokens, guaranteeing perfect load balance without an auxiliary loss.
• Jiang, A. Q., et al. (2024). Mixtral of Experts. The model that made MoE visibly competitive with dense frontier models in open weights — 8 experts, top-2, ~13B active / ~47B total.
• DeepSeek-AI. (2024). DeepSeek-V3 Technical Report. Fine-grained MoE with 256 routed experts, top-8, plus a set of always-active shared experts — an example of fine-grained routing with a shared expert.