Arc 4: Transformers from First Principles

Mixture-of-Experts Layers

14 min readbuilds on The Feed-Forward Block as Key-Value Memory, Distributions, Softmax, and the Chain Rule of Words, Self-Attention: Q, K, V from First Principles and Cross-Entropy, KL Divergence, and What Loss Functions Measure

A router sends each token to a small subset of expert FFNs. Total parameters balloon; compute per token barely moves.

Why can't we just make a model bigger and call it a day?

We talked in Post 4.5 about how the feed-forward block is where most of the parameters live. The FFN is the transformer's memory bank. So if you want a more capable model, the obvious play is to widen the FFN. More neurons, more capacity, more knowledge baked into those weight matrices.

But here's the problem. If you double the FFN's hidden dimension, you double the compute for every single token. Every token, no matter how trivial, pays the full cost of that giant matrix multiply. A comma gets the same treatment as a nuanced mathematical expression. That's wasteful. And at the scale we're talking about (hundreds of billions of parameters), it becomes prohibitively expensive.

Mixture-of-Experts breaks this coupling. You get to scale the total parameters way up while keeping per-token compute roughly constant. The core idea is straightforward.

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 smaller 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 burn no compute and consume no memory bandwidth 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:

g(x) = \text{softmax}(x \cdot W_{\text{router}})

where $W_{\text{router}}$ has shape $(\text{d\_model}, \text{num\_experts})$ . Select the top-k entries from $g(x)$ , zero out the rest, and compute:

y = \sum_{i \in \text{top-k}} g_i(x) \cdot \text{Expert}_i(x)

Each $\text{Expert}_i$ is a standard FFN (or SwiGLU FFN, as we saw in Post 4.5). The router learns which experts are useful for which kinds of input. In theory, anyway. We'll get to the caveats.

Tokens from different domains enter the transformer, get routed to their top-2 experts, and recombine. Most experts stay dark for any given token.

Mixtral 8x7B

This is easier to reason about with real numbers. Mixtral 8x7B has 8 expert FFNs, each with roughly 7B parameters' worth of weights. 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. You get the knowledge capacity of a 47B-parameter model at the inference cost of something closer to a 13B dense model. Mixtral matched or beat LLaMA 2 70B on most benchmarks while being dramatically cheaper to serve.

This is the framing that matters most. Sparsity in MoE buys you compute, not memory. Every expert sits in GPU memory whether or not it fires this token. We'll come back to that distinction in the serving section.

The load-balancing problem

If this is so great, why isn't every model an MoE? Because the router is a trainable neural network, and trainable neural networks find shortcuts.

Without any intervention, the router quickly learns to send most tokens to the same one or two experts. This is called expert collapse. Early in training, random initialization causes the router to slightly prefer certain experts. Those experts see more data, so they learn faster. The router notices they've gotten better, sends them even more tokens, and the positive feedback loop crushes the remaining experts into irrelevance.

You end up with 8 experts but effectively using 2. All that extra capacity is wasted.

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:

Expert load distribution from a training run with gentle auxiliary pressure (weight 0.1). A few experts dominate while others receive almost no tokens. Step through the training to see the collapse develop.

Auxiliary loss

The original fix comes from Shazeer et al. (2017). You add a penalty term to the training loss that measures how imbalanced the load is across experts:

L_{\text{aux}} = \alpha \cdot N \cdot \sum_{i=1}^{N} f_i \cdot P_i

where $f_i$ is the fraction of tokens routed to expert $i$ in the current batch, and $P_i$ is the average routing probability assigned to expert $i$ . When load is perfectly uniform, this penalty is minimized. 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. Cranking it to 1.0 produced perfectly balanced load but flat, undifferentiated routing patterns. The sweet spot was narrow and fragile.

Worked example

Let me walk through the math for a single token hitting a 4-expert MoE with top-2 routing. Say 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.50, 0.09, 0.37, 0.05]$ (rounding to two places). The top-2 experts are Expert 0 (score 0.50) and Expert 2 (score 0.37). We renormalize so they sum to 1: Expert 0 gets weight $0.50 / 0.86 \approx 0.57$ , Expert 2 gets $0.37 / 0.86 \approx 0.43$ .

Now we run the token through both experts and combine:

y = 0.57 \cdot \text{Expert}_0(x) + 0.43 \cdot \text{Expert}_2(x)

Experts 1 and 3 never fire. 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 routed 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.31 = 1.24\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 the Switch Transformer introduced a more physical solution: hard capacity limits. Each expert has a maximum buffer size per batch. If an expert's queue overflows, excess tokens are dropped. They pass through via the residual connection without any expert processing at all.

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.

It's a brutal mechanism, but it works. Token dropping forces the router to learn about its own capacity constraints. When a token gets dropped, the resulting error gradient punishes the routing scores that sent it to the full queue. Over time, the router learns to spread the load more evenly to avoid drops.

But dropped tokens are lost computation. Early in training, when the router is still calibrating, a lot of tokens get dropped. And the capacity limit can actively interrupt good routing decisions. A math token that should go to the math expert might get diverted elsewhere because that expert's buffer was full for this particular micro-batch.

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 weighted combination, no second-choice fallback.

This is simpler and faster. With top-1, you avoid the overhead of combining two experts' outputs. The Switch Transformer achieved a 4x speedup over a dense model at equivalent compute budget, with competitive quality.

The trade-off is robustness. With top-2, a bad routing decision is partially mitigated by the second expert. With top-1, if the router sends a token to the wrong expert, that's all the computation it gets. You're betting everything on the router making the right call.

Token-choice vs. expert-choice

Everything I've described so far is token-choice routing. Each token picks its favorite experts. But Zhou et al. (2022) flipped the problem entirely.

In expert-choice routing, each expert picks its favorite tokens.

The router still produces a score matrix of shape (batch_size, 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). Each expert selects exactly $c$ tokens:

c = \frac{\text{batch\_size} \times k}{\text{num\_experts}}

This guarantees perfect load balancing from the expert's perspective. Every expert processes exactly $c$ tokens, every GPU stays fully utilized, no hardware wasted.

But the token's perspective changes. Some highly salient tokens get claimed by three or four experts, receiving extra computation. Meanwhile, "easy" tokens like punctuation might not be claimed by any expert and pass through only via the residual connection.

So expert-choice doesn't eliminate the dropping problem, it reframes it. Instead of accidentally overloading popular experts and dropping overflow, each expert deliberately picks the tokens it can help most with. The tokens that get skipped are the ones no expert found interesting, which feels like a more principled trade-off than capacity-overflow drops.

Expert load with expert-choice routing. Every expert processes the same number of tokens by construction.

Do experts actually specialize?

This is the question everyone asks, and the answer is "sort of."

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 patterns were clearest with expert-choice routing, where the router was free to optimize without a balancing penalty pulling it in two directions.

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.

Domain affinity for an expert-choice run. Some experts show preferences for math or code tokens, but the boundaries are fuzzy.

DeepSeek-V3 took an interesting approach: fine-grained experts (256 routed experts, top-8 routing) plus one shared expert per layer that is always active. The shared expert handles the common-case patterns that every token needs. The routed experts handle whatever specialized computation is left over. This division of labor seems to work well in practice and has become influential.

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__()
        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.view(-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 tokens 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]
        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. But the logic is the same.

Serving challenges

MoE introduces a unique headache for inference. 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.

And in batched inference, different tokens in the same batch will route to different experts. If your batch has 256 tokens and each hits 2 of 8 experts, you probably need all 8 experts active. The sparsity savings that work great for single-token inference largely disappear in high-throughput serving.

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." Per-token FLOPs stay roughly flat, but everything else gets harder. Training cost in FLOPs is similar to a dense model with the same active-parameter count, but training a frontier MoE actually takes more wall-clock and money than the FLOP count suggests, because you're pushing huge weight matrices around an interconnect, juggling expert parallelism, and tuning load-balance pressure to avoid collapse. At serving time the bottleneck shifts to memory bandwidth and capacity, not compute. The cost is real, just rearranged.

"You only run the chosen experts, so MoE is much cheaper to serve." True per token, false at the system level. All N experts have to be in GPU memory (or hot in CPU memory and swappable). In a batched inference workload of any reasonable size, different tokens route to different experts, so almost every expert ends up active per batch. The single-token sparsity story doesn't survive batching.

"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

That closes out Arc 4. You now have a full mental model of the transformer: the residual stream as a backbone; multi-head attention as the communication primitive; causal masking, positional encoding, normalization, and FFNs as the supporting cast; encoder-decoder vs decoder-only as the architectural fork; and MoE as the dense-to-sparse generalization of the FFN block.

Arc 5 shifts from architecture to runtime, and this is where the picture you've been carrying starts to look very different from what the papers describe. The first post is Training View vs. Inference View, which sets up the prefill-vs-decode split that shapes everything downstream: KV caches, sampling strategies, speculative decoding, continuous batching, and the inference-engine concerns that turn an architecture into a served product.

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 — currently the dominant large-MoE template.