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

AI summary

Overview

Research area: Large Language Model pre-training, specifically data mixture optimization (deciding what proportion of general, mathematical, code, and multilingual data to feed a model during pre-training).

Technical level: Intermediate. The core idea is intuitive, but the justification rests on model-merging theory and rank-correlation metrics that assume familiarity with modern LLM training pipelines.

Scope: This paper introduces DeMix, a framework that replaces expensive proxy-model training with weighted model merging so that thousands of candidate data mixtures can be evaluated essentially for free, and releases a 22T-token pre-training corpus with validated mixture ratios.

What This Paper Is About

Choosing the right blend of training data is one of the highest-leverage decisions in LLM pre-training, because a model must simultaneously learn general language, math, and code. The standard way to find a good blend is to train many smaller "proxy" models on candidate mixtures and extrapolate — but small proxies are unreliable for hard tasks like math and code, while large proxies are prohibitively expensive. DeMix's goal is to break this trade-off by making each proxy model cost nothing to produce.

Key Contributions

1. The DeMix framework, which decouples data-mixture search from model training. Instead of training a new proxy for every candidate mixture ratio, it trains one "component model" per candidate dataset and constructs proxy models by weighted linear merging of those components. Merging weights correspond directly to data mixture ratios.

2. Empirical validation that merged models faithfully rank real mixtures. The authors show that model-merged proxies preserve the performance ordering of reference models actually trained on 50B tokens of the corresponding mixtures, and retain most of their absolute capability (a "capability recovery" rate up to 0.85).

3. A complete mixture-search pipeline built on these proxies: sample many mixture ratios from the simplex, score each merged proxy on a benchmark suite, fit a LightGBM regressor mapping ratios to average rank, and iteratively resample around the top candidates to converge on a final ratio.

4. DeMix Corpora, a 22T-token (15T original tokens) dataset spanning general, multilingual, math, and code domains, published together with validated mixture ratios — addressing the scarcity of benchmarked corpora that come with pre-validated mixtures.

Main Findings

• Merged proxies are ~6× cheaper than trained proxies for equal accuracy. DeMix reaches a macro-average Spearman correlation of 0.81 (and 0.59 on the top-25% of mixtures) using only 212B total tokens. The training-based baseline needs 1344B tokens to reach a comparable level of proxy accuracy — a 6.4× cost increase. Under matched budgets, DeMix scores 0.81 versus 0.53 for trained proxies.

• The merging approximation holds in practice. The ratio of parameter-update magnitude to total parameter magnitude is roughly 10% (δ ≈ 10%), well within the regime where the weight deltas from separately trained models approximately sum to the delta from joint training.

• Better final mixtures at lower cost. With 224 merged proxies, DeMix achieves the best macro-average rank (24.00) across general, code, and math benchmarks, beating RegMix and CLIMB — including RegMix/CLIMB configurations that use larger budgets or 8B-scale proxies instead of 2B.

• Capability recovery validates absolute performance, not just ordering. Merged proxies retain up to 0.85 of the reference model's benchmark scores, confirming they are usable proxies rather than merely rank-preserving abstractions.

• Simple linear merging wins. Compared against Multi-SLERP, Breadcrumbs, DARE, DELLA, and TIES, plain weighted linear merging achieves the best capability recovery (0.845) and macro-average ρ (0.787), while requiring no hyperparameter tuning.

• General data must be mixed into component datasets. Dropping the general-data portion of candidate datasets from 50% to 25% collapses proxy accuracy from ρ = 0.787 to 0.667; going to 0% drops it further to 0.652. Regularization with general data is essential.

• Proxy count matters, but overshoots hurt. Scaling proxies from 56 to 224 steadily improves final rank (29.33 → 24.00); pushing to 448 regresses, suggesting overfitting to noise in the regression surrogate.

• Results transfer to larger and finer-grained settings. Mixtures optimized on Qwen3-1.7B also rank well on Qwen3-8B, and a finer 15-category data partition reproduces the same trends as the main 7-category partition.

• DeMix Corpora is distinctive in composition. Among public corpora, it is the only one that combines multilingual coverage, math and code data, and a validated mixture ratio, and it achieves the best average rank (24.00) after 50B-token mid-training versus SmolLM-Corpus (31.33) and Nemotron-Pretrain (36.00).

Methodology in Plain English

The pipeline has four stages.

First, the team gathered and cleaned a very large pool of data from open sources and sorted it into a small number of candidate datasets — one for general web text, one for math, one for code, and so on. The main experiments use 7 categories; a finer 15-category variant is also tested.

Second, they trained a component model for each candidate dataset. All components start from a shared base model that was pre-trained from scratch on 50B tokens of general data, so they share a common origin. Each component is then further trained on its own domain-specific data blended with general data at a fixed 50/50 ratio, giving it domain specialization without losing general ability.

Third, they used model merging as a stand-in for training. Because each component's parameters can be written as base parameters plus a "delta" learned from its dataset, and because those deltas are small relative to the base parameters, the delta from training on a blend of datasets is approximately the weighted sum of the individual deltas. That means a weighted average of component model parameters approximates a model trained on the corresponding weighted blend of datasets. So any candidate mixture ratio — represented as weights summing to 1 — can be turned into an evaluable model instantly, with no training at all.

Fourth, they searched. They sampled mixture ratios broadly, merged a proxy for each, and scored each proxy on a benchmark suite covering general language, code, and math, converting scores into an average rank. A LightGBM regressor was trained on the (ratio, rank) pairs, then used to propose many new candidate ratios; the process repeated three times (64, 32, then 16 samples), and the top-ranked proposals were averaged to produce the final mixture. Because benchmarking (0.3 GPU-hours) costs roughly the same as training 0.013B tokens, evaluation is negligible relative to training.

Why This Matters

This work attacks a bottleneck that scales with compute budgets: the more capable your target model, the more expensive it is to figure out what it should be trained on. By making proxy evaluation essentially free, DeMix shifts the constraint from "how many proxies can we afford" to "how many candidate datasets do we want to consider" — which is a far better position for both large labs and resource-constrained researchers.

Real-world applications:

• Pre-training recipe design at scale. Teams training foundation models can search far more mixture configurations under a fixed budget, then commit the winning ratio to an expensive large run.
• Domain-specialized model development. Builders of coding assistants or math-reasoning models can use the framework to tune the balance between general data and their specialty domain without running dozens of throwaway training jobs.
• Open corpus construction. Releasing a large corpus with validated mixture ratios removes a major guesswork step for anyone reproducing or building on it.
• Small-lab and academic research. Groups without multi-million-dollar compute can now run a search procedure that previously required mid-sized proxy training runs, using the released 22T-token corpus as a starting point.

Industry relevance: the technique directly reduces the cost of one of the most expensive and least automated parts of building an LLM. It also complements data-quality filtering rather than replacing it — the corpus work shows that cleaning, deduplication, and classification remain prerequisites, while DeMix handles the proportioning question on top.

Future Directions

• Theoretical characterization of when merging fails. The method rests on a small-update assumption (δ ≪ 1). It is unclear at what model scale, token budget, or dataset divergence this approximation breaks down, or whether a more principled merging scheme is needed beyond that regime.
• Adaptive partition granularity. The 7-category and 15-category experiments expose a real trade-off: finer partitions enlarge the search space but weaken the benefit of merging similar datasets, and small categories yield unreliable ratio estimates. Automatically choosing the right granularity is an open problem.
• Better surrogates than LightGBM over benchmark ranks. The predictor is a simple regressor on rank scores. A more sample-efficient or theoretically grounded predictor could find better mixtures with fewer proxy evaluations.
• Extension beyond pre-training and beyond text. The authors note concurrent work applying similar ideas during mid-training; extending the approach to continued pre-training, post-training data mixtures, or multimodal corpora is a natural next step.

Target Audience

This paper is most valuable to LLM pre-training practitioners and ML engineers who own or influence data-mixing decisions and compute budgets. It also suits data-curation and dataset-construction teams interested in the released DeMix Corpora, researchers working on model merging who want to see the technique applied to a search problem rather than to fine-tuning, and graduate students or smaller labs looking for a tractable method to optimize mixtures without mid-sized proxy training runs. Readers should be comfortable with transformer pre-training terminology, Spearman rank correlation, and basic model-merging concepts.

Authors’ 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.