Mixture of Experts: Routing, Dispatch, and Load Balancing

A Mixture of Experts (MoE) layer replaces the standard MLP inside a decoder block with multiple feed-forward sub-networks — the experts — and a trainable router that selects which experts process each token. The architectural promise is simple: store many expert parameters but activate only a small subset per token, so total model capacity grows without a proportional increase in per-token compute. The engineering reality is more nuanced: dispatch, capacity limits, communication, and load balancing all change what “sparse” actually means at inference time.

This lesson treats Lesson 15 (mha-mqa-gqa) as an attention prerequisite, not as a source for MoE claims. We build the routing mechanism from primary sources and keep the dense-to-MoE conversion brief and optional.

The one-sentence version

A decoder-only MoE layer sends each token through a router that computes expert scores; selects top-k experts; dispatches the token to those experts; computes their outputs; and combines them into a normalised weighted result — with total stored parameters equal to all experts, but active per-token compute limited to the selected subset.

Everything else in this lesson makes that concrete.

1. What the router actually does

The router (also called the gate) is a small trainable network — usually a linear projection from the hidden dimension d to the number of experts N, followed by softmax. For an input token x of shape [d]:

router_scores = softmax(x @ W_router + b_router)   # [N]

Mathematically, W_router has shape [d, N] (input dimension d, output dimension N). In PyTorch, nn.Linear(d, N) stores the weight tensor as [N, d] — that is, [out_features, in_features]. The mathematical notation [d, N] describes the conceptual matrix; the stored tensor shape [N, d] is PyTorch’s convention. The output is a probability vector over N experts. Nothing in this step implies any expert acquires a human-readable speciality; the router is optimised jointly with the rest of the model for predictive performance, not for interpretable categories.

From Shazeer et al. (2017), the original sparsely-gated MoE used top-k routing with separate load-balancing and importance-weighting losses (§2.1, §2.3). GShard (Lepikhin et al., 2020) adopted top-2 routing with a capacity factor (§3.1) and showed results on multilingual translation, not universal benchmarks. Switch Transformer (Fedus et al., 2022) simplified to top-1 routing with a single auxiliary load-balancing loss (§2.2) and evaluated on masked language modelling for the C4 corpus. We cite the specific design choices from those papers; we do not generalise their quality results beyond the stated models and tasks.

No guaranteed expert specialities

There is no reliable evidence that individual experts learn stable, human-readable categories (syntax, entities, numbers). The router distributes tokens based on predictive utility, not semantic labels. Any claim that “Expert 3 handles dates” requires independent verification on the specific model and dataset; it is not a property of the architecture.

2. Token-to-output trace with shapes

Consider a single layer with N = 4 experts, hidden dimension d = 512, and k = 2 (top-2 routing, as in GShard). A token x enters the layer:

Step	Operation	Input shape	Output shape	Key note
1	x (token hidden state)	[d]	[d]	Residual stream entry
2	W_router projection	[d]	[N]	W: [d, N]
3	Softmax scores	[N]	[N]	Probability vector
4	Top-k selection	[N]	k indices	k = 2 here
5	Dispatch to expert i	[d]	[d]	Only for selected i
6	Expert FFN (i)	[d]	[d]	Same structure as MLP
7	Collect expert outputs	k outputs of [d]	k tensors of [d]	One per selected expert
8	Normalised weighted sum	k outputs + k scores	[d]	Normalise by score sum
9	Residual addition	[d] + [d]	[d]	x + output (pre-norm or post-norm)

The output of step 9 replaces the standard MLP output inside the decoder block; the attention sub-layer and residual connections remain unchanged.

MoE routing flow

Router scores, top-k selection, expert dispatch, and normalised weighted output. Only selected experts are active; others contribute nothing to this token's compute.

Read this diagram as text

Text equivalent: token enters router; router produces scores over experts; top-2 selected; token dispatched to those experts; expert outputs combined with weights from the router scores, normalised; result added back to the residual stream.

3. Minimal reproducible example (pseudocode, verified by inspection)

The following Python-style pseudocode defines a minimal router, dispatch, and combination. It is labelled as pseudocode and intended for inspection, not as a production training loop.

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

torch.manual_seed(42)

# Toy dimensions
B, T, d = 2, 4, 8      # batch, sequence, hidden
N = 4                 # number of experts
k = 2                 # top-k routing (GShard-style)

# Router: linear projection to N scores, no bias for simplicity
# Mathematical: W_router is [d, N]. PyTorch stores weight as [N, d] (out_features × in_features).
W_router = nn.Linear(d, N, bias=True)  # weight shape [N, d] in PyTorch storage

# Experts: simple FFN (like standard MLP but separate per expert)
experts = nn.ModuleList([nn.Linear(d, d) for _ in range(N)])

# Forward pass for one token position (illustrative)
x = torch.randn(B, T, d)  # [2, 4, 8]

# Step 1: router scores
router_logits = W_router(x)              # [B, T, N]
router_probs = F.softmax(router_logits, dim=-1)  # [B, T, N]

# Step 2: top-k selection (per token, per position)
# For simplicity we select per token independently
topk_vals, topk_idx = torch.topk(router_probs, k, dim=-1)
# topk_vals: [B, T, k], topk_idx: [B, T, k]

# Step 3: dispatch and expert computation
# In practice this requires a batched gather; here we loop for clarity
outputs = torch.zeros_like(x)  # will accumulate weighted results
for b in range(B):
    for t in range(T):
        token_input = x[b, t, :]  # [d]
        weighted_sum = torch.zeros(d)
        total_weight = 0.0
        for ki in range(k):
            expert_idx = topk_idx[b, t, ki].item()
            weight = topk_vals[b, t, ki].item()
            expert_out = experts[expert_idx](token_input)  # [d]
            weighted_sum += weight * expert_out
            total_weight += weight
        outputs[b, t, :] = weighted_sum / total_weight  # normalised combination

# Step 4: residual addition (assumes pre-norm placement, as in modern decoders)
# Norm(x) would come before the expert layer in a full block; we omit it here
# to focus only on the MoE mechanism.
final_output = x + outputs  # [B, T, d]

Expected result (inspection-based):

• router_probs.sum(dim=-1) is approximately 1.0 (softmax invariant).
• final_output.shape == x.shape == [2, 4, 8] (residual stream preserved).
• Only k = 2 experts contribute per token; the other N - k = 2 experts receive zero weight for that token.
• The division by total_weight ensures the weighted average is properly normalised; without it, the magnitude would scale with the number of selected experts.

This is pseudocode for mechanism illustration only. It does not include training, gradient accumulation, or distributed dispatch.

4. Capacity, overflow, and load imbalance

In GShard (§2.2), each expert has a fixed capacity per group of tokens, and capacity is shared out under top-2 routing. If a token’s selected expert exceeds its capacity, the token is either dropped or redirected. The capacity factor C defines the maximum tokens per expert as C * (total_tokens / N). A C = 1.25 allows a 25% overload per expert before overflow.

Capacity overflow case (labelled):

• Batch: 8 tokens, N = 4, C = 1.0, and top-1 routing so each token produces exactly one assignment (with top-2 the counts double and capacity must double as well).
• Uniform dispatch: 2 tokens per expert.
• Skewed dispatch (simulated load imbalance): 5 tokens to Expert 1, 2 to Expert 2, 1 to Expert 3, 0 to Expert 4 - eight assignments in total.
• Expert 1 capacity = 1.0 * (8 / 4) = 2. Only 2 of its 5 assigned tokens are accepted; 3 are dropped or overflowed.
• Quality effect: dropped tokens lose the expert’s contribution; if the router relies on that expert’s representation, the output changes. The original GShard paper reports that ensuring low drop rates (typically < 1%) is important for scaling (§3); we do not invent new thresholds.

Load imbalance is not only a capacity issue — it is also a communication and memory issue. If Expert 1 resides on a different device, more tokens must be routed to that device, increasing inter-device bandwidth. The total compute per layer remains bounded by k active experts per token, but the actual latency depends on the slowest device’s load.

From Switch Transformer (§2.3), the auxiliary load-balancing loss is:

loss_aux = alpha * N * sum_i (f_i * P_i)

where f_i is the fraction of tokens dispatched to expert i (not differentiable) and P_i is the mean routing probability (differentiable). The paper selects alpha = 1e-2; we state this as the reported value for the C4 masked-language-modelling setup, not as a universal recommendation.

5. Active vs. stored parameters; compute, memory, communication, latency

These distinctions are the core engineering insight of the architecture. We separate them explicitly.

Dimension	Dense MLP (same hidden size)	MoE with N experts, k = 2, same expert size	What changes
Total stored parameters	2 · d · d_ff (i.e. 8d² for d_ff = 4d)	N · (2 · d · d_ff) + (N · d + N) for router; expands linearly with N for experts plus fixed router params	Increases linearly with N
Per-token active params	Full MLP params	Only k experts’ params + router params	Scales with k / N of stored expert params
Shared parameters	All weights shared across tokens	Router shared; expert weights shared within each expert across tokens that select it	No new sharing within an expert
Per-token compute (FLOPs)	One full FFN pass	k expert FFN passes + router pass	Proportional to k, not N
Memory (weights)	One MLP stored	N MLPs + router stored	Linear in N
Memory (KV cache)	Unchanged	Unchanged (MoE affects only the MLP sub-layer)	No effect
Communication (multi-device)	Minimal (local layer)	Token routing across devices for selected experts; larger with N and load imbalance	Increases with expert placement and load
Latency	Predictable	Depends on device load, routing overhead, and overflow handling; not guaranteed faster despite fewer active params	Must be measured, not assumed

Mixtral 8x7B (Jiang et al., 2024) provides a concrete decoder-model reference: 8 experts, top-2 routing per layer, 47B total parameters, 13B active during inference (§Abstract). That is a result for Mixtral 8x7B on its training and evaluation setup; we do not claim it applies to all MoE configurations.

The key takeaway is that reducing active parameters does not automatically reduce latency. If the bottleneck is memory bandwidth for loading expert weights, a larger N can make it worse. If the bottleneck is compute and experts are already on-device, the reduction can help. The choice requires measurement on the target workload.

6. Dense-to-MoE conversion: brief and optional

Converting a dense model to MoE is an architecture change, not a weight edit: selected FFN layers are replaced by MoE layers that contain N experts (this layer-placement change is what GShard §2.1 and Figure 3, and Switch Transformer §3, describe). Those papers train MoE models in that form. Copying a trained dense FFN’s weights into the new experts is a practical initialisation choice, not a recipe established by those papers, and expert diversity has to come from continued training. Converting multi-head attention to grouped-query attention is a different, attention-side procedure from the GQA paper, not from GShard.

Important caveats:

• Copying dense weights gives every expert the same initial function; diversity must emerge during training through the router loss and gradient updates.
• Uptraining is not part of the GShard/Switch MoE conversion recipe. The separate GQA paper (Ainslie et al. 2023, §4.3) shows that uptraining on a small fraction of the original pretraining compute recovers most of the gap when converting MHA to grouped-query attention; do not credit that result to GShard.
• The conversion changes memory footprint, serving infrastructure (expert placement), and evaluation baselines; treat it as a separate architecture decision from the routing design itself.

We do not provide a full conversion recipe. If your goal is to compare architectures defensibly, the comparison framework in Lesson 20 (architecture-comparison) is the appropriate next step.

7. Exit exercise: a defensible architecture decision

Exercise: Routing choice and load-balancing design

1. 01A production decoder-only model serves at batch size 4, sequence length 2K, with latency budget 50 ms/token. Your team proposes replacing MLP layers with MoE: N=8 experts, top-2 routing, expert size equal to original MLP hidden width. State three measured constraints (not guesses) you must verify before accepting this design, and identify which primary source result applies to each.

   Show a hint

   Constraints could include: per-token active parameter count vs. dense; capacity-factor drop rate (<1% per GShard §3); inter-device routing overhead. Cite GShard §3.1 (capacity), Shazeer et al. §2.3 (load-balancing loss design), and Mixtral §Abstract (active/total parameter distinction).

2. 02Given the same N=8, k=2 proposal: describe a capacity-overflow scenario where load imbalance causes token drops. State the overflow condition numerically (capacity formula from GShard §3.1) and explain why the quality impact depends on which tokens are dropped, not only on the drop rate.

   Show a hint

   Capacity = C * (total_tokens / N). For C = 1.0, 8 tokens: 1 token per expert uniform. If 4 tokens route to Expert 1, 2 are dropped. The impact depends on whether the dropped tokens are critical for the current prediction.

3. 03Your evaluation shows a 2% perplexity increase for the MoE model relative to dense on your target task, with 3.5x fewer active MLP parameters per token but 8x more total stored parameters. State a defensible decision (adopt, reject, or conditional) with explicit conditions: memory budget, serving device layout, quality tolerance, and retraining budget. Do not claim expert specialisations justify the choice.

   Show a hint

   A defensible decision includes the measured delta (2%), the stated constraints (memory, latency, quality threshold), and a condition (e.g., "adopt only if multi-GPU placement reduces routing latency to under 10 ms/token").

8. Related lessons

• Lesson 15 (mha-mqa-gqa): attention architecture tradeoffs; read first for the prerequisite attention mechanics, but do not treat it as evidence for MoE routing.
• Lesson 8 (transformer-blocks): the decoder block structure that hosts the MoE layer; the MLP sub-layer replaced by experts is described there.
• Lesson 20 (architecture-comparison): the framework for comparing configurations defensibly, which applies to dense-vs-MoE decisions.

Related lessons

• LLM Internals & Architecture · Modern LLM architecturesMHA vs. MQA vs. GQA: Attention Architecture Tradeoffs14 min
• LLM Internals & Architecture · Building and training a language modelInside the Transformer Block: Residuals and Normalisation12 min

9. Sources and evidence limits

Source	Relevant section	What it supports (scope)
Shazeer et al. (2017), “Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer” (arXiv:1701.06538)	§2.1 (softmax gating), §2.3 (top-k selection, load-balancing and importance-weighting losses)	Original MoE routing design; results for language modelling and machine translation on stated architectures up to 137B parameters
Lepikhin et al. (2020), “GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding” (arXiv:2006.16668)	§2.1 and Figure 3 (MoE layer placement: replacing FFN layers), §2.2 (group capacity and capacity factor C under top-2 routing)	Capacity overflow definition and group capacity; MoE layer placement; results for multilingual translation
Fedus, Zoph & Shazeer (2022), “Switch Transformers” (JMLR 23, arXiv:2101.03961)	§2.2 (simplified auxiliary loss, alpha = 1e-2), §2.3 (top-1 routing, capacity, drop rates < 1%), §3 (C4 masked-language-modelling results)	Simplified load-balancing; capacity handling; reported drop rates; results for C4 masked LM, not universal benchmarks
Jiang et al. (2024), “Mixtral of Experts” (arXiv:2401.04088)	Abstract (8 experts, top-2, 47B total / 13B active), model description	Active vs. total parameter distinction for Mixtral 8x7B; results for Mixtral’s evaluation benchmarks

Verification gaps stated:

• We do not claim that any expert acquires a stable human-readable speciality; this is not supported by the cited papers.
• The minimal pseudocode is verified by shape inspection only; it is not a full training or distributed deployment.
• The dense-to-MoE conversion is described at the design level; actual conversion quality depends on the base model, dataset, and uptraining budget, which we do not measure here.
• Performance and latency claims are model-specific (Mixtral 8x7B, GShard translation, Switch C4); we label each result by its stated workload.