Optimal Splitting of Language Models from Mixtures to Specialized Domains

Abstract

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.

Figures (18)

• Figure 1: (a) Split model training: The pretraining corpus is clustered based on semantics, then a model is trained on all domains before being copied and subsequently trained independently on an individual cluster. (b) Compute multiplier: The multiplicative factor on the amount of data required for standard pretraining to match optimal split model training performance. At large compute budgets, standard pretraining requires up to 50$\times$ the amount of data to match the loss of optimal split training.
• Figure 2: Average fact accuracy (over two clusters) vs. split step for synthetic phonebook fact clusters with different overlaps under different total training budget. We say the total budget is sufficient if split training on two cluster each using (Total Steps/2) from step 0 reaches 100% average fact accuracy, and otherwise say the total budget is limited.
• Figure 3: Predicted loss vs. observed loss across multiple scenarios. The scaling laws are fit on smaller models with varying amount of PT and CPT token budgets. Loss is estimated over domain 8.
• Figure 4: Left Number of pretraining tokens after which split training becomes beneficial $D_{\mathrm{split}}$, as a function of model size. We compute it as the solution of equation \ref{['eq:split_id']}. We see that it takes more tokens for larger models to start benefiting from split training. We only use the loss on one domain to approximate the average loss $L$. Right Optimal fraction of pretraining tokens as a function of the total training budget $t_S/ D_T$: the fraction decays with budget, but the optimal number of pretraining tokens increases with budget.
• Figure 5: Loss vs. fraction of pretraining tokens for a 1.3B parameter models at 120, 240, and 360B total training tokens.