cs.LGMar 9, 2026

Capacity-Aware Mixture Law Enables Efficient LLM Data Optimization

Authors: Jingwei Li, Xinran Gu, Jingzhao Zhang

Organizations: Institute for Interdisciplinary Information Sciences, Tsinghua University · Shanghai Qizhi Institute

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

Abstract

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)(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 kk and evaluate the validation-loss prediction error. When kk is small, the model has limited capacity, resulting in higher prediction error. As kk increases, the error decreases and reaches its minimum around k=5k=5 . Further increasing kk 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.
ObjectiveMMLUARC-ChallengeBBHGSM8KMathHumanEval (pass@1)CEval
Balanced0.200.100.150.150.150.150.10
Math Specialized0.120.080.120.150.300.150.08
Code Specialized0.100.100.100.150.150.250.15
Knowledge Specialized0.300.160.140.100.040.040.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.
More in this topic

Data Mixture Optimization

Data mixture optimization studies how to allocate training examples across domains and tasks to improve model performance under data and compute constraints.

See all 40 papers

Explore similar work

CardsList

  1. InfoLaw: Information Scaling Laws for Large Language Models with Quality-Weighted Mixture Data and Repetition

    May 4, 2026Fengze Liu, Weidong Zhou, Binbin Liu +7Language Model PretrainingLanguage Model Scaling Laws

    1ByteDance · 2UC Santa Cruz

  2. CausalMix: Data Mixture as Causal Inference for Language Model Training

    Jul 1, 2026Zinan Tang, Yukun Zhang, Shaomian Zheng +7Data Mixture OptimizationLLM Training

    Tsinghua University · Ant Group · Renmin University of China

  3. Scaling Laws for Mixture Pretraining Under Data Constraints

    May 12, 2026Anastasiia Sedova, Skyler Seto, Natalie Schluter +1Language Model PretrainingLanguage Model Scaling Laws

    Apple · ITU