BASE Layers: Simplifying Training of Large, Sparse Models
This work simplifies training for large, sparse models, which is an incremental improvement for researchers and practitioners in efficient AI model development.
The paper tackles the challenge of learning balanced routing functions in large, sparse language models by introducing a BASE layer that formulates token-to-expert allocation as a linear assignment problem, resulting in optimal balanced assignments without needing hyperparameters or auxiliary losses.
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/