Mixture of Experts: Sparse Models Explained

A dense transformer normally applies the same feed-forward network to every token representation. A sparse mixture-of-experts (MoE) layer offers several feed-forward networks and chooses a subset for each token. This adds stored parameters without evaluating every expert on every token. The extra routing, memory traffic, and communication still have costs; sparse activation is not a guarantee of fixed latency as the model grows.

Experts and routing

An expert is a trainable subnetwork, usually a feed-forward network with its own weights. It is not assigned a human-readable specialty such as mathematics or translation by definition. A learned router maps the token’s current hidden representation to scores over experts. Different tokens in a sentence can select different experts, and the same word can select differently in another context. Here sparsity refers to which experts execute; their weight matrices can remain dense.

Let \(x\in\mathbb R^d\), let there be \(E\) experts, and let \(p(x)=\operatorname{softmax}(W_r x)\). The router matrix has shape \((E,d)\). If \(S(x)\) contains the indices of the \(k\) largest scores, the layer computes:

\[y=\sum_{i\in S(x)}w_i(x)E_i(x).\]

Each expert output has the same width as \(x\), so the weighted outputs can be added. The weights require a convention. For top-2 below, \(w_i=p_i/\sum_{j\in S}p_j\); for top-1 we retain the selected full-softmax probability, \(w_i=p_i\), to preserve a task-gradient path to the router. These are different gating choices. For probabilities [0.6, 0.3, 0.1], normalized top-2 uses weights [2/3, 1/3]; values 3 and 9 in one output coordinate would combine to 5. A top-1 gate retaining 0.6 would instead contribute 0.6×3=1.8.

With identical expert sizes, total expert parameters scale with E and expert arithmetic per token scales with k. For E=64 and k=2, the expert bank has 64 times one expert’s parameters and evaluates two experts per token. That does not give it the same learned function or quality as a 64-times-wider dense network. Router work also grows with E, and shared attention and other layers remain. The attention and residual block structure is introduced in Transformers from Scratch.

Dispatch, compute, and combine

The example accepts an unpadded tensor of shape (batch, sequence, width), flattens its first two axes, groups assignments by expert, and adds the weighted expert results back to their token positions. It has no capacity limit: every selected assignment is computed. This is a CPU float32 teaching implementation, not a distributed or optimized mixed-precision kernel. Run the blocks in order.

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.net = nn.Sequential(nn.Linear(d_model, d_ff), nn.GELU(),
                                 nn.Linear(d_ff, d_model))
    def forward(self, x):
        return self.net(x)

class MoELayer(nn.Module):
    def __init__(self, d_model, d_ff, n_experts=4, k=2):
        super().__init__()
        if not 1 <= k <= n_experts or d_model <= 0 or d_ff <= 0:
            raise ValueError("Use positive widths and 1 <= k <= n_experts")
        self.k, self.n, self.d = k, n_experts, d_model
        self.router = nn.Linear(d_model, n_experts, bias=False)
        self.experts = nn.ModuleList([Expert(d_model, d_ff) for _ in range(n_experts)])

    def forward(self, x):
        if x.ndim != 3 or x.shape[-1] != self.d or x.numel() == 0:
            raise ValueError("Expected nonempty (batch, sequence, width) inputs")
        flat = x.reshape(-1, self.d)
        logits = self.router(flat)
        probs = F.softmax(logits, dim=-1)
        selected_logits, indices = logits.topk(self.k, dim=-1)
        if self.k == 1:
            weights = probs.gather(1, indices)
        else:
            weights = selected_logits.softmax(dim=-1)
        out = torch.zeros_like(flat)
        for e, expert in enumerate(self.experts):
            token, slot = (indices == e).nonzero(as_tuple=True)
            if token.numel() == 0:
                continue
            contribution = expert(flat[token]) * weights[token, slot, None]
            out = out.index_add(0, token, contribution)
        return out.reshape_as(x), probs, indices

torch.manual_seed(0)
moe = MoELayer(8, 16, n_experts=4, k=2)
x = torch.randn(2, 3, 8)
y, probs, indices = moe(x)
per_expert = sum(p.numel() for p in moe.experts[0].parameters())
print("shape", list(y.shape), "assignments", indices.numel())
print("expert bank", 4 * per_expert, "selected per token", 2 * per_expert)
print("router parameters", sum(p.numel() for p in moe.router.parameters()))
# shape [2, 3, 8] assignments 12
# expert bank 1120 selected per token 560
# router parameters 32

Six tokens produce twelve assignments because each selects two experts. If both experts contribute to one token, index_add sums their contributions at that token’s row. Grouping can create larger matrix operations, but the Python loop, indexing, and small expert batches add overhead. Production implementations optimize dispatch and grouped matrix multiplication. Padding tokens should be excluded from dispatch and balancing statistics; this example assumes there are none.

Top-k indices are discrete. Autograd differentiates selected weights and expert computations while treating the chosen indices as fixed; it does not differentiate the change in expert identity at a selection boundary. With normalized top-2, relative selected scores receive task gradients. With top-1, normalizing the sole selected weight to one would remove this path. Retaining its original probability, as in Switch-style gating, allows a gradient through that probability even with one selected expert. The final exercise checks this distinction directly.

Load balancing

A router can send too much traffic to a few experts. This may overload those devices and leave other experts with little task-gradient data. Feedback between routing and expert training can reinforce imbalance, but collapse is not inevitable. Monitor assignment counts and task quality; uneven counts alone do not show that experts have learned useful specialties.

For \(M\) tokens and fixed top-k routing, define \(f_i\) as expert i’s assignment count divided by \(Mk\), so \(\sum_i f_i=1\). Let \(P_i\) be its full-softmax probability averaged over the M tokens. A commonly used auxiliary objective, with this top-k normalization, is:

\[L_{\mathrm{aux}}=\alpha E\sum_i f_iP_i.\]

Here alpha is a nonnegative coefficient weighting the auxiliary term against the task loss. The counts are treated as fixed in backpropagation; gradients flow through P. This encourages lower probability mass on heavily assigned experts locally. Uniform f and P give alpha, independent of E. That value is a reference, not a general lower bound: f and P are different statistics, and their dot product can fall below 1/E. Different implementations normalize top-k counts differently, so their auxiliary-loss coefficients are not automatically comparable.

def load_balance_loss(probs, indices, alpha=0.01):
    if (probs.ndim != 2 or indices.ndim != 2 or probs.size(0) == 0
            or probs.size(0) != indices.size(0) or indices.size(1) == 0):
        raise ValueError("Expected matching nonempty token dimensions")
    n = probs.size(1)
    counts = torch.bincount(indices.reshape(-1), minlength=n).to(probs.dtype)
    fraction = counts / indices.numel()
    mean_probability = probs.mean(dim=0)
    return alpha * n * (fraction * mean_probability).sum()

base = torch.tensor([0.4, 0.3, 0.2, 0.1])
balanced = torch.stack([base.roll(i) for i in range(4)])
concentrated = torch.tensor([[0.7, 0.2, 0.06, 0.04]]).repeat(4, 1)
for name, p in [("balanced", balanced), ("concentrated", concentrated)]:
    chosen = p.topk(2, dim=-1).indices
    print(name, round(load_balance_loss(p, chosen).item(), 6))

moe.zero_grad(set_to_none=True)
y, probs, indices = moe(x)
task_loss = y.square().mean()
loss = task_loss + load_balance_loss(probs, indices)
loss.backward()
print("finite router gradient", bool(torch.isfinite(moe.router.weight.grad).all()))
# balanced 0.01
# concentrated 0.018
# finite router gradient True

The balanced example rotates a distribution over four tokens, so both assignment fractions and mean probabilities are uniform. The concentrated example sends every token to two distinct experts; its assignment fractions are [0.5, 0.5, 0, 0] and its auxiliary loss is 0.018, compared with 0.01. Sending the same token twice to one expert is not valid top-2 routing. The final lines connect the auxiliary term to a toy task loss and backpropagate; they do not train a language model or establish balanced routing after training.

Capacity and router stability

A capacity-limited implementation might allocate \(C=\lceil cMk/E\rceil\) assignment slots per expert, where c is the capacity factor. With M=8 tokens, k=2, E=4, and c=1.25, C=5. There are twenty slots for sixteen assignments. This is 25% more than average demand, but unused space would be 4/20=20% even with no overflow, before rounding effects in other examples.

Total spare space does not prevent one expert from overflowing. For counts [7, 5, 3, 1] and capacity five, two assignments exceed capacity while six slots elsewhere are unused. Policies include dropping excess expert contributions, rerouting, or using a dropless implementation with variable-size computation. Dropping a contribution is not deleting the sequence token. In a residual block the residual path still exists, but attention, normalization, other selected experts, and later layers can still change the representation. The minimal code above drops nothing.

Router z-loss is a separate stability term, often \(\beta M^{-1}\sum_t[\log\sum_i\exp(z_{ti})]^2\). Here z is the matrix of router scores before softmax, t indexes tokens, and beta weights this stability term. It penalizes the squared log-partition value, not every large logit magnitude independently, and does not directly enforce balanced assignments. Compute it with a stable logsumexp operation. Higher precision for router calculations can also help, but neither measure guarantees stable training under all configurations.

Design choices

Choose the number and width of experts, the number selected, and the locations of MoE layers together. More experts increase stored weights and router work. Larger k increases expert computation and dispatch volume, while also changing the representation and training signal. Shared experts, evaluated for every token, add a dense contribution and must be counted in both parameters and active work. These choices vary across architectures; top-1, top-2, and other k values are not interchangeable defaults.

Switch Transformer demonstrated an effective top-1 design; it did not establish that top-1 is optimal for every model. Likewise, a balancing-loss coefficient or capacity factor from one experiment is a starting point to validate, not a general recommendation. Compare task quality, overflow, expert utilization, memory, and measured throughput at the relevant batch sizes. Expert-choice routing is another family in which experts select tokens; equal per-expert quotas do not guarantee that every token receives the same number of experts.

Expert parallelism

Experts can be placed on different devices. Dispatch sends each token representation to the devices holding its selected experts; combine returns their outputs for weighting and summation. Distributed implementations often use all-to-all communication for these exchanges, with additional exchanges in backpropagation. A single-device MoE needs no network all-to-all, and local assignments need no remote transfer.

The cost depends on token volume, k, hidden width, placement, network topology, and imbalance. A slow or overloaded expert can delay the group. Communication can become a bottleneck, but it does not increase according to a universal device-count rule and can sometimes overlap other work. Expert parallelism may be combined with tensor, data, or pipeline parallelism. Measure compute and dispatch time separately when tuning these layouts. Pretraining costs are discussed in LLM Pretraining: Objectives, Data, and Scaling Laws.

What sparse activation saves

At fixed expert size and k, increasing E leaves the selected-expert arithmetic approximately fixed while increasing the parameter bank. Whether that improves quality per training FLOP is empirical and depends on data, optimization, and the dense comparison. Equal active parameter counts do not imply equal wall-clock latency or equal quality.

Low-latency serving commonly keeps the expert bank resident across devices. Offloading or caching experts can reduce accelerator residency but adds transfer and scheduling costs; all experts need not reside on one GPU. At the same time, a batch may collectively visit many or all experts even when each token visits only two. Per-token active counts therefore do not directly describe weight traffic for the whole batch. Working buffers, KV caches, and training optimizer states are separate memory costs.

During fine-tuning, inspect changes in routing and held-out quality across tasks. Rarely selected experts receive little task-gradient information; that does not by itself mean their weights overfit or drift. Optimizer state, weight decay, and other updates can also affect parameters. Serving choices should be compared on target hardware and request lengths; LLM Inference Optimization explains the distinction between memory, arithmetic, and latency.

Exercises

1. Expert weights and active work. For 64 experts with width 4096 and hidden width 16384, compute expert-only parameters and fp16 payload for top-2. Ignore biases, the router, and shared layers.

d, hidden, experts, k = 4096, 16384, 64, 2
per_expert = 2 * d * hidden
active = k * per_expert
total = experts * per_expert
print("per expert", per_expert)
print("selected per token", active)
print("expert bank", total)
print("fp16 expert bank GiB", total * 2 / 2**30)
print("equal-active dense GiB", active * 2 / 2**30)
print("payload ratio", total / active)
# per expert 134217728
# selected per token 268435456
# expert bank 8589934592
# fp16 expert bank GiB 16.0
# equal-active dense GiB 0.5
# payload ratio 32.0

2. Is the balanced loss a minimum? For two experts and top-1, route three tokens with probabilities [0.51, 0.49] and one with [0.01, 0.99]. Compare the auxiliary loss with the uniform reference alpha.

p = torch.tensor([[0.51, 0.49]] * 3 + [[0.01, 0.99]], dtype=torch.float64)
chosen = p.topk(1, dim=-1).indices
f = torch.bincount(chosen.flatten(), minlength=2) / chosen.numel()
print("assignment fractions", f.tolist())
print("mean probabilities", [round(v, 4) for v in p.mean(0).tolist()])
print("auxiliary", round(load_balance_loss(p, chosen).item(), 6))
print("uniform reference", 0.01)
# assignment fractions [0.75, 0.25]
# mean probabilities [0.385, 0.615]
# auxiliary 0.00885
# uniform reference 0.01

3. The top-1 gate gradient. Compare a selected full-softmax probability with a selected weight renormalized to one. Use one fixed expert output of 2 and squared error to target zero.

def gate_gradient(renormalize):
    logits = torch.tensor([1., 0., -1.], dtype=torch.float64, requires_grad=True)
    p = logits.softmax(0)
    selected = p.topk(1).values
    gate = selected / selected.sum() if renormalize else selected
    loss = (2 * gate).square().sum()
    gradient, = torch.autograd.grad(loss, logits)
    return gradient

for normalize in (False, True):
    print("renormalize", normalize,
          [round(v, 6) for v in gate_gradient(normalize).tolist()])
# renormalize False [1.185169, -0.866428, -0.318741]
# renormalize True [0.0, 0.0, 0.0]

References

• Shazeer et al. (2017). Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer.
• Fedus, Zoph, and Shazeer (2022). Switch Transformers.
• Zoph et al. (2022). ST-MoE: Designing Stable and Transferable Sparse Expert Models.
• Gale et al. (2023). MegaBlocks: Efficient Sparse Training with Mixture-of-Experts.
• Zhou et al. (2022). Mixture-of-Experts with Expert Choice Routing.