Optimal Splitting of Language Models from Mixtures to Specialized Domains
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
| 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 |
| 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 |
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 |
| 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 |
| 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 |
| Parameter | Transform | Effective range |
| softplus | ||
| scaled sigmoid | ||
| softplus | ||
| sigmoid | ||
| sigmoid | ||
| sigmoid |
| 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 |
| 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 |
| Train | Test | |||||||
| Domain | MAE | MRE | MAE | MRE | ||||
| 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 |
| 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 |
| 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 |
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 161 papers