Optimal Splitting of Language Models from Mixtures to Specialized Domains

Skyler Seto; Pierre Ablin; Anastasiia Filippova; Jiayuan Ye; Louis Bethune; Angelos Katharopoulos; David Grangier

cs.CLMar 19, 2026

Optimal Splitting of Language Models from Mixtures to Specialized Domains

Authors: Skyler Seto, Pierre Ablin, Anastasiia Filippova, Jiayuan Ye, Louis Bethune, Angelos Katharopoulos, David Grangier

Organizations: Apple · National University of Singapore

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 & tables

Panel of Figure 1 — 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.

Panel of Figure 2 — 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.

	Train Size	Test Size	MAE	R2
Scenario 1	652	144	0.012	0.934
Scenario 2	347	449	0.021	0.982
Scenario 3	156	640	0.023	0.979

Table 1: MAE and

R^{2}

for different model size on varying held out test conditions.

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 5.2 . 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.

Panel of Figure 5 — Figure 5: Loss vs. fraction of pretraining tokens for a 1.3B parameter models at 120, 240, and 360B total training tokens.

Panel of Figure 6 — Figure 6: Zero shot QA accuracy vs. fraction of pretraining tokens for a 1.3B parameter models at 120, 240, and 360B total training tokens.

Panel of Figure 7 — Figure 7: Performance comparison across model sizes for split training. (a) Zero-shot QA accuracy for standalone (base) and split (cluster) models at various model sizes. Dashed lines indicate where the best performing model changes in size as a function of the total compute (FlOPs). (b) Zero-shot QA accuracy for various standalone and split training models at 1.3B and 2.7B scale. All models are trained for a total compute budget equivalent to roughly 140B tokens of the 6.7B model, however 1.3B and 2.7B models have lower inference costs.

	Subset Routing		Full
PT Tokens	MMLU Subset	Task Subset	Full
10B	53.85	55.85	56.94
20B	55.68	57.05	58.11
40B	57.47	59.08	59.53
80B	59.41	60.70	61.26
120B	59.62	60.97	61.53

Table 2: Routing comparisons for split training of a 1.3B model at 150B tokens. Subset routing uses only two models per task at inference time whereas full routing uses all domain models.

Figure 8: Accuracy vs. number of training tokens for varying number of domains and varying number of PT steps. All domain models are 1.3B parameter models. The total number of parameters varies with the number of domains, but the number of active parameters remains constant at a single domain model.

Appendix figures & tables18 assets

Supplementary material from the paper’s appendix.

Appendix

Task	Routing Text
ARC	question
SciQ	question
HellaSwag	ctx
PIQA	goal
BoolQ	question
Winogrande	sentence

Appendix

Table 3: Routing Text Field by Task

Panel of Figure 9 — Figure 9: Distribution of samples from Pile domains within the DCLM clustering.

Figure 10: Loss for five domains of the pile and common crawl data for each split model. The models are 1.3B models pretrained for 340B tokens before CPT.

Domain	Pretrained	Split Model
ArXiv	11.01	9.92
DM Math	8.14	7.81
FreeLaw	5.76	4.85
Github	3.74	3.42
PubMed Central	7.90	7.25

Appendix

Table 4: PPL for different domains of the Pile for a standalone pretrainged model only, and split models. The standalone model with pretraining only is trained for a total of 720B tokens. The split models are 1.3B parameters and trained for the same total compute budget with

t_{s}=340

Prefix Length (tokens)	R@1	R@5
1	0.0128	0.0362
2	0.0955	0.1753
4	0.2664	0.3765
8	0.4802	0.5884
16	0.7112	0.8078
32	0.8775	0.9330

Appendix

Table 5: Retrieval accuracy by prefix length (N=10000)

Figure 11: Average loss vs. fraction of compute for pretraining for different pairs of clusters in the DCLM-base dataset. Total compute budget is 90000 steps or roughly 90B tokens.

Parameter	Transform	Effective range
$E_{p}$	softplus	$(0,\infty)$
$E_{0}$	scaled sigmoid	$[1.0,\;3.0]$
$A$	softplus	$(0,\infty)$
$\alpha_{1}$	sigmoid	$[0.1,\;1.0]$
$\alpha_{2}$	sigmoid	$[0.1,\;1.0]$
$c$	sigmoid	$[0.5,\;4.0]$

Appendix

Table 6: Parameter constraints for the single-size CPT scaling law. Transformations are applied to unconstrained optimizer variables to enforce bounded, physically meaningful parameter values.

Figure 12: Estimated optimal fraction of pretraining tokens as a function of total compute budget, estimated from per-size scaling laws (solid lines) and extrapolation via power-law fit across smaller sizes (dotted lines). Shaded intervals represent the range of allocations that achieve loss within 0.005 of the optimum, capturing the flatness of the loss surface near the optimal split.

Figure 13: Estimated fits with the scaling law from ( Liew and Kato, 2025 )

	Train Size	Test Size	R2 (ours)	R2 ( Liew and Kato, 2025 )
Scenario 1	649	142	0.934	0.910
Scenario 2	344	447	0.982	0.946
Scenario 3	156	640	0.979	0.953

Appendix

Table 7: Comparison of the test R2 scores with our proposed scaling law and that of ( Liew and Kato, 2025 ) .

Figure 14: 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 0.

	Train Size	Test Size	MAE	R2
Scenario 1	649	142	0.009	0.934
Scenario 2	344	447	0.025	0.960
Scenario 3	156	640	0.022	0.973

Appendix

Table 8: Fitting the scaling laws on domain 0 with multiple held out test conditions.

			Train			Test
Domain	$n_{\text{train}}$	$n_{\text{test}}$	MAE	MRE	$r$	MAE	MRE	$r$
0	113	60	0.0021	0.0008	0.9993	0.0014	0.0005	0.9959
1	114	60	0.0023	0.0008	0.9945	0.0024	0.0009	0.9988
2	113	60	0.0020	0.0008	0.9995	0.0022	0.0009	0.9992
3	113	60	0.0026	0.0010	0.9989	0.0032	0.0013	0.9965
4	112	59	0.0020	0.0007	0.9993	0.0014	0.0005	0.9995

Appendix

Table 9: Per-domain scaling law fit metrics for the 1.3B model. Scaling laws are trained with the CPT steps

<20

B tokens, and tested with CPT steps

\geq 20

B tokens.

Panel of Figure 15 — Figure 15: Performance across eight zero-shot QA benchmarks for a 125M model at the same training compute with varying amount of pretraining steps. Each subfigure has a different total number of training tokens.

Panel of Figure 16 — Figure 16: Performance of 1.3B models at the same training compute with different amount of pretraining steps. We compare split training with pretraining from scratch, split training with no pretraining, and finetuning (split training for only 1B tokens).

Figure 17: Accuracy vs. number of training tokens for varying number of domains and varying number of PT steps. All domain models are 2.7B parameter models. The total number of parameters varies with the number of domains, but the number of active parameters remains constant at a single domain model.

Benchmark	Pretrain 1.3B	Split 1.3B	Pretrain 2.7B	Split 2.7B	Pretrain 7B
ARC-Challenge	39.25	42.75	40.70	45.82	42.06
ARC-Easy	68.06	72.31	70.62	73.74	72.94
BoolQ	64.01	64.19	65.78	66.21	63.12
HellaSwag	63.91	67.30	67.60	69.37	69.39
MMLU	35.59	37.04	37.42	38.54	38.74
PIQA	75.35	78.29	78.24	78.29	77.75

Appendix

Table 10: QA benchmark accuracy for split model training with optimal token budget allocation.

Cluster ID	Base PPL	Split MoE PPL
0	20.190	19.819
1	18.075	17.731
2	18.600	18.354
3	16.980	16.903
4	21.036	18.414
5	19.639	19.550

Appendix

Table 11: Per-cluster perplexity (PPL) comparing split training of Mixture of Expert (MoE) models with a single model pretrained on all clusters.

Efficient Language Model Training

Efficient language model training studies how to reduce the data, compute, memory, time, energy, and cost needed to build capable language models.

See all 162 papers

Scaling Laws for Mixture Pretraining Under Data Constraints
May 12, 2026Anastasiia Sedova, Skyler Seto, Natalie Schluter +1Language Model Pretraining Language Model Scaling Laws
Apple · ITU
The Finetuner's Fallacy: When to Pretrain with Your Finetuning Data
Mar 17, 2026Christina Baek, Ricardo Pio Monti, David Schwab +31Language Model Pretraining Domain-Adaptive Pretraining
DatologyAI
Train Separately, Merge Together: Modular Post-Training with Mixture-of-Experts
Apr 20, 2026Jacob Morrison, Sanjay Adhikesaven, Akshita Bhagia +3Continual Learning for LLMs Mixture-of-Experts Language Models
University of Washington · Allen Institute for AI · University of California, Berkeley

May 12, 2026J/K move · Enter open · S save

Scaling Laws for Mixture Pretraining Under Data Constraints

Anastasiia Sedova, Skyler Seto, Natalie Schluter, Pierre Ablin

Apple · ITU

Language Model Pretraining Language Model Scaling Laws Data Mixture Optimization LLM Training Data Repetition in Language Model Training

Open paper AskPDF arXiv

As language models scale, the amount of data they require grows -- yet many target data sources, such as low-resource languages or specialized domains, are inherently limited in size. A common strategy is to mix this scarce but valuable target data with abundant generic data, which presents a fundamental trade-off: too little target data in the mixture underexposes the model to the target domain, while too much target data repeats the same examples excessively, yielding diminishing returns and eventual overfitting. We study this trade-off across more than 2,000 language-model training runs spanning multiple model and target dataset sizes, as well as several data types, including multilingual, domain-specific, and quality-filtered mixtures. Across all settings, we find that repetition is a central driver of target-domain performance, and that mixture training tolerates much higher repetition than single-source training: scarce target corpora can be reused 15-20 times, with the optimal number of repetitions depending on the target data size, compute budget, and model scale. Next, we introduce a repetition-aware mixture scaling law that accounts for the decreasing value of repeated target tokens and the regularizing role of generic data. Optimizing the scaling law provides a principled way to compute effective mixture configurations, yielding practical mixture recommendations for pretraining under data constraints.

Abstract

Efficient Language Model Training

Explore similar work

Scaling Laws for Mixture Pretraining Under Data Constraints

The Finetuner's Fallacy: When to Pretrain with Your Finetuning Data

Train Separately, Merge Together: Modular Post-Training with Mixture-of-Experts