When do data mixtures improve scaling laws? Insights from high-dimensional regression
Authors: Diyuan Wu, Lehan Chen, Theodor Misiakiewicz, Marco Mondelli
Organizations: Institute of Science and Technology Austria (ISTA). · Department of Applied and Computational Mathematics, Yale University. · Department of Statistics and Data Science, Yale University.
Modern machine learning systems are trained on mixtures of data from different domains, and choosing the right mixture can substantially improve downstream performance. Despite an extensive literature on data mixing and reweighting, existing work is largely empirical and it remains unclear when auxiliary data genuinely improves scaling laws rather than merely providing more samples. To gain insight into this question, we study a high-dimensional mixed-data regression model with a shared regression function, heterogeneous covariances and noise levels, and dataset sizes that may grow at different rates. We establish the minimax risk under an ellipsoidal parameter constraint for the general covariance structure and derive deterministic equivalents for the test error of ridge regression under commutative covariances. We then specialize to a target domain and an auxiliary domain with aligned power-law covariance spectra, where the theory yields explicit scaling laws in terms of spectral decay, target regularity, and the relative growth of the two datasets. These laws identify regimes in which combining data mixtures provably yields a faster scaling rate than using either dataset alone. In particular, improving the scaling law requires a specific interplay between spectra and relative sample sizes of the domains. Our numerical experiments on language models exhibit the same qualitative phenomenon: appropriate data mixtures yield a faster decrease in target-domain test loss than training on either domain alone.
Figures & tables
Figure 1: Test error (dots) of mixed ridge regression under power-law covariance, together with its corresponding deterministic equivalent from Theorem 2 (solid lines). The grey , golden and black dash lines are the minimax rates of target-only, auxiliary-only and mixture datasets from Theorem 3 . The blue , green and red lines correspond to different regularizations with the optimal one in red . Here xi∼N(0,Ci), with Ci as in ( 6 ). We set α1=3.2 , α2=1.2 , β=0.6 , and d=512 .
Figure 2: Test error (dots) of mixed ridge regression, together with its corresponding deterministic equivalent (solid lines). For CIFAR-10, features are extracted via a pretrained CLIP ViT-B/32 encoder, classes 0 – 4 form the target domain and classes 5 – 9 form the auxiliary domain. For ImageNet-100, features are extracted via a pretrained ResNet-18 encoder, and the corresponding class splits are 0 – 49 and 50 – 99 . In both settings, we take d=512 and generate labels according to y=x⊤θ∗+ξ , with ξ∼N(0,0.1) . The shared θ∗ is supported on the top- 10 eigenvectors of (C1+C2)/2 , with random signs and coefficient magnitudes proportional to j−1/2 for j=1,…,10 .
Figure 3: Target-domain test negative log-likelihood (NLL) scaling of an 81.5M-parameter GPT-2-style language model trained from scratch on SlimPajama. Panels (a)–(c) report mean NLL on a fixed target-domain test set as a function of the reference target budget n . We use n∈{4,6,8,12,16,20}Mi , except for the γ2=1.75 experiment that uses n≤16Mi . Panel (d) compares the normalized target and auxiliary token frequencies. Token types are ordered by decreasing target-domain frequency. The light auxiliary curve shows the raw frequencies, and the darker, thicker curve shows their 100-token running-window average.
Appendix figures & tables2 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 4: Target-domain test negative log-likelihood (NLL) scaling of an 81.5M-parameter GPT-2-style language model trained from scratch on SlimPajama. Panel (a) reports mean NLL on a fixed target-domain test set as a function of the reference target budget n . We use n∈{4,6,8,12,16,20}Mi , except for the γ2=1.75 experiment that uses n≤16Mi . Panel (b) compares the normalized target and auxiliary token frequencies. Token types are ordered by decreasing target-domain frequency. The light auxiliary curve shows the raw frequencies, and the darker, thicker curve shows their 100-token running-window average.
γ2
Target-only
Auxiliary-only
Mixture
Mixture improvement (%) Target / Auxiliary
0.75
-0.8319
-0.6085
-0.8742
+5.1% / +43.7%
1
-0.8319
-0.8389
-1.0272
+23.5% / +22.4%
1.25
-0.8319
-1.1155
-1.1225
+34.9% / +0.6%
1.5
-0.8319
-1.4015
-1.2513
+50.4% / -10.7%
1.75
-0.8319
-1.5971
-1.3208
+58.8% / -17.3%
2
-0.8319
-1.7332
-1.4361
+72.6% / -17.1%
Appendix
Table 1: Linear OLS scaling slopes on the target test set for the 81.5M-parameter model over 4–20 M target datasets tokens. More-negative values indicate a faster decrease in target NLL, i.e. better scaling law. Each slope is obtained by fitting NLL=α+sln(n) over the six Target dataset sizes n∈{4,6,8,12,16,20} M within each seed and then averaging over three seeds. Improvement is 100(∣sMixture∣−∣sX∣)/∣sX∣ , reported against target-only and auxiliary-only, respectively.
Recent research has established empirical scaling laws to predict model performance on multi-domain data mixtures. However, a theoretical understanding of these model loss behaviors remains absent. In this work, we propose a unified framework to explain the underlying mechanics of data mixing. Our approach extends theoretical perspectives originally developed for standard neural scaling laws (e.g., Kaplan and Chinchilla) to the multi-domain setting. Based on the distributional assumption that domains overlap on fundamental skills while diverging on specialized skills, we identify two key factors that govern the domain losses of models trained on different data mixtures: \textit{Capacity Competition}, where the allocation of finite model capacity couples domain losses globally, and \textit{Noise Reduction}, where optimal weights shift toward harder-to-learn domains to minimize overall noise. Empirical evaluations show that our framework outperforms existing baselines by fitting the loss landscape with a lower Mean Relative Error and identifying higher-performing training mixtures. Most importantly, our model successfully extrapolates across scales, predicting highly effective mixtures for large, unseen scales using parameters fitted on smaller ones. In addition, our model achieves these results using significantly fewer parameters compared to previous empirical laws. Our code is available at https://github.com/meiqwq/Explaining-Data-Mixing-Scaling-Laws.
Rui Dai, Shuran Zheng
Beijing Institute of Technology · IIIS, Tsinghua University
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.
Pre-training data mixtures are commonly tuned by running small-scale experiments and extrapolating to the target training budget. When high-quality data is scarce and must be repeated, this extrapolation frequently fails, but the source of the failure has not been isolated. We show that a primary culprit is a repetition mismatch: because high-quality datasets are small, their repetition rate changes as the training budget grows, shifting the optimal mixture in ways that small-scale proxy experiments do not anticipate. A subsampling procedure that matches the target repetition rate controls for this effect. In a two-source setting combining limited high-quality data with web crawl, a single repetition-controlled experiment using only 1/16 of the target tokens recovers a mixture within 0.10 of the optimum on Wiki-Text for a 1.17B parameter model, compared to an error of 0.85 without repetition control. Achieving comparable accuracy without repetition control requires multiple training horizons, consuming 19%, 44%, and 94% of the target token budget when using the results from two, three, and four horizons respectively. With three data sources, the larger mixture space requires more than a single experiment to constrain, but the approach remains effective: at the 757M scale, just two repetition-controlled horizons recover the optimal mixture, outperforming baselines that instead require the full two-source experiments to construct. Our results reveal that repetition dynamics, not scale alone, shape whether small-scale mixture experiments generalize. More broadly, they suggest that data repetition deserves treatment as a first-class variable in mixture optimization, rather than an inconvenient side effect of limited data.