  We introduce a new balanced assignment of experts (BASE) layer for large
language models that greatly simplifies existing high capacity sparse layers.
Sparse layers can dramatically improve the efficiency of training and inference
by routing each token to specialized expert modules that contain only a small
fraction of the model parameters. However, it can be difficult to learn
balanced routing functions that make full use of the available experts;
existing approaches typically use routing heuristics or auxiliary
expert-balancing loss functions. In contrast, we formulate token-to-expert
allocation as a linear assignment problem, allowing an optimal assignment in
which each expert receives an equal number of tokens. This optimal assignment
scheme improves efficiency by guaranteeing balanced compute loads, and also
simplifies training by not requiring any new hyperparameters or auxiliary
losses. Code is publicly released at https://github.com/pytorch/fairseq/
