31. Mixture-of-Experts: Top-2 Routing with Load-Balancing Auxiliary Loss

Problem

Mixtral 8x7B replaces the single feed-forward block in each transformer layer with 8 independent "expert" feed-forward networks and a small router: for every token, the router scores all 8 experts, picks the top-2, and combines their outputs weighted by the (renormalized) router probabilities. This gives the model far more total parameters than it pays compute for per token, since only 2 of 8 experts actually run. Naive top-k routing has a well-known failure mode though: without a corrective signal, the router can collapse onto a small favorite subset of experts, leaving the rest undertrained. The load-balancing auxiliary loss fixes this by directly penalizing uneven routing, added to the main training loss during pretraining.

Given a batch of token vectors, a router weight, and a stack of per-expert weight matrices, compute: (1) the combined output using top-2 routing with renormalized gate weights, and (2) the load-balancing auxiliary loss num_experts * sum_e(f_e * P_e), where f_e is the fraction of tokens that route to expert e (counting both top-2 slots) and P_e is the average full-softmax router probability assigned to expert e across the batch.

Source: src/31_moe_top2_routing_load_balancing.py

def moe_top2_forward(
    x: torch.Tensor, gate_weight: torch.Tensor, expert_weights: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]: ...

>>> x = torch.randn(6, 4)
>>> gate_weight = torch.randn(4, 3)
>>> expert_weights = torch.randn(3, 4, 4)
>>> output, aux_loss = moe_top2_forward(x, gate_weight, expert_weights)
>>> output.shape
torch.Size([6, 4])

Step-by-Step Approach

1. Compute router logits x @ gate_weight and the full softmax distribution probs, shape (num_tokens, num_experts) — this full distribution is needed for the aux loss even though only 2 experts get used per token.
2. probs.topk(2, dim=-1) gives each token's top-2 expert indices and their (un-normalized) probabilities in one call.
3. Renormalize the top-2 weights so they sum to 1: top2_vals / top2_vals.sum(dim=-1, keepdim=True) — without this, tokens the router is more "confident" about would get systematically larger-magnitude outputs.
4. Gather each token's two chosen expert weight matrices via fancy indexing, expert_weights[top2_idx[:, 0]], then apply with torch.bmm(x.unsqueeze(1), w0).squeeze(1) — a per-token batched matmul, not a loop over experts.
5. Combine: output = weight0 * y0 + weight1 * y1.
6. For the aux loss: P_e = probs.mean(dim=0) (mean full-softmax probability per expert); build a one-hot (T, E) mask of each token's top-2 picks with scatter_, then f_e = onehot.mean(dim=0) (fraction of tokens routed to each expert); finally aux_loss = num_experts * (f_e * P_e).sum().

The key insight behind the loss formula is that it multiplies a hard routing signal (f_e, which expert actually got picked) by a soft one (P_e, what the router thinks of every expert including ones it didn't pick) — this makes the loss differentiable through the router even though the top-2 selection itself is not, since gradients flow through P_e.

Reference solution

import torch


def moe_top2_forward(
    x: torch.Tensor, gate_weight: torch.Tensor, expert_weights: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    num_tokens, d_model = x.shape
    num_experts = gate_weight.shape[1]

    logits = x @ gate_weight                            # (T, E) router logits
    probs = torch.softmax(logits, dim=-1)                # (T, E) full router distribution
    top2_vals, top2_idx = probs.topk(2, dim=-1)           # (T, 2) each
    # renormalize so the two chosen weights sum to 1 -- otherwise the output
    # magnitude would depend on how "confident" the router happened to be
    top2_weights = top2_vals / top2_vals.sum(dim=-1, keepdim=True)

    # gather each token's two chosen expert matrices directly via fancy
    # indexing, then apply with a batched matmul -- no loop over experts
    w0 = expert_weights[top2_idx[:, 0]]                   # (T, d_model, d_model)
    w1 = expert_weights[top2_idx[:, 1]]
    y0 = torch.bmm(x.unsqueeze(1), w0).squeeze(1)          # (T, d_model)
    y1 = torch.bmm(x.unsqueeze(1), w1).squeeze(1)
    output = top2_weights[:, 0:1] * y0 + top2_weights[:, 1:2] * y1

    # Mixtral-style load-balancing loss: num_experts * sum_e(f_e * P_e).
    # P_e uses the FULL softmax (not just top-2) so the loss sees the
    # router's raw preference even for experts it didn't select.
    p_e = probs.mean(dim=0)                                # (E,)
    onehot = torch.zeros_like(probs)
    onehot.scatter_(1, top2_idx, 1.0)                      # (T, E) 1 at each token's 2 chosen experts
    f_e = onehot.mean(dim=0)                                # (E,) fraction of tokens routing to e
    aux_loss = num_experts * (f_e * p_e).sum()

    return output, aux_loss

Key Functions & Tricks

• probs.topk(2, dim=-1) — returns both the top-2 values and their indices in one call, exactly what's needed for both the output combine and the aux-loss one-hot mask.
• expert_weights[top2_idx[:, 0]] — fancy/advanced indexing gathers a per-token weight matrix in one vectorized op, avoiding a python loop over tokens or experts.
• torch.bmm(x.unsqueeze(1), w0).squeeze(1) — applies a different weight matrix to each token in the batch simultaneously; ordinary x @ W can't do this since W would need to be shared across the batch.
• onehot.scatter_(1, top2_idx, 1.0) — writes a 1 at each token's two selected expert columns without a python loop; the in-place scatter_ (not gather) is the direction that matches "index tells you where to write."
• probs.mean(dim=0) / onehot.mean(dim=0) — both aggregate over the token axis to produce one scalar per expert, the P_e and f_e terms of the loss formula respectively.

How to Recognize This Pattern