Chapter 27

Mixture of Experts

Top-k routing, load-balancing losses, capacity, shared experts, and auxiliary-loss-free balancing.

A dense feed-forward block spends the same parameters on every token. A mixture-of-experts (MoE) block keeps many feed-forward blocks but activates only a few for each token, so the model has more parameters than it pays for on one forward pass. In 2026, MoE is one of the main ways large language models raise capacity without raising per-token compute by the same factor.

27.1 From one FFN to many experts

A Transformer FFN applies one function to every token:

\operatorname{FFN}(\vx) = \mW_2\,\sigma(\mW_1\vx + \vb_1) + \vb_2 .\tag{27.1}

An MoE layer replaces that one function by $N$ expert FFNs $E_1,\ldots,E_N$ and a router. The router produces logits $\vr$ and probabilities $\vp = \softmax(\vr)$ . If $S$ is the set of selected experts, the combine weights renormalize the selected probabilities:

a_i = \frac{p_i}{\sum_{j\in S} p_j}, \quad i\in S .\tag{27.2}

The output is a weighted sum of only the selected experts:

\vy = \sum_{i\in S} a_i E_i(\vx) .\tag{27.3}

The selection is discrete, so the common engineering path is simple: learn the router from the probabilities, route tokens with top-k, and let the main loss train the selected experts. The listing returns the top-k expert indices, the renormalized combine weights, and the full router probabilities used by auxiliary losses.

The important separation is between total parameters and activated parameters. A dense layer must read every FFN weight for every token. An MoE layer may store many more FFNs, but a token touches only its selected experts. That makes the router part of the model architecture, not just a scheduler: different tokens can learn to use different parameter subspaces. The cost is that tokens in the same sequence no longer follow the same compute path.

Listing 27.1 Top-k routing with renormalized combine weights

def softmax(logits):
    """Row-wise softmax."""
    logits = np.asarray(logits)
    shifted = logits - logits.max(axis=-1, keepdims=True)
    exp = np.exp(shifted)
    return exp / exp.sum(axis=-1, keepdims=True)


def top_k_router(logits, k, selection_bias=None):
    """Return top-k expert indices and renormalized combine weights."""
    logits = np.asarray(logits)
    scores = logits if selection_bias is None else logits + selection_bias
    order = np.argsort(scores, axis=-1)[:, ::-1]
    experts = order[:, :k]
    probabilities = softmax(logits)
    chosen = np.take_along_axis(probabilities, experts, axis=-1)
    weights = chosen / chosen.sum(axis=-1, keepdims=True)
    return experts, weights, probabilities

27.2 Forward pass and capacity

A practical implementation groups tokens by expert, runs each expert on its local batch, and scatters the weighted results back. The code below is intentionally direct: it loops over experts, not over tokens, so the tests can compare it with a dense token-by-token reference. With no capacity limit, this sparse implementation and the dense loop compute exactly the same $\vy$ from (27.3).

Listing 27.2 Sparse MoE forward pass with optional token dropping

def expert_mlp(x, w1, b1, w2, b2):
    """One ReLU feed-forward expert."""
    hidden = np.maximum(x @ w1 + b1, 0)
    return hidden @ w2 + b2


def sparse_moe_forward(x, w1, b1, w2, b2, experts, weights, capacity=None):
    """Apply selected experts; optionally drop assignments beyond capacity."""
    num_tokens, _ = x.shape
    num_experts, _, out_dim = w2.shape
    y = np.zeros((num_tokens, out_dim), dtype=x.dtype)
    dropped = np.zeros(experts.shape, dtype=bool)
    for expert in range(num_experts):
        token, slot = np.nonzero(experts == expert)
        if capacity is not None and len(token) > capacity:
            dropped[token[capacity:], slot[capacity:]] = True
            token, slot = token[:capacity], slot[:capacity]
        if len(token) == 0:
            continue
        out = expert_mlp(x[token], w1[expert], b1[expert], w2[expert], b2[expert])
        y[token] += weights[token, slot, None] * out
    return y, dropped

Routing creates a systems problem: one expert may receive far more tokens than another. MoE layers therefore set an expert capacity, usually a per-batch token budget. Assignments beyond that budget are dropped or sent to a backup path. Dropping is ugly but useful: it bounds memory, all-to-all communication, and latency even when the router is temporarily skewed.

Capacity is counted on assignments, not on original tokens. With top-k routing, one token can consume several expert slots. The combine weights should not be renormalized after a drop unless the implementation explicitly wants the remaining experts to compensate; the code here treats a dropped assignment as a zero contribution. That convention makes the test easy to reason about: without dropping, the grouped implementation must match a dense loop; with dropping, only the overflowing assignments disappear.

27.3 Balancing the router

The Switch Transformer uses top-1 routing and adds a load-balancing loss [fedus2021switch]. Let $f_i$ be the fraction of tokens actually sent to expert $i$ , and let $P_i$ be the mean router probability for that expert. The loss is

\mathcal{L}_{\text{lb}} = N \sum_{i=1}^{N} f_i P_i .\tag{27.4}

Why does this prefer uniform routing? In the intended fixed point, assignments match the router probabilities, so $f_i = P_i = p_i$ . Then $\mathcal{L}_{\text{lb}} = N\sum_i p_i^2$ . Since $(\sum_i p_i)^2 \le N\sum_i p_i^2$ , the loss is at least $1$ , with equality at $p_i = 1/N$ . A collapsed router has loss $N$ . The test suite checks both cases.

The same listing includes the router z-loss used by ST-MoE [zoph2022stmoe]. It penalizes a large router log-partition, $\log \sum_i e^{r_i}$ , keeping router logits small enough that bfloat16 and softmax do not fight each other.

The two balancing terms act at different places. The Switch loss changes where probability mass goes, because it multiplies the realized load by the router’s mean probability. The z-loss does not prefer one expert over another; it only discourages the router from making all logits large in magnitude. In a tiny NumPy model both look like scalar regularizers, but in a distributed model the load term protects devices from idle-or-overloaded extremes while the z-loss protects numerics.

Listing 27.3 Switch load balancing and router z-loss

def switch_load_balancing_loss(probabilities, chosen_expert):
    """Switch loss E * sum_i f_i P_i for top-1 routing."""
    probabilities = np.asarray(probabilities)
    num_experts = probabilities.shape[-1]
    counts = np.bincount(chosen_expert, minlength=num_experts)
    load = counts / chosen_expert.size
    mean_probability = probabilities.mean(axis=0)
    loss = num_experts * np.sum(load * mean_probability)
    return float(loss), load, mean_probability


def router_z_loss(logits):
    """Mean squared log-partition of router logits."""
    logits = np.asarray(logits)
    maximum = logits.max(axis=-1, keepdims=True)
    log_z = np.log(np.exp(logits - maximum).sum(axis=-1)) + maximum[:, 0]
    return float(np.mean(log_z ** 2))

27.4 Fine-grained and auxiliary-loss-free experts

DeepSeekMoE splits each large expert into finer-grained experts and adds shared experts that all tokens use, so some capacity is specialized and some is always available [dai2024deepseekmoe]. DeepSeek-V2 applies that idea in a language model setting [shao2024deepseekv2]. Auxiliary-loss-free balancing takes a different route: keep a per-expert bias that is used only for top-k selection, not for the combine probabilities, and update it from load error [wang2024auxiliarylossfree]:

b_i \leftarrow b_i - \eta\,\sign(f_i - 1/N) .\tag{27.5}

Overloaded experts get a smaller selection score next batch; underloaded experts get a larger one. Because the bias is not in $\softmax(\vr)$ , it steers routing without changing the probability weights that train the router. The simulation in the tests starts with a skewed router and verifies that this sign update sharply reduces the load gap.

This bias rule is deliberately coarse. It does not estimate a gradient of the language-model loss, and it does not say which expert would have produced the best token representation. It only says the current batch used some experts too often. That is enough for a feedback controller: lower selection scores for overloaded experts, raise them for underloaded experts, and leave the main objective to decide what the experts learn once tokens arrive.

Listing 27.4 Auxiliary-loss-free selection bias update

def update_selection_bias(bias, load, target, rate):
    """Lower overloaded experts and raise underloaded experts."""
    return bias - rate * np.sign(load - target)


def simulate_bias_balancing(base_logits, steps=80, rate=0.05):
    """Balance a fixed skewed router by changing only its selection bias."""
    bias = np.zeros(base_logits.shape[1], dtype=base_logits.dtype)
    target = np.full(base_logits.shape[1], 1 / base_logits.shape[1])
    history = []
    for _ in range(steps):
        experts, _, probabilities = top_k_router(base_logits, 1, bias)
        _, load, _ = switch_load_balancing_loss(probabilities, experts[:, 0])
        history.append(load)
        bias = update_selection_bias(bias, load, target, rate)
    return np.array(history), bias

In practice

Early MoE layers showed that conditional computation could grow parameter count without a matching increase in activated FLOPs [shazeer2017outrageously]. GShard made expert routing a sharded Transformer primitive [lepikhin2020gshard], and Switch simplified routing to one expert per token with an explicit balancing loss [fedus2021switch]. Recent MoE language models often combine sparse experts with stabilizers such as z-loss, capacity rules, shared experts, and balancing mechanisms that reduce or remove auxiliary losses [dai2024deepseekmoe][wang2024auxiliarylossfree].

Key equations

\vp = \softmax(\vr), \qquad S = \operatorname{topk}(\vp)

a_i = \frac{p_i}{\sum_{j\in S}p_j}, \quad \vy = \sum_{i\in S} a_i E_i(\vx)

\mathcal{L}_{\text{lb}} = N\sum_i f_iP_i, \qquad N\sum_i p_i^2 \ge 1

\mathcal{L}_z = \E\Big[\big(\log\sum_i e^{r_i}\big)^2\Big]

b_i \leftarrow b_i - \eta\,\sign(f_i - 1/N)

27.5 Teach it

The one-sentence version. An MoE layer is a bank of FFNs plus a router that chooses a few of them per token and averages their outputs with renormalized router probabilities.

An analogy. A dense FFN is one generalist doctor for every patient; MoE is a clinic that sends each patient to a few specialists, while watching that no specialist’s queue explodes.

At the board.

Write one dense FFN, then replace it by $E_1,\ldots,E_N$ .
Compute $\vp = \softmax(\vr)$ , circle the top-k entries, and renormalize them.
Show $\vy = \sum a_iE_i(\vx)$ and then write the Switch loss.
Use $N\sum_i p_i^2 \ge 1$ to explain why uniform routing is the balanced point.

Misconceptions to address. Sparse does not mean cheap communication; expert parallelism moves token batches between devices. The router bias in auxiliary-loss-free balancing is not a model logit; it is only a selection nudge. Dropped tokens are a capacity mechanism, not a goal.

Check for understanding. If every token picks expert 0, what happens to $\mathcal{L}_{\text{lb}}$ and why does capacity matter?

27.6 Exercises

Exercise 27.1 ★ Router arithmetic

Given router probabilities $(0.5, 0.3, 0.2)$ and top-k set $\{1,2\}$ , compute the combine weights and explain why they sum to one.

Solution in Appendix D

Exercise 27.2 ★★ Uniform minimizes the Switch fixed point

Assume $f_i = P_i = p_i$ . Prove that $N\sum_i p_i^2 \ge 1$ and identify the equality case. Then compute the loss for a collapsed router.

Solution in Appendix D

Exercise 27.3 ★★ Capacity and dropping

For top-1 routing with two experts and capacity two, four tokens choose experts $(0,0,0,1)$ in order. Which assignment is dropped? What output contribution does it make?

Solution in Appendix D

Exercise 27.4 ★★★ Implement and compare

Write a dense token-by-token MoE reference and compare it with sparse_moe_forward on random small tensors. Why is this a stronger test than checking shapes?

Solution in Appendix D

References

[dai2024deepseekmoe] D. Dai et al. DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models. 2024. arXiv:2401.06066
[fedus2021switch] W. Fedus, B. Zoph, and N. Shazeer. Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity. 2021. arXiv:2101.03961
[lepikhin2020gshard] D. Lepikhin et al. GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding. 2020. arXiv:2006.16668
[shao2024deepseekv2] Z. Shao et al. DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model. 2024. arXiv:2405.04434
[shazeer2017outrageously] N. Shazeer et al. Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer. 2017. arXiv:1701.06538
[wang2024auxiliarylossfree] L. Wang et al. Auxiliary-Loss-Free Load Balancing Strategy for Mixture-of-Experts. 2024. arXiv:2408.15664
[zoph2022stmoe] B. Zoph et al. ST-MoE: Designing Stable and Transferable Sparse Expert Models. 2022. arXiv:2202.08906