Mixture-of-Experts Routing - LLM Internals

Mixture-of-Experts (MoE) Routing

What you'll build

A top-k mixture-of-experts layer: a router picks k experts per token, each expert runs only on the tokens routed to it, the outputs are combined with the gate weights, and a load-balancing auxiliary loss is returned.

Background

MoE layers (Switch Transformer, Mixtral, DeepSeek-MoE) replace one big FFN with E smaller experts and a router. Parameters grow with E, but compute per token grows only with k: Mixtral 8x7B has about 47B parameters but uses about 13B per token. The cost is engineering: tokens must be dispatched to experts and gathered back, and the router must be kept from collapsing onto a few experts. Here each expert is a linear map expert_e(x) = x @ expert_weights[e] (real experts are SwiGLU FFNs; the routing is identical).

Algorithm

With T tokens and E experts:

1. Router: logits = x @ W_router (T, E); probs = softmax(logits, axis=-1) (used by the aux loss).
2. Top-k gating (Mixtral style): expert_indices = the k experts with the largest logits per token, sorted by descending logit; gates = softmax(top-k logits), renormalized over the chosen k only.
3. Dispatch and combine: for each expert e, gather the tokens routed to it, run the expert only on those tokens, and add the outputs back weighted by their gate:

y_t = sum_{j=1..k} gates[t, j] * expert_{expert_indices[t, j]}(x_t)

4. Load-balancing loss (Switch Transformer):

f_e      = (number of tokens whose top-k includes e) / (T * k)    # fraction of routing slots
P_e      = mean over tokens of probs[:, e]                        # mean router probability
aux_loss = E * sum_e f_e * P_e

It is 1.0 for perfectly uniform routing and approaches E when everything collapses onto one expert. In training it is added to the LM loss times a small coefficient (e.g. 0.01); f is not differentiable, so the gradient flows through P.

Inputs

• x: (T, d)
• W_router: (d, E)
• expert_weights: (E, d, d_out)
• top_k: int, 1 <= top_k <= E

Expected Return Format

A dictionary with these keys:

• 'output': combined expert outputs, (T, d_out)
• 'expert_indices': integer array (T, k), sorted by descending router logit
• 'gates': (T, k), each row sums to 1, aligned with 'expert_indices'
• 'aux_loss': the load-balancing loss as a Python float

Edge cases and conventions

• top_k = 1 (Switch): the gate is 1 and the output is the chosen expert's output.
• Tests use random logits, so ties between experts don't occur; a stable sort of -logits is a safe choice.

Constraints

NumPy only. Don't run every expert on every token and mask afterwards: run each expert on its routed tokens.

Follow-up

• Real systems give each expert a capacity (capacity_factor · T · k / E slots) and drop overflow tokens. How would you implement that, and what happens to dropped tokens?
• Why does expert parallelism need an all-to-all communication step?

Input: x (6, 4), W_router (4, 4), expert_weights (4, 4, 3), top_k=2
Output: {'output': shape (6, 3), 'expert_indices': shape (6, 2), 'gates': shape (6, 2), 'aux_loss': float}

Input: W_router = np.zeros((4, 4)), any x, top_k=2
Output: aux_loss = 1.0

Input: x = np.ones((6, 4)), router strongly prefers expert 0, top_k=1, E=4
Output: aux_loss = 4.0

import numpy as np

def moe_layer(x, W_router, expert_weights, top_k):
    """
    Top-k mixture-of-experts layer with a load-balancing loss.

    Args:
        x: (T, d) token representations
        W_router: (d, E) router weights
        expert_weights: (E, d, d_out); expert e computes x @ expert_weights[e]
        top_k: experts per token

    Returns:
        dict with 'output' (T, d_out), 'expert_indices' (T, k) sorted by
        descending logit, 'gates' (T, k) rows summing to 1, and
        'aux_loss' (float) = E * sum_e f_e * P_e
    """
    # Your code here
    pass