Inside the Mixture-of-Experts layer: softmax gating, top-k routing, capacity, and the load-balancing loss

A dense Transformer runs every token through the same feed-forward block, so capacity (parameters) and cost (FLOPs) are welded together — to get a smarter model you pay more per token. A Mixture-of-Experts layer breaks that weld. It replaces the one shared FFN with N independent expert MLPs plus a tiny gating network, and keeps only a few experts per token. Params scale with N; FLOPs per token stay fixed. More knowledge, roughly the same compute. That decoupling is how Switch Transformer and Mixtral-8x7B (≈47B params, only ≈13B active per token) work.

The gate — a softmax over the experts

The gate is a single linear layer W_g that maps a token to one logit per expert; softmax turns those into a distribution — how much this token wants each expert.

def softmax(z):
    z = z - z.max()                 # stability, doesn't change the result
    e = np.exp(z)
    return e / e.sum()

def gate(x, Wg):                    # Wg: (N_experts, d)
    return softmax(Wg @ x)          # p[e] = P(expert e | token x)

Top-k routing — the sparsity

Sparsity is the whole point: do not run all N experts, keep only the top-k gate entries (k=1 in Switch, k=2 in Mixtral). Renormalize just those k weights so they sum to 1 again — they become the mixing coefficients for combining the chosen experts' outputs, y = Σ_{e∈top-k} g̃_e · E_e(x).

def top_k(p, k):
    idx = np.argsort(p)[::-1][:k]   # indices of the k largest gate probs
    w   = p[idx]
    w   = w / w.sum()               # renormalize over the kept experts
    return idx, w                   # experts to run + their weights

Because k is fixed, the FLOPs per token stay constant no matter how many experts you add. Grow N and you add parameters, not compute.

Capacity — a hardware budget that drops overflow

Real hardware wants dense, equal-sized batches per expert, so each expert gets a fixed number of slots C = f·(T·k / N). Dispatch tokens in arrival order; once an expert is full, extra tokens overflow and are dropped — they skip the layer through the residual connection. The capacity factor f buys slack against imbalance; shrink it toward 0.25 and even balanced routing starts dropping tokens. Capacity is a budget, not a correctness knob.

def dispatch(idxs, T, k, N, f=1.25):
    C = int(np.ceil(f * T * k / N))     # slots per expert
    load = np.zeros(N, int); kept = []
    for token_experts in idxs:          # tokens in arrival order
        keep = []
        for e in token_experts:
            if load[e] < C:
                load[e] += 1; keep.append(e)   # accepted
            # else: OVERFLOW -> dropped, token uses the residual
        kept.append(keep)
    return kept, load, C

The auxiliary load-balancing loss

A raw gate collapses: whichever experts get picked early get trained more, get picked more — rich-get-richer, until a few experts dominate and the rest die. Switch adds an auxiliary loss to the training objective:

$$L_{aux} = N \cdot \sum_e f_e \cdot P_e$$

where f_e is the fraction of tokens dispatched to expert e (a hard count) and P_e is the mean soft router probability for e. It is minimized at 1 exactly when usage is uniform, and climbs toward N under total collapse — so pulling it down is load balancing.

def aux_loss(P_soft, dispatch_counts, T, k, N):
    f = dispatch_counts / (T * k)   # fraction of routed tokens per expert
    P = P_soft.mean(axis=0)         # mean router prob per expert (soft)
    return N * np.sum(f * P)        # = 1 when uniform, ~N when collapsed

The gradient flows through the soft P_e (the hard count f_e is detached), and through the softmax Jacobian it reduces to a clean rule: push the logit of an over-used expert down. Take a step and the aux loss falls, the load bars level off, and dropped tokens go to zero.

Put it together — gate every token, take top-k, dispatch under capacity, run the kept experts, gate-weight the outputs, add λ·L_aux during training — and you have the layer. Total params scale with N, FLOPs per token scale with k.

Watch a collapsed router genuinely rebalance by gradient descent, live, at https://dev48v.infy.uk/dl/day56-mixture-of-experts.html