Mixture of Experts: top-k gating, route each token to its best expert FFN

The problem, from first principles

[01-transformer-block/04-feedforward-sublayer]'s FFN applies the exact SAME weights to every single token, regardless of what that token actually is: every token pays the full compute cost of the entire FFN, whether or not it needed all of that capacity. Mixture of Experts (MoE) architectures ask a pointed question: what if, instead of one large FFN every token must pass through, a model had MANY smaller "expert" FFNs, and each token only needed to pass through a FEW of them, chosen dynamically based on the token's own content? A learned "gating" network scores every expert for every token, and only the top-scoring experts actually run for that token, letting the model have a HUGE total parameter count (many experts) while keeping the actual COMPUTE per token comparable to a much smaller dense model, since most experts are simply never invoked for any given token.

This is a genuinely different kind of scaling lever than everything earlier in this curriculum: [01-transformer-block]'s FFN scales compute and parameters TOGETHER (a bigger d_ff costs more compute for every token), while MoE scales PARAMETERS (more experts) largely independently of compute (still only top_k experts run per token, regardless of how many total experts exist).

From theory to code

Implement moe_gate(x, gate_weight, top_k), scoring every expert and keeping only the top_k highest-scoring ones (softmax-renormalized among just those top_k), and moe_ffn_forward(x, gate_weight, expert_params, top_k), routing each token through its own selected experts and combining their outputs via the gate's weights.

Constraints

• moe_gate scores ALL experts (x @ gate_weight.T), then keeps only the top_k highest scores per token, discarding the rest ENTIRELY (not merely down-weighting them).
• The kept top_k scores are re-normalized via softmax AMONG THEMSELVES, so they still sum to 1, even though the full expert population was much larger.
• moe_ffn_forward runs each token's top_k selected experts' OWN independent FFN (via [01-transformer-block/04-feedforward-sublayer]'s feedforward_sublayer, one call per selected expert), and combines their outputs as a WEIGHTED SUM using the gate weights.
• A non-selected expert's weights must have NO effect on a token's output; only the top_k experts actually chosen for that specific token contribute anything.

Hints

gate_logits = x @ gate_weight.T          # (..., num_experts)
top_indices = np.argsort(gate_logits, axis=-1)[..., -top_k:]
top_logits = np.take_along_axis(gate_logits, top_indices, axis=-1)
top_weights = softmax_last_axis(top_logits)