Switch Transformers: Top-1 Routing
Intermediate
The problem, from first principles
Switch Transformers (Fedus, Zoph & Shazeer, 2021) send each token to just one expert, chosen by a router. Routing to a single expert keeps compute per token constant, even as the total number of parameters grows with the number of experts.
From theory to code
Implement switch_route(logits), which returns each token's top expert and that expert's softmax probability.
Constraints
- Use a stable softmax over the expert axis.
Hints
Hint 1
Softmax each row, take the argmax as the expert, then read that expert's probability as the gate.
Keywords
research-paperstransformersllmmoerouting
Python
Run your code to see results
No results yet
Run your code to check it against the examples, or submit to run every test.