Hard Research

Mixture of Experts Routing

A mixture-of-experts layer chooses a few expert functions for an input, then combines their outputs using routing weights. Implement moe_forward for one token and simple linear experts.

Inputs and output

x has shape (d,): one token vector, with d positive.
gate_weights has shape (d, E), where E is the number of experts.
expert_weights is a list of E matrices, each with shape (d, d). List index i identifies expert i.
top_k is an integer with 1 <= top_k <= E.

Inputs are floating-point tensors with a shared dtype and device. Assume the computed gate logits are finite. Return one vector of shape (d,), preserving dtype and device and leaving all inputs unchanged.

Route, evaluate and combine

Compute logits = x @ gate_weights, giving one score per expert.
Select the top_k largest scores. Break ties by choosing lower expert indices first.
Apply a stable softmax over those selected scores only. The selected weights sum to one; unselected experts contribute nothing.
For each selected expert i, compute expert_output = x @ expert_weights[i]. x is a row vector here. Do not transpose the expert matrix.
Add each expert output multiplied by its routing weight.

For example, if two selected experts have equal scores, each gets weight 0.5. If their outputs are [2, -1] and [4, 3], return [3, 1].

With top_k=1, the selected expert has routing weight one, regardless of its score. With top_k=E, combine all experts. Do not return the routing weights or concatenate expert outputs.

This exercise uses deterministic routing and linear experts, with no bias, activation, routing noise or auxiliary loss. It is a simplified version of the sparse expert combination described by Shazeer et al. in The Sparsely-Gated Mixture-of-Experts Layer.

Understand the ideas behind this exercise. The Problem tab contains the exact implementation requirements.

Choose which experts process each token

A mixture-of-experts layer contains several networks, called experts, and a router. For each token, the router scores the experts, selects some of them, and weights their outputs. Sparse routing lets a large model use only a few experts for each token.

For scores [3, 1, 0], top-1 routing selects expert 0. If selected weights are normalized to sum to one, its weight is 1. Top-2 routing selects experts 0 and 1, with weights determined by the task's rule.

Follow the weighting convention

Two common choices are softmax over selected logits and full softmax followed by selection. Full softmax followed by selection gives the same selected weights as the first choice if they are renormalized. Without renormalization, their sum is generally less than one and the output differs.

Setting excluded logits to zero is not a valid way to remove them before softmax: exp(0) = 1, so they would still get positive weight. Excluded experts must contribute zero probability to the combination.

Put each expert result back at the right token

When grouping tokens for expert computation, keep enough indexing information to restore their original order. If several experts process the same token, add their weighted outputs; do not overwrite one contribution with another.

Real systems also limit expert capacity and encourage balanced workloads so one expert does not receive nearly all tokens. Implement those policies only when the exercise includes them; they can change which outputs are returned.

Detecting runtime…

Mixture of Experts Routing

Inputs and output

Route, evaluate and combine

Choose which experts process each token

Follow the weighting convention

Put each expert result back at the right token

Your code, saved here

Stuck?

PyTorch reference solution