Switch Transformers: The Load-Balancing Loss
Advanced
The problem, from first principles
Left alone, a learned router tends to send most tokens to a few favourite experts, leaving the rest unused. Switch Transformers add an auxiliary loss that penalizes uneven routing, so that the experts share the work.
From theory to code
Implement load_balancing_loss(expert_index, probs, n_experts), returning E * sum_i f_i P_i.
Constraints
f_icounts the tokens assigned to expertidivided byT.
Hints
Hint 1
Use np.bincount for the token fractions, take the mean of probs over tokens, then combine.
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.