PaperlayerView exact revision on arXiv

Source-linked AI summary

Decouple Searching from Training: Scaling Data Mixing via Model Merging for Large Language Model Pre-training

Shengrui Li, Fei Zhao, Kaiyan Zhao, Jieying Ye, Haifeng Liu, Fangcheng Shi, Zheyong Xie, Yao Hu, Shaosheng Cao

arXiv:2602.00747v3cs.CLcs.AI

TL;DR

Pre-training data-mixture selection must balance general ability with difficult math and code tasks, while existing searches trade off proxy reliability against computational cost. DeMix trains component models once and uses weighted model merging to evaluate mixtures through training-free proxies. Across broad evaluations, it identifies mixtures that balance capabilities with greater efficiency, and the 22T-token DeMix Corpora provides validated mixtures for large-scale pre-training.

  • Problem

    Data-mixture optimization must balance general language ability with mathematical reasoning and code generation, but tiny proxies can fail on complex tasks while large-scale exploration is expensive.

  • Method

    DeMix trains component models on candidate datasets and creates mixture proxies through weighted model merging instead of separately training proxies for each ratio.

  • Results

    DeMix achieves higher proxy accuracy than training-based proxies at the same budget and comparable accuracy with approximately 6× less computation, using 200 versus 1200 units.

  • Takeaways & Limitations

    DeMix provides a mixture-search framework balancing sufficiency, accuracy, and efficiency, alongside a 22T-token corpus with validated mixture ratios for LLM pre-training.

  • Takeaways & Limitations

    Finer data partitioning expands the search space, reduces the benefit of merging similar datasets, increases computational cost, and complicates estimation for tiny mixture proportions.

Abstract

from arXiv · show

Determining an effective data mixture is a key factor in Large Language Model (LLM) pre-training, where models must balance general competence with proficiency on hard tasks such as math and code. However, identifying an optimal mixture remains an open challenge, as existing approaches either rely on unreliable tiny-scale proxy experiments or require prohibitively expensive large-scale exploration. To address this, we propose Decouple Searching from Training Mix (DeMix), a novel framework that leverages model merging to predict optimal data ratios. Instead of training proxy models for every sampled mixture, DeMix trains component models on candidate datasets at scale and derives data mixture proxies via weighted model merging. This paradigm decouples search from training costs, enabling evaluation of unlimited sampled mixtures without extra training burden and thus facilitating better mixture discovery through more search trials. Extensive experiments demonstrate that DeMix breaks the trade-off between sufficiency, accuracy and efficiency, obtaining the optimal mixture with higher benchmark performance at lower search cost. Additionally, we release the DeMix Corpora, a comprehensive 22T-token dataset comprising high-quality pre-training data with validated mixtures to facilitate open research. Our code and DeMix Corpora is available at https://github.com/Lucius-lsr/DeMix.

1. Introduction

Optimizing pre-training data mixtures requires balancing general language ability with mathematical reasoning and code generation, but existing proxy-based searches are either expensive or unreliable. DeMix addresses this by using weighted model merging to create training-free proxies, improving mixture search efficiency while preserving performance signals.

  • Data-mixture optimization must balance general-purpose language abilities with mathematical reasoning and code generation.
  • Large-scale proxy experiments can provide accurate signals but are computationally expensive and insufficient for systematically identifying optimal mixtures.
  • Tiny-scale proxy methods such as RegMix and CLIMB reduce training cost but often fail to generalize to complex math and code tasks.
  • DeMix constructs proxy models by weighted-merging component models trained on candidate datasets, enabling unlimited mixture proxies without training each one.
  • DeMix achieves comparable proxy accuracy with approximately 6× less computation budget, using 200 versus 1200 units.
  • DeMix outperforms state-of-the-art mixture methods across language understanding, mathematical reasoning, and code-generation benchmarks while requiring less computational budget.
  • Model-merging proxies preserve the performance ordering of reference models trained on real mixtures, providing a reliable signal for mixture selection.
  • The released DeMix Corpora contains 22T tokens of high-quality pre-training data with validated mixtures for direct reuse.

2. Method

DeMix constructs proxy models by merging component models trained on candidate datasets, then searches mixture weights through benchmarked iterative prediction rather than repeated proxy training. Its approximation relies on small parameter updates and weighted parameter averaging to represent models trained on mixed data.

  • Pipeline: DeMix organizes candidate datasets, trains specialized component models, merges them into proxies, and iteratively predicts improved mixture ratios.The pipeline moves from data preparation and component training to weighted merging, benchmarking, regression, and resampling.
  • Component Model Preparation: The method defines training from a shared base model, with each component parameterized as Θ_i = Θ_base + Δ(D_i).This formulation separates the common initialization from each candidate dataset’s parameter update.
  • Model Merging as Proxy: Under a small-update assumption, weighted parameter averaging approximates training on a weighted mixture of candidate datasets.The paper reports δ ≈ 10% in its experiments and cites empirical support for summing separate weight deltas.
  • Mixture Weight Optimization: Inference on merged proxies incurs no additional training cost, allowing mixture search through iterative optimization of the weights.The search therefore evaluates newly sampled ratios without training a new proxy model for each ratio.
  • Mixture Weight Optimization: For sampled mixture weights, DeMix weighted-merges component models, benchmarks each proxy, and records its average ranking across general, math, and code tasks.Ranking is used as the evaluation signal because it is robust to scale mismatches and directly relevant to mixture selection.
  • Mixture Weight Optimization: A predictor maps sampled mixture weights to ranking scores, selects promising ratios, and repeats sampling and prediction three times before averaging top candidates.The final averaged candidates define the mixed dataset used for pre-training.

3. Experimental Settings

Experiments use Qwen3 models, standardized training and evaluation settings, and comparisons with RegMix and CLIMB. Proxy consistency measures ranking and capability retention, while mixture quality evaluates the final model across general, math, and code benchmarks.

  • Models: Qwen3-1.7B determines the mixture, with validation on Qwen3-1.7B and Qwen3-8B after each component receives 50B tokens of general-data pre-training.The general-data stage is intended to maintain fundamental capabilities before candidate-dataset training.
  • Evaluation: Evaluation covers five general benchmarks, two code benchmarks, and two math benchmarks, with one benchmarking run costing 0.013B training tokens.General tasks are ARC-E, HellaSwag, WinoGrande, PIQA, and SIQA; code tasks are HumanEval and MBPP; math tasks are GSM8K and MATH.
  • Baselines: RegMix and CLIMB serve as baselines, using 112 proxies by default to match CLIMB’s configuration while aligning model, data, and budget settings.The study also varies proxy counts and excludes earlier methods described as inferior.
  • Proxy Consistency: Proxy accuracy measures Spearman rank correlation between 96 merged proxies and 96 large-scale reference models, including a separate correlation for the top 25%.The reference models are trained on a 50B-token corpus, and higher ρ indicates stronger ranking alignment.
  • Proxy Consistency: Capability Recovery Rate compares average proxy and reference benchmark scores to quantify retained absolute performance without additional training costs.Mixture quality instead evaluates the final 50B-token model using benchmark scores and macro-averaged rank across general, math, and code domains.

4. Experimental Results

Across proxy consistency, mixture quality, and robustness studies, DeMix identifies strong data mixtures with lower training costs than training-based baselines, while performance depends on proxy count, merging strategy, general-data proportion, scale, and partition granularity.

  • 4.1. Proxy Consistency: DeMix achieves a macro avg. ρ of 0.81 and top-25% ρ of 0.59 with 212B tokens, while training-based proxies reach 0.53 and 0.20 and require 1344B tokens for similar performance.The comparable-cost training-based approach uses 6.4× more tokens than DeMix.
  • 4.2. Mixture Quality: DeMix with 224 merged proxies achieves the top average rank of 24.00, outperforming RegMix and CLIMB under comparable or substantially larger training budgets.Neither 112 2B-trained proxies nor 28 8B-trained proxies surpass DeMix; even at 448B, RegMix and CLIMB remain behind.
  • 4.2. Mixture Quality: Increasing DeMix proxies from 56 to 224 improves average rank from 29.33 to 24.00, but performance declines at 448 proxies, indicating possible overfitting to noise.The gains from broader proxy searches may require substantial additional computation.
  • 4.3.1. Merging Methods: Linear merging achieves the best capability recovery rate and macro-average ρ among the evaluated merging methods while remaining hyperparameter-free.The comparison includes Multi-SLERP, DARE, Breadcrumbs, DELLA, and TIES.
  • 4.3.2. General Data Proportion: Reducing mixed general data to 25% lowers proxy accuracy to ρ = 0.667 and capability recovery to 0.796, with further declines at 0%.These results support retaining general data as a regularization component in candidate datasets.
  • 4.3.3. Validity on Larger Scale Models: DeMix mixture ratios transfer to Qwen3-8B, achieving a better average rank than alternative ratios and matching the observed effectiveness on Qwen3-1.7B.This indicates that the reported effectiveness extends beyond the smaller target model evaluated in the main experiments.
  • 4.3.4. Finer Partitioning of Training Data: Under a finer 15-category partition, DeMix retains the same proxy-accuracy trend and comparable Spearman correlations at lower cost than 12B-trained proxies.The main conclusions are reported as insensitive to partition granularity.
  • 4.3.4. Finer Partitioning of Training Data: Finer partitions enlarge the search space and computational burden, reduce the benefit of merging similar datasets, and make small-category ratios difficult to estimate reliably.These trade-offs can make optimized ratios for tiny categories less meaningful under a fixed search budget.

5. DeMix Corpora

DeMix Corpora combines high-quality general data with a validated mixture designed to balance general and domain-specific capabilities. Its HQ general subset outperforms other general datasets, while the mixed corpus achieves the best reported average rank.

  • General-data quality: The HQ general subset outperforms all other general datasets on general benchmarks when training Qwen3-1.7B from scratch.The comparison is shown in Figure 3.
  • Corpus and mixture design: DeMix Corpora includes high-quality general data and a validated mixture balancing multi-domain pre-training capabilities.Its final mixture is organized across domain, data category, and data origin levels.
  • Multi-domain performance: DeMix Corpora achieves the best average rank of 24.00 across general and domain-specific capabilities.The comparison is reported against existing mixed datasets in Table 9.
  • Multi-domain performance: Nemotron-Pretrain shows poor general capability because of insufficient general-domain training data, whereas SmolLM-Corpus has severely deficient math capability despite strong general data.These comparisons illustrate different trade-offs among existing mixed datasets.

6. Related Works

Data-mixture optimization has shifted from costly large-scale proxy training toward training-free model merging, but continuous pre-training mixture weights make the search problem more difficult than subset selection.

  • Data-mixture optimization: Large-scale proxy experiments can provide accurate mixture estimates but are computationally impractical for many research institutes.They use mid-sized models with substantial token budgets.
  • Model merging: Model merging combines structurally identical LLMs through arithmetic operations without additional training.Prior uses include regularization, generalization improvement, mitigation of catastrophic forgetting, and replacing fine-tuning.
  • Model merging: Pre-training mixture optimization is harder than binary dataset subset selection because mixing weights are continuous and the feasible space is unbounded.Merge-to-Mix enumerates binary subset choices, whereas pre-training requires searching continuous-valued weights.
  • Model merging: Recent work indicates that merging weight deltas from models trained on different datasets can substitute for directly merging their training processes when the models share a base.The passage describes this as highly effective under the stated condition.

7. Conclusion

The paper introduces DeMix, which uses weighted model merging to separate data-mixture search from costly proxy training. It reports strong multi-domain mixture selection and releases a 22T-token corpus with validated ratios.

  • Conclusion: DeMix decouples mixture search from proxy training by constructing proxy models through weighted model merging.This provides unlimited proxy models while maintaining fixed token budgets.
  • Conclusion: Across evaluations, DeMix finds a mixture balancing general language understanding, mathematical reasoning, and code generation.The conclusion frames this as the best data mixture across these capability demands.
  • Conclusion: DeMix Corpora is a 22T-token high-quality dataset accompanied by validated mixture ratios for large-scale LLM pre-training research.The release is intended to support reproducible research and practical pre-training.

Impact Statement

The paper presents a method for optimizing LLM pre-training data mixtures and a dataset assembled from filtered and merged sources under valid licenses. No associated legal or ethical risks were identified.

  • Impact Statement: The proposed method is designed to facilitate data-mixture optimization for large language model pre-training.
  • Impact Statement: The dataset is constructed by filtering and merging multiple datasets held under valid and appropriate licenses.
  • Impact Statement: The authors identify no associated legal or ethical risks for the dataset construction described.

A.1. Data Curation Pipeline

The DeMix Corpora pipeline curates heterogeneous general, mathematical, code, multilingual, and Chinese web data into candidate datasets, using filtering, labeling, evaluation, and mixture preparation for optimization.

  • The corpus combines general-domain, mathematical, code, multilingual, and Chinese web sources into candidate datasets for mixture optimization.
  • Global exact and fuzzy deduplication removes redundant content, while dataset-specific perplexity filtering removes extremely low-perplexity samples.Perplexity filtering reduces the data volume by approximately 2%.
  • A FastText classifier distinguishes higher- and lower-quality samples using positive and negative training sets containing over one million examples each.
  • Chinese web data is quality-labeled, classified, and resampled by removing extremely low-quality or undetermined samples and upsampling extremely high-quality samples.
  • A three-level domain taxonomy supports corpus annotation, retaining 3.63M high-confidence instances to train a model that labels the full dataset.
  • Data evaluation trains 1.7B-scale models on 50B tokens and uses stable general, math, and code benchmarks to assess candidate datasets.
  • Candidate preparation reduces mixture dimensionality to seven categories, mixes each candidate with 50% general data, and trains seven component models for later search.

A.2. Data Composition Across Stages

DeMix uses a dynamic three-stage pre-training composition: early training emphasizes broad general knowledge, while later stages increase high-quality math and code data and optimize their ratios.

  • Early training prioritizes data diversity and general knowledge, whereas later stages increasingly emphasize high-quality data for advanced capabilities.
  • Stage 1 adjusts mixtures using basic general benchmarks and quality-based upsampling or downsampling.
  • The three pre-training stages contain approximately 14T, 6T, and 2T tokens, respectively, with mixtures organized by domain, data category, and data origin.
  • Stages 2 and 3 use DeMix to optimize mixture ratios as data quality increasingly affects final training performance.
  • Stage 2 lowers the allowable repetition threshold to produce a more balanced data distribution than Stage 3.
  • DeMix Corpora composition and stage-specific token allocations are reported in Figure 5 and Table 10.
  • RegMix and CLIMB differ from DeMix in proxy-model construction and predictor training, while using the same large mixture-sampling set for comparison.

C. Detailed Mixtures in Experiments.

Table 11 provides the detailed mixture ratios used across the reported experiments.

  • Table 11 lists detailed mixtures from different experiments.

D. Additional Verification of the Model-Merging Approximation

The verification compares merged proxy models with models trained directly on mixed math and code data across mixture ratios and token budgets. Merging closely approximates direct mixed-data training when parameter updates are small, with consistency declining moderately as updates grow.

  • The approximation in Eq. 5 is motivated by prior work on model merging and small-update regimes.
  • The experiment tests one math dataset D1 and one code dataset D2 across several D1/D2 mixture proportions.
  • Model merge trains separate component models on D1 and D2, then merges their parameters according to the target mixture ratio.
  • Data mix directly trains a model on D1 + D2 using the same target mixture ratio.
  • Consistency is the average merged-model score divided by the data-mix score across tested mixture proportions, with values closer to 1 indicating better matching.
  • When δ is small, merged proxies show high consistency with directly trained mixed-data models.
  • As token budgets increase and δ grows, consistency moderately decreases, but merged proxies retain reasonably high consistency while avoiding separate training for every mixture.
Loading 2602.00747v3…