244  Mixture of Experts

244.1 1. Introduction

A dense transformer applies every parameter to every token. Doubling the width of the feed-forward blocks doubles both the parameter count and the cost of every forward pass, so capacity and compute are chained together. Mixture of Experts (MoE) breaks that chain. It replaces a single dense feed-forward network with a bank of parallel networks, the experts, and a small learned router that sends each token to only a few of them. The model can hold an enormous number of parameters while spending compute proportional to the few experts each token actually visits.

This decoupling of total parameters from active parameters is the central reason MoE sits underneath most frontier-scale language models shipped through 2025. The Switch Transformer scaled to over a trillion parameters while keeping per-token compute fixed, and open releases such as Mixtral 8x7B and DeepSeek-V3 made sparse MoE the default architecture for the largest open-weight models. The technique is conceptually simple, a gate plus a set of experts, but it introduces a genuinely new engineering problem: the router must spread load evenly across experts, or the model wastes most of its capacity on a handful of overused experts while the rest starve.

This chapter develops the mathematics of sparse top-k routing and the load-balancing loss that makes training stable, then builds a complete sparse MoE layer from scratch in PyTorch on the CPU, small enough to run in seconds yet faithful to the production design. It closes with the GPU-scale invocations a reader would actually use (Switch Transformer, Mixtral, expert-parallel training), shown but not executed because they require accelerators, and with the pitfalls that decide whether an MoE model trains well or collapses.

244.2 2. The Core Idea: Conditional Computation

In a dense transformer block the feed-forward sublayer computes, for a token representation \(x \in \mathbb{R}^{d}\),

\[ \mathrm{FFN}(x) = W_2\,\sigma(W_1 x), \]

and this is applied identically to every token. The MoE block replaces \(\mathrm{FFN}\) with \(E\) experts \(\{f_1, \dots, f_E\}\), each itself a feed-forward network with its own parameters, plus a router (or gate) \(g\) that produces a distribution over experts. The output for a token is a weighted combination of the experts the router selects,

\[ y = \sum_{i \in \mathcal{T}(x)} g_i(x)\, f_i(x), \]

where \(\mathcal{T}(x)\) is the set of experts chosen for \(x\). If \(\mathcal{T}(x)\) is all \(E\) experts the layer is dense MoE and costs \(E\) times a single FFN. The frontier design is sparse: \(\mathcal{T}(x)\) is the top \(k\) experts by gate score, with \(k\) much smaller than \(E\) (commonly \(k = 1\) or \(k = 2\)). A token then pays for only \(k\) experts no matter how large \(E\) grows, which is exactly the property that lets total parameters scale independently of per-token compute.

The distinction between total and active parameters is the vocabulary of MoE. Mixtral 8x7B has eight experts per layer and routes to two, so it carries roughly 47 billion total parameters but activates about 13 billion per token. Capacity that large at the cost of a 13 billion parameter forward pass is the whole point.

244.3 3. The Gating Network

The router is a single linear map followed by a softmax over experts. Given a token \(x\), the router logits are

\[ h(x) = W_g\, x \in \mathbb{R}^{E}, \qquad W_g \in \mathbb{R}^{E \times d}, \]

and the gate weights are the softmax of those logits,

\[ p_i(x) = \frac{\exp\!\big(h_i(x)\big)}{\sum_{j=1}^{E} \exp\!\big(h_j(x)\big)}. \]

For sparse routing we keep only the top \(k\) logits. Let \(\mathcal{T}(x)\) index the \(k\) largest entries of \(h(x)\). There are two standard conventions for the combine weights \(g_i\). The first softmaxes over the full set of experts and then keeps the selected weights; the second, used in Switch and many later models, renormalizes the softmax over the selected experts alone so the surviving weights sum to one,

\[ g_i(x) = \frac{\exp\!\big(h_i(x)\big)}{\sum_{j \in \mathcal{T}(x)} \exp\!\big(h_j(x)\big)}, \qquad i \in \mathcal{T}(x). \]

Renormalizing keeps the output on the same scale regardless of how confident the router is, which stabilizes training. The reference implementation below uses this renormalized top-k form.

The top-k selection is not differentiable in the indices, the choice of which experts to use is a hard, discrete decision. Gradients still flow to the router through the weights \(g_i(x)\) of the experts that were selected, because those weights are smooth functions of the logits. The router therefore learns: if an expert produces output that reduces the loss, the gradient increases that expert’s weight, which raises its logit, which makes the router more likely to select it again. This indirect learning signal is enough in practice, but it is also the source of the load-balancing problem in the next section.

244.4 4. The Load-Balancing Problem

Left to itself, the router tends to collapse. Early in training a few experts are slightly better by chance, the router sends them more tokens, those experts receive more gradient and improve faster, and the router sends them still more tokens. This positive feedback loop concentrates traffic on a small subset of experts. The unused experts receive almost no gradient and never improve, so the effective capacity of the model shrinks toward that of a much smaller dense network. This failure is called routing collapse, and preventing it is the defining engineering challenge of MoE training.

The standard remedy, introduced with the Switch Transformer, is an auxiliary load-balancing loss added to the task loss. Consider a batch of \(T\) tokens routed across \(E\) experts. Define two quantities per expert. Let \(f_i\) be the fraction of tokens dispatched to expert \(i\),

\[ f_i = \frac{1}{T}\sum_{t=1}^{T} \mathbb{1}\!\left[\, i \in \mathcal{T}(x_t)\,\right], \]

and let \(P_i\) be the average router probability mass assigned to expert \(i\) over the batch,

\[ P_i = \frac{1}{T}\sum_{t=1}^{T} p_i(x_t). \]

The auxiliary loss is the scaled dot product of these two vectors,

\[ \mathcal{L}_{\text{aux}} = \alpha \cdot E \cdot \sum_{i=1}^{E} f_i\, P_i, \]

with \(\alpha\) a small coefficient (Switch used \(\alpha = 10^{-2}\)). The factor \(E\) makes the loss scale-independent of the number of experts: when load is perfectly uniform, \(f_i = P_i = 1/E\) for every expert, the sum equals \(E \cdot (1/E)^2 = 1/E\), and \(E \cdot \mathcal{L}_{\text{aux}}/\alpha = 1\). Any imbalance pushes the value above one.

The design is subtle and worth dwelling on. The term \(f_i\) counts hard dispatch decisions and is not differentiable, but it acts only as a per-expert weight on the differentiable probability \(P_i\). The gradient therefore flows through \(P_i\), and it is largest for experts that are both overloaded (large \(f_i\)) and receiving high router confidence (large \(P_i\)). Minimizing \(\sum_i f_i P_i\) pushes the router to lower the probability it assigns to whichever experts are currently overworked, which spreads tokens toward the idle experts. The product form is what couples the hard load counts to the soft, trainable router so a single differentiable loss can flatten the distribution.

A complementary signal sometimes added is the router z-loss, \(\mathcal{L}_z = \frac{\beta}{T}\sum_t \big(\log \sum_j e^{h_j(x_t)}\big)^2\), which penalizes large router logits and improves numerical stability of the softmax. The implementation below focuses on the load-balancing loss, which is the one that prevents collapse.

244.5 5. Expert Capacity and Token Dropping

A second practical constraint appears when MoE runs on real hardware. Experts are implemented as fixed-size tensors, so each expert can process only a bounded number of tokens per batch. This bound is the expert capacity,

\[ C = \Big\lceil \text{capacity\_factor} \cdot \frac{T \cdot k}{E} \Big\rceil, \]

where \(T \cdot k / E\) is the number of tokens an expert would receive under perfectly uniform routing and the capacity factor (commonly 1.0 to 1.5) adds headroom for imbalance. If more than \(C\) tokens route to one expert, the overflow tokens are dropped: they skip the expert entirely and pass through only the residual connection. Capacity factor thus trades memory and compute (higher factor, fewer drops, more padding) against efficiency (lower factor, more drops). Token dropping is a feature of the batched GPU implementation, not of the mathematics, and our CPU reference computes the true routing without dropping so the mechanism is visible in isolation; the production code section returns to capacity as a deployment concern.

244.6 6. Compressing a Hybrid MoE for Deployment

A trained MoE is not the model you serve. Conditional computation cuts the FLOPs per token but not the memory footprint: every expert must be resident, and at frontier scale the parameter count (together with the KV cache) dominates the interactive-serving budget rather than arithmetic does. A production line of work therefore compresses a large MoE after training to meet tighter latency and concurrency targets. Bercovich et al. (2026)’s Nemotron-Labs-3-Puzzle-75B-A9B is a representative case: a hybrid MoE with 75B total parameters but only about 9B active per token (the “A9B”), derived from a larger parent by a multi-stage pipeline that pairs an iterative Puzzle compression search with knowledge distillation (the subject of the distillation chapter), reinforcement learning, quantization, and a multi-token-prediction head. Crucially, the compression jointly optimizes three coupled budgets rather than one: heterogeneous MoE pruning (dropping or shrinking experts unevenly across layers, since the load-balancing view of Section 4 implies different layers carry different amounts of routed information), the active-parameter budget (how many experts fire per token), and, because the architecture interleaves attention with Mamba state-space layers, Mamba-layer pruning.

The payoff is a serving result, not an accuracy one: on a single 8xB200 node the compressed model reaches roughly 2x the server throughput of its uncompressed parent at matched per-user throughput, and on a single H100 it lifts 1M-token-context concurrency from one request to eight, while retaining most of the parent’s quality across reasoning, coding, multilingual, long-context, and agentic benchmarks. The lesson for MoE practitioners is a division of labor: the routing and load-balancing machinery of this chapter sets a model’s quality, but a separate post-training compression stage sets whether that quality is economical to serve.

244.7 7. Production Code

We now build a sparse top-k MoE layer end to end in PyTorch, running on the CPU. The implementation is deliberately small and deterministic so every number below is reproducible, yet it contains the real pieces: a linear router, renormalized top-k gating, per-expert dispatch, the load-balancing auxiliary loss, and a measurement of how evenly tokens spread across experts. Mature open-source building blocks (PyTorch, NumPy) carry all the weight here; nothing in this section needs a GPU.

244.7.1 7.1 A Sparse Top-k MoE Layer

The expert is an ordinary two-layer feed-forward network. The router is a single linear layer. The forward pass computes router logits, selects the top \(k\) experts per token, renormalizes their gate weights, dispatches each token’s representation to its chosen experts, and accumulates the gated outputs. We also compute the load-balancing loss from the same routing decisions so it can be added to a task loss during training.

import torch
import torch.nn as nn
import torch.nn.functional as F

torch.manual_seed(0)


class Expert(nn.Module):
    """A single feed-forward expert."""

    def __init__(self, d_model: int, d_hidden: int):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(d_model, d_hidden),
            nn.GELU(),
            nn.Linear(d_hidden, d_model),
        )

    def forward(self, x):
        return self.net(x)


class SparseMoE(nn.Module):
    """Sparse top-k mixture-of-experts feed-forward layer."""

    def __init__(self, d_model: int, d_hidden: int, num_experts: int,
                 k: int, aux_alpha: float = 1e-2):
        super().__init__()
        self.num_experts = num_experts
        self.k = k
        self.aux_alpha = aux_alpha
        self.router = nn.Linear(d_model, num_experts, bias=False)
        self.experts = nn.ModuleList(
            [Expert(d_model, d_hidden) for _ in range(num_experts)]
        )

    def forward(self, x):
        # x: (tokens, d_model). Flatten any batch/seq dims before calling.
        tokens, d_model = x.shape
        logits = self.router(x)                       # (tokens, num_experts)
        probs = F.softmax(logits, dim=-1)             # full router distribution

        # Top-k selection and renormalized gate weights over the chosen experts.
        topk_vals, topk_idx = probs.topk(self.k, dim=-1)   # (tokens, k)
        gate = topk_vals / topk_vals.sum(dim=-1, keepdim=True)

        # Dispatch: accumulate gated expert outputs into the result.
        y = torch.zeros_like(x)
        # Per-expert dispatch fraction f_i (hard counts), for the aux loss.
        f = torch.zeros(self.num_experts, device=x.device)
        for e in range(self.num_experts):
            # Which (token, slot) pairs selected expert e?
            mask = topk_idx == e                       # (tokens, k) bool
            token_sel = mask.any(dim=-1)               # (tokens,) bool
            if token_sel.any():
                # Gate weight this expert contributes to each selected token.
                w = (gate * mask).sum(dim=-1)[token_sel]   # (n_sel,)
                out = self.experts[e](x[token_sel])        # (n_sel, d_model)
                y[token_sel] += w.unsqueeze(-1) * out
            f[e] = token_sel.float().mean()

        # Load-balancing auxiliary loss: alpha * E * sum_i f_i * P_i.
        P = probs.mean(dim=0)                          # (num_experts,)
        aux_loss = self.aux_alpha * self.num_experts * torch.sum(f * P)
        return y, aux_loss, {"f": f.detach(), "P": P.detach()}


d_model, d_hidden = 32, 64
num_experts, k = 8, 2
moe = SparseMoE(d_model, d_hidden, num_experts, k)

n_params = sum(p.numel() for p in moe.parameters())
expert_params = sum(p.numel() for p in moe.experts[0].parameters())
print(f"experts: {num_experts}, top-k: {k}")
print(f"total params:  {n_params:,}")
print(f"active params per token (router + {k} experts): "
      f"{moe.router.weight.numel() + k * expert_params:,}")

experts: 8, top-k: 2
total params:  33,792
active params per token (router + 2 experts): 8,640

The active-parameter count is far below the total: every token uses the router plus \(k\) of the \(E\) experts, so adding experts grows capacity without growing per-token compute. This is the decoupling from Section 2 made concrete.

244.7.2 7.2 Routing and Utilization on Synthetic Input

We push a batch of synthetic tokens through the layer and inspect the routing. Two checks matter. First, the output shape must match the input, the MoE is a drop-in replacement for a dense FFN. Second, we measure expert utilization, the fraction of tokens each expert receives, which is what the load-balancing loss is trying to flatten. On a randomly initialized router the load is already somewhat uneven, which is exactly the situation the auxiliary loss is designed to correct during training.

torch.manual_seed(1)
T = 512                                  # tokens in the batch
x = torch.randn(T, d_model)

y, aux_loss, stats = moe(x)

print(f"input shape:  {tuple(x.shape)}")
print(f"output shape: {tuple(y.shape)}")
print(f"aux load-balancing loss: {aux_loss.item():.4f}")

f = stats["f"]                           # fraction of tokens per expert
util_pct = (f / f.sum() * 100).tolist()
print("\nexpert utilization (% of dispatched tokens):")
for e, u in enumerate(util_pct):
    bar = "#" * int(round(u))
    print(f"  expert {e}: {u:5.1f}%  {bar}")

ideal = 100.0 / num_experts
spread = max(util_pct) - min(util_pct)
print(f"\nideal uniform share: {ideal:.1f}% per expert")
print(f"observed spread (max - min): {spread:.1f} percentage points")

input shape:  (512, 32)
output shape: (512, 32)
aux load-balancing loss: 0.0201

expert utilization (% of dispatched tokens):
  expert 0:  11.7%  ############
  expert 1:  13.5%  #############
  expert 2:  14.3%  ##############
  expert 3:  11.0%  ###########
  expert 4:  14.8%  ###############
  expert 5:  12.6%  #############
  expert 6:  10.9%  ###########
  expert 7:  11.1%  ###########

ideal uniform share: 12.5% per expert
observed spread (max - min): 3.9 percentage points

The spread between the busiest and idlest expert is the imbalance the auxiliary loss penalizes. A perfectly balanced router would put every bar at the ideal share; a collapsed router would pile almost all tokens onto one or two experts.

244.7.3 7.3 Training Drives the Router Toward Balance

To show the load-balancing loss working, we train the layer for a short loop on a trivial reconstruction objective (make the output match a fixed linear target), adding the auxiliary loss to the task loss. We track the imbalance over steps. The point is not the task itself but that the combined loss flattens the expert distribution, the mechanism from Section 4 in action.

torch.manual_seed(2)

# A small, fixed reconstruction task so the run is fast and deterministic.
target_map = nn.Linear(d_model, d_model)
for p in target_map.parameters():
    p.requires_grad_(False)

model = SparseMoE(d_model, d_hidden, num_experts, k, aux_alpha=1e-2)
opt = torch.optim.Adam(model.parameters(), lr=1e-3)

data = torch.randn(256, d_model)
with torch.no_grad():
    targets = target_map(data)


def imbalance(f):
    # Coefficient of variation of the per-expert load; 0 means perfectly even.
    return (f.std(unbiased=False) / (f.mean() + 1e-9)).item()


print(f"{'step':>5} {'task_loss':>10} {'aux_loss':>9} {'imbalance':>10}")
for step in range(201):
    y, aux, stats = model(data)
    task_loss = F.mse_loss(y, targets)
    loss = task_loss + aux
    opt.zero_grad()
    loss.backward()
    opt.step()
    if step % 50 == 0:
        print(f"{step:5d} {task_loss.item():10.4f} {aux.item():9.4f} "
              f"{imbalance(stats['f']):10.4f}")

 step  task_loss  aux_loss  imbalance
    0     0.3657    0.0200     0.0630
   50     0.0903    0.0200     0.0816
  100     0.0175    0.0201     0.1376
  150     0.0055    0.0200     0.1185
  200     0.0031    0.0200     0.1159

The task loss falls steadily while the auxiliary loss holds the routing in check, which is the behavior we want: the primary objective improves without the router being free to collapse onto a few experts. The cleaner demonstration that the auxiliary term matters is the controlled ablation in the next subsection, where the same run with the loss removed drifts to a markedly worse imbalance. In a real model the auxiliary coefficient \(\alpha\) is tuned so that it balances load without overwhelming the language-modeling loss.

244.7.4 7.4 The Effect of the Auxiliary Loss

A final controlled comparison isolates the auxiliary loss. We train two routers from the same initialization, one with the load-balancing loss and one without, and compare the final imbalance. Turning the auxiliary loss off should let the router drift toward collapse.

torch.manual_seed(3)


def train_router(use_aux: bool, steps: int = 200):
    torch.manual_seed(7)               # identical init for a fair comparison
    m = SparseMoE(d_model, d_hidden, num_experts, k, aux_alpha=1e-2)
    o = torch.optim.Adam(m.parameters(), lr=1e-3)
    d = torch.randn(256, d_model)
    with torch.no_grad():
        t = target_map(d)
    final = None
    for _ in range(steps):
        y, aux, stats = m(d)
        loss = F.mse_loss(y, t) + (aux if use_aux else 0.0)
        o.zero_grad()
        loss.backward()
        o.step()
        final = stats["f"]
    return imbalance(final)


imb_with = train_router(use_aux=True)
imb_without = train_router(use_aux=False)
print(f"final imbalance WITH    aux loss: {imb_with:.4f}")
print(f"final imbalance WITHOUT  aux loss: {imb_without:.4f}")
print(f"the auxiliary loss reduces imbalance by "
      f"{(1 - imb_with / imb_without) * 100:.1f}%"
      if imb_without > imb_with else
      "the runs ended at similar imbalance on this toy task")

final imbalance WITH    aux loss: 0.0690
final imbalance WITHOUT  aux loss: 0.1642
the auxiliary loss reduces imbalance by 58.0%

This is the load-balancing loss earning its place: with the same data, optimizer, and initialization, the only difference is the auxiliary term, and it produces a measurably more even distribution of work across experts.

244.7.5 7.5 Scaling Up: Switch Transformer and Mixtral (requires GPU / large model)

The CPU layer above is the same design that runs at frontier scale, but training or serving a real MoE model needs accelerators, sharded experts, and optimized dispatch kernels. The code below is the genuine library invocation a reader would use on a GPU; it is shown, not executed, because no GPU is available in this environment.

Loading and running Mixtral 8x7B (a sparse MoE with eight experts and top-2 routing) through Hugging Face Transformers:

# requires GPU / large model: Mixtral 8x7B needs multi-GPU or heavy offload.
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

model_id = "mistralai/Mixtral-8x7B-v0.1"
tok = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.bfloat16,
    device_map="auto",            # shards experts across available GPUs
)

inputs = tok("Mixture of experts routes each token to", return_tensors="pt").to(model.device)
out = model.generate(**inputs, max_new_tokens=40)
print(tok.decode(out[0], skip_special_tokens=True))

Hugging Face also exposes the Switch Transformer directly, including its top-1 routing and capacity-based token dropping:

# requires GPU / large model: Switch-Base has 8 to 128 experts per layer.
from transformers import SwitchTransformersForConditionalGeneration, AutoTokenizer

tok = AutoTokenizer.from_pretrained("google/switch-base-8")
model = SwitchTransformersForConditionalGeneration.from_pretrained(
    "google/switch-base-8", device_map="auto", torch_dtype="bfloat16"
)
inputs = tok("The <extra_id_0> walks in the park", return_tensors="pt").to(model.device)
print(tok.decode(model.generate(**inputs)[0], skip_special_tokens=True))

Memory pressure is the binding constraint for MoE inference, because all experts must be resident even though each token uses only \(k\) of them. Quantization with bitsandbytes is the usual mitigation; it is shown here only, since quantized kernels target the GPU:

# requires GPU: bitsandbytes 4-bit quantization to fit a large MoE in memory.
from transformers import AutoModelForCausalLM, BitsAndBytesConfig
import torch

quant = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.bfloat16,
    bnb_4bit_quant_type="nf4",
)
model = AutoModelForCausalLM.from_pretrained(
    "mistralai/Mixtral-8x7B-v0.1", quantization_config=quant, device_map="auto"
)

244.7.6 7.6 Expert-Parallel Training (requires GPU cluster)

At training scale the experts of a single layer are distributed across devices, a scheme called expert parallelism. Each device holds a subset of the experts, and an all-to-all collective shuffles tokens to the device that owns their selected expert, runs the expert, then shuffles the results back. The mature open-source tooling for this is DeepSpeed-MoE and Megatron-LM; the snippet shows the DeepSpeed MoE layer wrapping the same expert-and-gate structure built above, configured for expert parallelism across a GPU cluster.

# requires GPU cluster: expert-parallel MoE across multiple devices.
import deepspeed
import torch.nn as nn

expert = nn.Sequential(nn.Linear(d_model, d_hidden), nn.GELU(),
                       nn.Linear(d_hidden, d_model))

moe_layer = deepspeed.moe.layer.MoE(
    hidden_size=d_model,
    expert=expert,
    num_experts=64,            # total experts across the cluster
    ep_size=8,                 # experts sharded over 8 GPUs (expert parallel)
    k=2,                       # top-2 routing
    use_residual=False,
)
# The all-to-all token shuffle and the load-balancing loss are handled
# inside the layer; the returned aux loss is added to the task loss.
output, aux_loss, _ = moe_layer(hidden_states)

The conceptual model is unchanged from the CPU implementation: route to top-k, gate, combine, and balance the load. Expert parallelism is purely a systems answer to the fact that the experts no longer fit on one device.

244.7.7 7.7 Extreme Sparsity at Frontier Scale

The toy layer above uses \(E = 8\) experts with \(k = 2\), the Mixtral-style setting that most tutorials adopt. A genuine frontier MoE looks nothing like that. Kimi Team (2026) introduce Kimi K3, a 2.8 trillion parameter mixture of experts with roughly 104 billion activated parameters per token, native vision, and a one-million-token context window. Its Stable LatentMoE block, in the architecture description of that report, activates 16 of 896 routed experts per token, and the team reports about a 2.5x improvement in overall scaling efficiency over its predecessor together with fully released model weights, which makes it a rare case where the architecture of a frontier-class system is inspectable rather than inferred.

The number worth pausing on is the sparsity. Write the routing activation ratio as \(\rho_{\text{route}} = k/E\) and the active parameter fraction as \(\rho_{\text{param}} = P_{\text{act}}/P_{\text{total}}\). For K3,

\[ \rho_{\text{route}} = \frac{16}{896} = \frac{1}{56} \approx 1.79\%, \qquad \rho_{\text{param}} = \frac{104 \times 10^{9}}{2.8 \times 10^{12}} \approx 3.71\%, \]

against the textbook baseline of \(2/8 = 25\%\) and Mixtral’s \(13/47 \approx 27.7\%\). Each comparison has its own baseline, so it is worth stating them separately. At the router, K3 is about fourteen times sparser than the 2-of-8 layer our code cells simulate (\(25\% / 1.79\% \approx 14.0\)). In parameters, K3 is about seven and a half times sparser than Mixtral (\(27.7\% / 3.71\% \approx 7.5\)), which is the right comparison because the toy layer has no full model attached and therefore no parameter fraction of its own. The active parameter fraction sits above the routing ratio for the usual reason: attention, embeddings, and any always-on dense components are paid by every token regardless of routing, so they dilute the sparsity of the expert bank.

Extreme sparsity is not free, and the first bill arrives at load balancing. Take a batch of \(N\) tokens routed top-\(k\) over \(E\) experts. Even with a perfectly unbiased router, the indicator that token \(t\) selects expert \(i\) is Bernoulli with probability \(k/E\), so the per-expert load \(L_i\) is marginally binomial:

\[ \mu = \mathbb{E}[L_i] = \frac{Nk}{E}, \qquad \mathrm{Var}[L_i] = N\frac{k}{E}\Big(1 - \frac{k}{E}\Big), \qquad \mathrm{CV} = \sqrt{\frac{1 - k/E}{Nk/E}} \approx \frac{1}{\sqrt{\mu}}. \]

Holding \(N\) and \(k\) fixed and growing \(E\) shrinks \(\mu\) proportionally, so the coefficient of variation grows like \(\sqrt{E}\). At \(N = 8192\) tokens, the 2-of-8 layer gives \(\mu = 2048\) and, from the exact expression rather than the \(1/\sqrt{\mu}\) approximation, \(\mathrm{CV} = 1.9\%\); 16-of-896 gives \(\mu \approx 146.3\) and \(\mathrm{CV} = 8.2\%\). The approximation is fine once \(k/E\) is small, but it overstates the spread by a factor \(1/\sqrt{1 - k/E} \approx 1.15\) at \(k/E = 0.25\), which is how \(1.9\%\) turns into the \(2.2\%\) one gets from \(1/\sqrt{\mu}\); the table below prints the exact formula next to the measured value. This is pure sampling noise, before any of the routing collapse of Section 4 enters the picture. It is exactly why very sparse MoE needs stronger balancing than the auxiliary loss alone was designed for.

The payoff is a straggler argument for expert parallelism. In the all-to-all scheme of Section 7.6, a training step cannot finish until the most heavily loaded expert finishes, so step time is set by \(\max_i L_i\), not by the mean. Treating the loads as approximately Gaussian and weakly dependent, the expected maximum of \(E\) such variables is about \(\mu\,(1 + \mathrm{CV}\sqrt{2\ln E})\), giving a straggler factor

\[ S = \frac{\max_i L_i}{\mu} \approx 1 + \sqrt{\frac{2 \ln E}{\mu}}. \]

Both terms move the wrong way as the model gets sparser: \(\ln E\) rises in the numerator while \(\mu = Nk/E\) falls in the denominator. At \(N = 8192\) the formula gives \(S \approx 1.04\) for the 2-of-8 layer and \(S \approx 1.30\) for 16-of-896; the simulation below measures \(1.025\) and \(1.263\), because the Gaussian maximum runs slightly high for counts this discrete and for the mild negative dependence top-\(k\) routing induces between experts. Read \(S\) as an overhead on step time, not as an idle fraction: a measured \(S = 1.263\) means every step takes twenty-six percent longer than the balanced ideal, so the average device spends \(1 - 1/1.263 \approx 21\%\) of the step waiting on the slowest expert. A fifth of a frontier training fleet sitting idle is an enormous amount of hardware to leave on the table, and it is why balanced expert-parallel training is a real engineering claim in that work rather than a throwaway phrase. The same arithmetic sets the capacity factor of Section 5: the headroom needed to avoid dropping tokens is essentially \(S\), so extreme sparsity forces either a larger capacity factor (padding waste) or a smarter balancer.

The second bill arrives at inference, where extreme sparsity shifts the memory-to-FLOP ratio. An expert GEMM processing \(m\) tokens in bfloat16 does about \(2mP_e\) FLOPs while reading \(2P_e\) bytes of weights, so its arithmetic intensity is simply \(m\) FLOPs per byte, and \(m = Nk/E\) is the tokens-per-expert figure above. A modern accelerator has a ridge point near 300 FLOPs per byte, so staying compute bound requires roughly \(N \gtrsim 300\,E/k\) tokens in flight per expert-parallel group: about 1,200 for 2-of-8 but about 16,800 for 16-of-896. Below that, the layer is bandwidth bound and the sparsity buys latency but not throughput. Sparse models are therefore batch-hungry in a way dense models are not.

The simulation below routes synthetic tokens with an unbiased router at four sparsity settings and measures the load spread and the straggler factor directly. The printed table is taken at the \(N = 8192\) batch quoted above, with the measured coefficient of variation averaged over trials and the closed-form predictions alongside it.

import numpy as np
import matplotlib.pyplot as plt

rng = np.random.default_rng(0)

SETTINGS = [(8, 2, "2 of 8 (Mixtral-style)"),
            (64, 8, "8 of 64"),
            (256, 8, "8 of 256"),
            (896, 16, "16 of 896 (Kimi K3)")]
BATCHES = [1024, 4096, 8192, 16384, 65536]
REF = 8192                      # the batch size quoted in the text above
TRIALS, CHUNK = 4, 4096


def expert_loads(n_tokens, E, k, rng):
    """Token counts per expert for an unbiased top-k router."""
    counts = np.zeros(E, dtype=np.int64)
    done = 0
    while done < n_tokens:
        m = min(CHUNK, n_tokens - done)
        logits = rng.standard_normal((m, E), dtype=np.float32)
        top = np.argpartition(-logits, k - 1, axis=1)[:, :k]
        counts += np.bincount(top.ravel(), minlength=E)
        done += m
    return counts


curves, ref_runs = {}, {}
for E, k, name in SETTINGS:
    row = []
    for n in BATCHES:
        runs = [expert_loads(n, E, k, rng) for _ in range(TRIALS)]
        row.append(float(np.mean([c.max() / c.mean() for c in runs])))
        if n == REF:
            ref_runs[name] = runs
    curves[name] = row

hdr = (f"{'setting':>23} {'k/E':>7} {'tok/exp':>8} {'CV':>7} {'CV thy':>7} "
       f"{'S':>6} {'S thy':>6}")
print(f"at N = {REF:,} tokens per batch, {TRIALS} trials each\n{hdr}")
for E, k, name in SETTINGS:
    runs = ref_runs[name]                       # CV averaged over trials
    mu = REF * k / E
    cv = float(np.mean([c.std(ddof=1) / c.mean() for c in runs]))
    s = float(np.mean([c.max() / c.mean() for c in runs]))
    cv_thy = np.sqrt((1 - k / E) / mu)
    s_thy = 1 + cv_thy * np.sqrt(2 * np.log(E))
    print(f"{name:>23} {100 * k / E:6.2f}% {mu:8.1f} {cv:7.3f} {cv_thy:7.3f} "
          f"{s:6.3f} {s_thy:6.3f}")

print(f"\nstraggler factor at N = {BATCHES[-1]:,}: " + ", ".join(
    f"{name} {curves[name][-1]:.3f}" for _, _, name in SETTINGS))
print(f"active parameter fraction, Kimi K3:      {100 * 104e9 / 2.8e12:.2f}%")
print(f"active parameter fraction, Mixtral 8x7B:  {100 * 13 / 47:.2f}%")
for E, k, name in SETTINGS:
    print(f"batch needed to stay compute bound, {name}: {int(300 * E / k):,} tokens")

fig, ax = plt.subplots(figsize=(7.0, 4.2))
for _, _, name in SETTINGS:
    ax.plot(BATCHES, curves[name], marker="o", label=name)
ax.axhline(1.0, ls="dashed", lw=1, color="gray")
ax.set_xscale("log", base=2)
ax.set_xlabel("tokens per batch (per expert-parallel group)")
ax.set_ylabel("straggler factor (max load / mean load)")
ax.set_title("Load imbalance from routing noise alone")
ax.legend(frameon=False)
fig.tight_layout()
plt.show()

at N = 8,192 tokens per batch, 4 trials each
                setting     k/E  tok/exp      CV  CV thy      S  S thy
 2 of 8 (Mixtral-style)  25.00%   2048.0   0.019   0.019  1.025  1.039
                8 of 64  12.50%   1024.0   0.029   0.029  1.068  1.084
               8 of 256   3.12%    256.0   0.061   0.062  1.179  1.205
    16 of 896 (Kimi K3)   1.79%    146.3   0.082   0.082  1.263  1.302

straggler factor at N = 65,536: 2 of 8 (Mixtral-style) 1.011, 8 of 64 1.023, 8 of 256 1.064, 16 of 896 (Kimi K3) 1.092
active parameter fraction, Kimi K3:      3.71%
active parameter fraction, Mixtral 8x7B:  27.66%
batch needed to stay compute bound, 2 of 8 (Mixtral-style): 1,200 tokens
batch needed to stay compute bound, 8 of 64: 2,400 tokens
batch needed to stay compute bound, 8 of 256: 9,600 tokens
batch needed to stay compute bound, 16 of 896 (Kimi K3): 16,800 tokens

Figure 244.1: Straggler factor (max expert load divided by mean expert load) under a perfectly unbiased top-k router. Sparser configurations need far larger batches before routing noise averages out.

The 2-of-8 curve falls from 1.082 at 1,024 tokens to about 1.03 by 4,096, sits near that level through 16,384, and only reaches 1.011 at 65,536. Getting inside a couple of percent takes tens of thousands of tokens even in the easy configuration, but the tax is small enough throughout that the balancing problem looks solved in small demonstrations. The 16-of-896 curve starts at 1.914, near a factor of two, and is still paying a nine percent tax at 65,536 tokens per batch, a batch size no single device will see on its own. That gap is why frontier systems push balance into the router design, the capacity policy, and the parallelism strategy rather than relying on the auxiliary loss alone.

This closes the loop on the tradeoff that opened the chapter. MoE decouples total capacity from per-token compute, and \(\rho_{\text{route}}\) is the dial that sets how aggressively you exploit that decoupling. Turning it down toward two percent buys enormous capacity per unit of compute, which is precisely the scaling efficiency K3 reports. What you pay for it is not accuracy but engineering surface: a routing distribution whose noise floor grows like \(\sqrt{E}\), a training step whose duration is set by the unluckiest expert, an inference regime that only reaches peak throughput at very large batch, and a memory footprint that still holds all 2.8 trillion parameters even though any one token touches under four percent of them. Sparsity is a compute discount financed by a systems debt, and the size of that debt is what separates a textbook MoE layer from a frontier one.

244.8 8. Pitfalls and When to Use

Routing collapse is the failure mode to watch. Without an effective load-balancing loss the router concentrates tokens on a few experts and the rest never train. Monitor per-expert utilization during training, not just the loss curve; a healthy run keeps every expert busy. The auxiliary coefficient \(\alpha\) is the main knob: too small and load drifts toward collapse, too large and balancing fights the language-modeling objective and degrades quality. The Switch value of \(10^{-2}\) is a reasonable starting point.

Capacity factor trades memory against dropped tokens. A capacity factor near 1.0 minimizes padding but drops tokens whenever load is uneven, and dropped tokens get no expert computation at all, which hurts quality. A higher factor wastes compute on padding. The right value depends on how well balanced the router is, so capacity and the auxiliary loss must be tuned together.

MoE is memory-bound at inference even though it is compute-light. Every expert must be resident in memory, so a model with 47 billion total parameters needs memory for all of them even while computing as if it had 13 billion. This is why MoE shines for training-time efficiency and high-throughput serving but is awkward on memory-constrained single devices. Quantization and expert offloading are the standard mitigations.

Fine-tuning behaves differently from dense models. Sparse MoE models can overfit faster on small downstream datasets because each expert sees only a slice of the data, and the router can shift its routing under fine-tuning in ways that destabilize training. Lower learning rates, careful regularization, and sometimes freezing the router are common practice.

When to reach for MoE. Use it when you want to grow model capacity under a fixed per-token compute budget, typically at large scale where the parameter count is the bottleneck and you have the memory to hold all experts and the infrastructure for expert-parallel training or high-throughput batched serving. Prefer a dense model when memory is the binding constraint, when the deployment is latency-sensitive on a single small device, or when the model is small enough that the routing machinery and its training instabilities are not worth the added capacity. MoE is a scaling tool: it pays off precisely when you have more parameters to spend than compute to spend on them, and it asks in return that you take load balancing seriously.

244.9 References

1. Shazeer, N. et al. “Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer.” 2017. https://arxiv.org/abs/1701.06538
2. Fedus, W., Zoph, B., and Shazeer, N. “Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity.” Journal of Machine Learning Research, 2022. https://arxiv.org/abs/2101.03961
3. Lepikhin, D. et al. “GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding.” 2020. https://arxiv.org/abs/2006.16668
4. Jiang, A. Q. et al. “Mixtral of Experts.” 2024. https://arxiv.org/abs/2401.04088
5. DeepSeek-AI. “DeepSeek-V3 Technical Report.” 2024. https://arxiv.org/abs/2412.19437
6. Zoph, B. et al. “ST-MoE: Designing Stable and Transferable Sparse Expert Models.” 2022. https://arxiv.org/abs/2202.08906
7. Rajbhandari, S. et al. “DeepSpeed-MoE: Advancing Mixture-of-Experts Inference and Training to Power Next-Generation AI Scale.” 2022. https://arxiv.org/abs/2201.05596
8. Zhou, Y. et al. “Mixture-of-Experts with Expert Choice Routing.” Advances in Neural Information Processing Systems, 2022. https://arxiv.org/abs/2202.09368
9. Bercovich, A. et al. “Nemotron-Labs-3-Puzzle-75B-A9B: Compressing Hybrid MoE LLMs.” 2026. https://arxiv.org/abs/2607.04371
10. Kimi Team. “Kimi K3: Open Frontier Intelligence.” 2026. https://arxiv.org/abs/2607.24653

Kimi Team. 2026. “Kimi K3: Open Frontier Intelligence.” arXiv Preprint arXiv:2607.24653.