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

TL;DR

This work tackles the challenge of selecting optimal pre-training data mixtures for large language models by decoupling search from training. It introduces DeMix, a framework that builds training-free proxies through weighted model merging of component models trained on candidate datasets, enabling extensive search with fixed budgets. By proving proxy rankings align with fully trained references and delivering a high-quality 22T-token DeMix Corpora, the approach achieves better mixture quality at substantially lower cost than prior proxy-based methods. Empirical results show DeMix outperforms RegMix and CLIMB across general, math, and code benchmarks, and the corpora provide a practical resource for open research in efficient LLM pre-training.

Abstract

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.

Figures (5)

• Figure 1: Methods such as RegMix and CLIMB require extensive proxies. Scaling each proxy leads to unaffordable overall budget. While our DeMix only require a few component models to merge unlimited training-free proxies.
• Figure 2: Pipeline for DeMix. After (1) cleaning and categorizing massive data, (2) component models are trained on individual candidate datasets. Instead of large-scale training for every ratio, (3) weighted model merging serves as a computationally efficient proxy to estimate performance for various mixture ratios. Finally, (4) a predictor is trained on the benchmarked proxy models to regress the relationship between mixing ratios and performance, utilizing iterative resampling to converge on the optimal mixture.
• Figure 3: General performance of high-quality general datasets.
• Figure 4: The data mixture constructed using our DeMix framework. The three hierarchical levels from the inside out are domain, data category, and data origin.
• Figure 5: Data mixtures for the three stages of pre-training in the DeMix Corpora, with approximately 14T, 6T, and 2T tokens, respectively. The three hierarchical levels from the inside out are domain, data category, and data origin.