A data mixture refers to how different data sources are combined to train large language models, and selecting an effective mixture is crucial for optimal downstream performance. Existing methods either conduct costly searches directly on the target model or rely on mixture scaling laws that fail to extrapolate well to large model sizes. We address these limitations by introducing a compute-efficient pipeline for data mixture scaling. First, we propose CAMEL, a capacity-aware mixture law that models validation loss with the nonlinear interplay between model size and mixture. We also introduce a loss-to-benchmark prediction law that estimates benchmark accuracy from validation loss, enabling end-to-end performance prediction for the target model. Next, we study how to allocate a fixed compute budget across model scales to fit the law and reduce prediction error. Finally, we apply our method to Mixture-of-Experts models with up to 7B-A150M parameters to fit the law, and verify the optimal mixture derived from the law by extrapolating to a 55B-A1.2B target model. Compared to prior methods, we reduce mixture optimization costs by 50% and improves downstream benchmark performance by up to 3%.
Figures & tables
Figure 1: Mixture optimization on the target model under different compute budgets. We evaluate different mixture extrapolation methods by applying them to a larger target model with varying training FLOPS as mixture optimization costs. CAMEL, our proposed method, identifies high-quality data mixtures with even less than the cost of one full training pass on the target model. As the optimization budget increases, CAMEL achieves higher average benchmark accuracy than baseline methods while using less than 50% of the compute cost of the baseline.
Figure 2: End-to-end framework for data mixture extrapolation under model scaling. We first fit a loss-to-benchmark mapping to relate validation loss to downstream benchmark accuracy (Section 2.2 ). We then model validation loss as a function of model size and data mixtures using sampled (M,r) pairs from smaller models (Section 2.1 ). These components together enable extrapolation to large models and direct optimization of data mixtures for target-scale performance (Section 3 ).
Figure 3: Training loss observations for each domain across model sizes. We train on a mixed dataset of math and knowledge and log the training loss for each domain. While larger models reduce loss in both areas, the rates of reduction differ significantly. This non-uniform scaling implies that the effective parameters allocated to each domain are redistributed dynamically rather than proportionally as the model scales.
Figure 4: Error of loss-to-benchmark prediction. We model each downstream benchmark accuracy as a function of multiple validation losses. The scatter plots show predicted versus ground-truth scores on training and validation splits. The low prediction error demonstrates that validation losses can reliably predict downstream benchmark accuracy. See Appendix B.3 for details and results on other benchmarks.
Figure 5: Comparison between CAMEL and baseline scaling laws. We compare the fitting error of our proposed Capacity-Aware Mixture Law (CAMEL) with two baseline methods, DML ( Ye et al., 2025 ) and SODM ( Shukor et al., 2025 ) . CAMEL achieves consistently lower fitting error and exhibits more stable extrapolation behavior across model scales.
Figure 6: Sampling strategies for fitting the scaling law. We illustrate several strategies for selecting training configurations to fit the scaling law. Each subfigure corresponds to one sampling strategy, showing how combinations of model size and data mixture are chosen. Each circle represents one sampling point, and the circle area is proportional to the model size.
Figure 7: Comparison of sampling strategies under fixed compute. For each expression, five strategies are compared under varying mixture optimization cost. The Hourglass strategy consistently achieves the lowest prediction error.
Table 1: Benchmark performance under the Balanced objective. Our method achieves the highest weighted average score, computed using the benchmark optimization objective weights, and performs better than baseline mixtures on most benchmarks, indicating stable extrapolation to larger model scales.
Figure 8: Performance on held-out benchmarks. Our method achieves the highest average accuracy on benchmarks not used during optimization, indicating strong generalization beyond the proxy objectives. This observation is consistent with prior work ( Mizrahi et al., 2025 ) that mixtures optimized on diverse benchmarks generalize to unseen tasks.
Table 2: Benchmark performance under other targets. Our method achieves the highest weighted average score across all other specialized targets, outperforming the baselines and demonstrating robust mixture optimization under diverse training targets.
Figure 9: Optimal mixture weights across model sizes. The x-axis denotes activated model parameters and the y-axis shows dataset mixture weights. As model size increases, the proportion of Knowledge data grows, while Math and Code decrease, indicating stronger knowledge absorption in larger models.
Appendix figures & tables3 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 10: Ablation study of intrinsic domains. We vary the number of intrinsic domains k and evaluate the validation-loss prediction error. When k is small, the model has limited capacity, resulting in higher prediction error. As k increases, the error decreases and reaches its minimum around k=5 . Further increasing k leads to higher error, indicating overfitting and reduced robustness.
Figure 11: Predictions of all benchmarks. This figure shows predicted versus observed accuracies on all benchmarks. The close alignment indicates strong generalization of the proposed law.
Objective
MMLU
ARC-Challenge
BBH
GSM8K
Math
HumanEval (pass@1)
CEval
Balanced
0.20
0.10
0.15
0.15
0.15
0.15
0.10
Math Specialized
0.12
0.08
0.12
0.15
0.30
0.15
0.08
Code Specialized
0.10
0.10
0.10
0.15
0.15
0.25
0.15
Knowledge Specialized
0.30
0.16
0.14
0.10
0.04
0.04
0.22
Appendix
Table 3: Weights assigned to benchmarks under different training objectives. Each objective emphasizes a specific domain: the Balanced objective balances overall performance, the Math Specialized objective prioritizes mathematical reasoning, the Code Specialized objective emphasizes programming ability and the Knowledge Specialized objective emphasizes general knowledge ability.
Data mixture optimization studies how to allocate training examples across domains and tasks to improve model performance under data and compute constraints.
Upweighting high-quality data in LLM pretraining often improves performance, but in datalimited regimes, especially under overtraining, stronger upweighting increases repetition and can degrade performance. However, standard scaling laws do not reliably extrapolate across mixture recipes or under repetitions, making the selection for optimal data recipes at scaling underdetermined. To solve this, we introduce InfoLaw (Information Scaling Laws), a data-aware scaling framework that predicts loss from consumed tokens, model size, data mixture weights, and repetition. The key idea is to model pretraining as information accumulation, where quality controls information density and repetition induces scaledependent diminishing returns. We first collect the model performance after training on datasets that vary in scale, quality distribution, and repetition level. Then we build up the modeling for information so that information accurately predicts those model performance. InfoLaw predicts performance on unseen data recipes and larger scale runs (up to 7B, 425B tokens) with 0.15% mean and 0.96% max absolute error in loss, and it extrapolates reliably across overtraining levels, enabling efficient data-recipe selection under varying compute budgets.
In Large Language Model (LLM) training, data mixing plays a pivotal role in determining model performance. Recent methods optimize mixture weights via proxy models, but they rely on the assumption of static data distributions. As a result, when the underlying data pool shifts, these methods require costly retraining from scratch. This limitation restricts their ability to scale seamlessly from small settings to larger data pools and model sizes. In this paper, we propose CausalMix to address this limitation by casting data mixture optimization as a causal inference problem. We formulate the statistical features of the data pool as covariates and the domain mixture as the treatment. After fitting a causal model on 512 runs of Qwen2.5-0.5B to estimate the Conditional Average Treatment Effect (CATE), we extrapolate the optimal mixture for an 800K data pool and apply it to train a 7B model. Furthermore, we successfully generalize the framework to long chain-of-thought data on Qwen3-4B-Base. By leveraging causal modeling to isolate confounding biases, CausalMix infers state-dependent data mixtures offline, with the selected mixture fixed during each training run. Extensive experiments show that the mixture guided by CausalMix achieves competitive performance across downstream tasks compared with other baselines. The CATE Interpreter visualizes the learned mixing strategy. Overall, CausalMix offers a causal and interpretable framework for optimizing LLM data mixtures.
Zinan Tang, Yukun Zhang, Shaomian Zheng +7
Tsinghua University · Ant Group · Renmin University of China
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.