A Mixture-of-Experts layer replaces the single feed-forward network in a Transformer block with many parallel "expert" networks and a small learned router that sends each token to only a few of them. The model gains parameters (capacity) without a proportional increase in compute per token. Sparse MoE is used in several large decoder-only LLMs, including openly documented ones such as Mixtral, DeepSeek-MoE and the Switch Transformer.
This page covers routing, the parameter arithmetic, load balancing, a working PyTorch layer, and the systems trade-offs.
For token representation h, the router computes p = softmax(h Wr) over E experts, keeps the top-k, renormalizes their weights gi, and outputs y = Σi∈top-k gi Ei(h). Each expert is an ordinary FFN. Attention layers are usually shared (dense); only the FFNs are expert-ized, often in every layer or every other layer.
Switch Transformer used k = 1; Mixtral and many others use k = 2; fine-grained designs (DeepSeek-MoE) use many smaller experts plus one or more always-on "shared" experts.
Quoting an MoE model by total parameters overstates its per-token compute; quoting it by active parameters understates its memory footprint. All experts must sit in (GPU) memory because any token may route to any expert, so serving an MoE needs memory like a dense model of its total size while generating at roughly the FLOPs of a dense model of its active size.
Left alone, routers collapse: a few experts get most tokens, train faster, and attract even more tokens. Fixes:
α · E · Σi fi Pi, where fi is the fraction of tokens dispatched to expert i and Pi is its mean router probability; it is minimized by uniform routing.C · tokens / E tokens per batch; overflow tokens skip the layer via the residual path (dropped tokens).import torch, torch.nn as nn, torch.nn.functional as F
class SwiGLU(nn.Module):
def __init__(self, d, h):
super().__init__()
self.w1, self.w3 = nn.Linear(d, h, bias=False), nn.Linear(d, h, bias=False)
self.w2 = nn.Linear(h, d, bias=False)
def forward(self, x):
return self.w2(F.silu(self.w1(x)) * self.w3(x))
class MoE(nn.Module):
def __init__(self, d=512, hidden=1024, n_experts=8, k=2):
super().__init__()
self.k, self.n = k, n_experts
self.router = nn.Linear(d, n_experts, bias=False)
self.experts = nn.ModuleList(SwiGLU(d, hidden) for _ in range(n_experts))
def forward(self, x): # x: (tokens, d)
probs = F.softmax(self.router(x), dim=-1) # (N, E)
gate, idx = probs.topk(self.k, dim=-1) # (N, k)
gate = gate / gate.sum(-1, keepdim=True)
y = torch.zeros_like(x)
for e in range(self.n): # dispatch per expert
tok, slot = (idx == e).nonzero(as_tuple=True)
if tok.numel():
y[tok] += gate[tok, slot, None] * self.experts[e](x[tok])
# Switch-style auxiliary loss: n * sum_i f_i * P_i
f = F.one_hot(idx[:, 0], self.n).float().mean(0) # fraction routed (top-1)
P = probs.mean(0) # mean router prob
aux = self.n * (f * P).sum()
return y, aux
x = torch.randn(64, 512)
y, aux = MoE()(x)
print(y.shape, float(aux)) # aux ~= 1.0 when perfectly balancedThis loop-over-experts version is clear but slow. Production implementations sort tokens by expert and use grouped GEMM kernels (e.g. MegaBlocks-style block-sparse matmuls) and, across GPUs, expert parallelism: experts are sharded over devices and tokens are exchanged with all-to-all communication.
| Benefit | Cost |
|---|---|
| More capacity at roughly constant FLOPs per token | Memory for all experts; harder to fit on few GPUs |
| Faster pretraining to a target loss for a compute budget | All-to-all communication in distributed training and serving |
| Experts can specialize (often on token types/syntax more than topics) | Training instability; balancing hyperparameters |
| Good throughput at large batch | Small-batch decode is still memory-bound across many weights |
Consider a 32-layer decoder with model width d = 4096 and an expert FFN of hidden size 14,336 using a gated (SwiGLU) design, which has three weight matrices per expert. These are the published dimensions of Mixtral 8x7B, so the arithmetic can be checked against its reported totals.
Total ≈ 45.1 + 1.3 + 0.26 ≈ 46.7B parameters. With top-2 routing each token touches two experts per layer: 2 × 176M × 32 ≈ 11.3B, plus the shared attention and embeddings, giving ≈ 12.9B active parameters. The name "8x7B" is therefore misleading in both directions: the model is not 56B (attention and embeddings are shared), and per-token compute resembles a 13B dense model, while memory resembles a 47B one. In 16-bit precision the weights alone need about 93 GB.
Suppose a router produces logits (2.0, 1.0, 0.5, −1.0) over four experts for one token. The softmax probabilities are approximately (0.61, 0.22, 0.14, 0.03). With top-2 routing the token goes to experts 1 and 2; renormalising their probabilities gives gates 0.61 / 0.83 ≈ 0.73 and 0.22 / 0.83 ≈ 0.27, and the layer output is 0.73·E1(h) + 0.27·E2(h). (Some implementations skip the renormalisation and use the raw softmax weights; the choice is part of the model definition and must match at inference.)
Now the capacity limit. With 4,096 tokens in a batch, E = 4 experts, top-2 routing and capacity factor C = 1.25, each expert accepts at most 1.25 × (4,096 × 2) / 4 = 2,560 token slots. If routing is skewed and expert 1 is chosen by 3,000 tokens, 440 of those assignments overflow; those tokens get only their other expert's contribution, or pass through on the residual path if both choices overflow. Monitoring the fraction of dropped assignments per layer is one of the most useful health checks during MoE training.
Because experts are independent, they can be placed on different devices. Each MoE layer then performs two all-to-all exchanges: tokens are sent to the devices that own their chosen experts, processed, and sent back to be combined with their gate weights. Communication rather than arithmetic often limits throughput, which is why designs that restrict how many devices a token can be routed to, and fast interconnects within a node, matter so much at scale.
Serving considerations differ from dense models:
| Situation | Lean toward | Reason |
|---|---|---|
| Fixed training compute budget, ample memory | MoE | Reaches a lower loss for the same training FLOPs in published comparisons (Switch Transformer, GLaM) |
| Single GPU or edge deployment | Dense | Memory, not compute, is the binding constraint |
| High-throughput serving across many GPUs | MoE | Large batches amortise expert weights; fewer FLOPs per token |
| Small team, small fine-tuning dataset | Dense | Simpler training, fewer instability modes |
The idea is older than transformers: Jacobs et al. (1991) introduced adaptive mixtures of local experts, and Shazeer et al. (2017) made sparsely gated MoE layers practical in large LSTM language models before the approach moved into decoder-only transformers.