Optimal Splitting of Language Models from Mixtures to Specialized Domains

Language models achieve impressive performance on a variety of knowledge, language, and reasoning tasks due to the scale and diversity of pretraining data available. The standard training recipe is a two-stage paradigm: pretraining first on the full corpus of data followed by specialization on a subset of high quality, specialized data from the full corpus. In the multi-domain setting, this involves continued pretraining of multiple models on each specialized domain, referred to as split model training. We propose a method for pretraining multiple models independently over a general pretraining corpus, and determining the optimal compute allocation between pretraining and continued pretraining using scaling laws. Our approach accurately predicts the loss of a model of size $N$ with $D$ pretraining and $D'$ specialization tokens, and extrapolates to larger model sizes and number of tokens. Applied to language model training, our approach improves performance consistently across common sense knowledge and reasoning benchmarks across different model sizes and compute budgets.

Training language models (LMs) to master multiple specialized subjects including math, code, or science is highly inefficient when forcing a single model to learn everything simultaneously, and it is difficult to know exactly how much compute to spend on general knowledge versus specialized training. This work discusses "split model training" an approach that trains multiple copies of a model on different specialized knowledge domains and routes to different domains depending on a user query. To understand when to specialize, we propose a mathematical scaling law, which calculates the exact, optimal moment to stop general training on all data, and split to multiple independent models that each specialize in a specific domain. This optimized splitting strategy consistently boosts performance on reasoning and knowledge tests across multiple topics including math, science, and general knowledge compared to training on all data jointly. Since the independent expert models are trained in parallel on standard hardware, this further has the potential to enable broader access to training LLMs and help reduce the cost of training.