MoE: Router Load-Balancing Loss
Medium · Solved in PyTorch · Machine learning coding practice
Problem
Implement a top-k extension of the auxiliary loss in Switch Transformers (Fedus et al., 2021). This exercise normalizes assignment counts over all selected slots, so k=1 recovers the Switch formulation.
logits contains finite scores of shape (T,E); expert_indices contains fixed, distinct zero-based expert IDs per row of shape (T,k); and boolean token_mask of shape (T,) marks real tokens. Dimensions are positive, 1 \leq k \leq E, and at least one token is valid. Use the supplied assignments without recomputing them. Let V be the valid-token set, M=|V|, and p_{te} the stable softmax of logits over all E experts.
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: Router Load-Balancing Loss in the browser and run it against the tests.