Scaling Expert Language Models with Unsupervised Domain Discovery
This addresses the scalability issue for researchers and practitioners training large language models, offering an incremental improvement over existing sparse methods by automating domain discovery.
The paper tackles the problem of high communication overhead in training large language models by introducing a method that clusters a corpus into domains, trains separate expert models on each cluster, and combines them in a sparse ensemble for inference, outperforming dense baselines on multiple corpora and few-shot tasks.
Large language models are typically trained densely: all parameters are updated with respect to all inputs. This requires synchronization of billions of parameters across thousands of GPUs. We introduce a simple but effective method to asynchronously train large, sparse language models on arbitrary text corpora. Our method clusters a corpus into sets of related documents, trains a separate expert language model on each cluster, and combines them in a sparse ensemble for inference. This approach generalizes embarrassingly parallel training by automatically discovering the domains for each expert, and eliminates nearly all the communication overhead of existing sparse language models. Our technique outperforms dense baselines on multiple corpora and few-shot tasks, and our analysis shows that specializing experts to meaningful clusters is key to these gains. Performance also improves with the number of experts and size of training data, suggesting this is a highly efficient and accessible approach to training large language models.