MoE: Top-k Expert Routing

Easy · Solved in PyTorch · Machine learning coding practice

Problem

Implement deterministic top-k gating, the noise-free form of the sparse gate in The Sparsely-Gated Mixture-of-Experts Layer (Shazeer et al., 2017).

logits has shape (T, E) with finite floating-point scores, T,E \geq 1, and 1 \leq k \leq E. For each token, select the k largest scores in descending order; break ties by smaller expert ID. Apply a numerically stable softmax over those selected scores only.

Topics: mixture-of-experts, transformers, routing, pytorch, numpy

The full statement, worked examples, hints and the test suite are available once you sign in. You can then solve MoE: Top-k Expert Routing in the browser and run it against the tests.

Related problems

• MoE: Dispatch and Weighted Combine Medium
• MoE: Router Load-Balancing Loss Medium
• Bradley Terry Ranking Medium
• Looped Transformer: Adaptive Early Exit Medium
• Looped Transformer: Shared-Weight Unrolling Easy
• Accumulate and Average Replica Gradients Easy