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

TL;DR

DeMix proposes using model merging as a cheap proxy for data mixing during LLM pretraining data-mixture search: train one component model per candidate data subset, then evaluate weighted linear combinations of those weights on benchmarks to score mixture ratios — bypassing the train-a-tiny-proxy-per-mixture step that prior work (RegMix, CLIMB) relies on. A regression predictor maps mixture ratios → benchmark performance and is iteratively resampled to converge on an optimum. The authors argue tiny-scale proxies (e.g. 100M-param model trained on a candidate mix) do not reliably transfer to 1.7B-scale training, motivating the at-scale-component-model approach; they validate proxy ranking against 96 reference models trained on sampled mixtures at 50B tokens, and ship a 22T-token “DeMix Corpora” with a validated mixture.

Key claims

Existing tiny-scale proxy models (RegMix-style) do not reliably predict optimal mixtures at the target scale; the paper argues the capability gap between proxy and target causes downstream-task ranking to break, particularly for math/code [§1, §2].
DeMix’s pipeline is four stages: (1) clean and categorize a multi-source corpus into candidate datasets, (2) train one component model per candidate dataset at scale, (3) construct mixture proxies by weighted linear merging of component weights, (4) regress mixture ratio → benchmark score and iteratively resample [Fig. 2].
Weighted model merging serves as a high-fidelity proxy for mixture-pretrained models: the cost of a benchmarking run on a merged proxy is reported as equivalent to training a 0.013B-parameter model [Table 1].
Proxy ranking consistency is measured by Spearman correlation between merged-proxy scores and 96 reference models trained from scratch on 96 sampled mixtures of 50B tokens each — this is the load-bearing validation experiment [§Proxy Consistency].
“Capability recovery” is reported as a second proxy-quality metric — how much of absolute target-model performance the merged proxy preserves, not just rank [§Proxy Consistency].
The released DeMix Corpora is 22T tokens; the high-quality general subset is reported to outperform all other general-corpus baselines when training Qwen3-1.7B from scratch on 50B tokens [§Evaluation, Fig. 3].
The validated optimal mixture is presented hierarchically — domain → data category → data origin [Fig. 4] — and is shipped alongside the corpus so downstream users do not re-run the search.
For math/code candidates, component models are trained from a pretrained checkpoint with 80%/60% general-data dilution, not from scratch on pure math/code — i.e. component models are not identical-data specialists, they are dilution-controlled specialists [§Data validation].

Method

The setup formalizes pretraining as Θ_mix = T(D_mix, Θ_base): a training operator T taking a mixed dataset and a base init to a final model. The paper’s central bet is the approximation T(Σ λ_i D_i, Θ_base) ≈ Σ λ_i T(D_i, Θ_base) — i.e. training on a mixture is approximately a linear combination, in weight space, of training-on-each-component runs. Concretely: pick K candidate datasets D_1…D_K (here a hierarchical taxonomy across general/math/code/etc.); train K component models Θ_i = T(D_i, Θ_base) once, at production scale; then for any candidate mixture ratio λ ∈ Δ^{K-1}, the proxy is the linear-merged Θ_λ = Σ λ_i Θ_i, evaluated cheaply on benchmarks. A regression predictor (DeMix calls this the “trained proxy” / predictor in the iterative loop) is fit on (λ, benchmark_score) pairs from the merged proxies and used to propose the next batch of mixture candidates, converging on an optimum without ever training a proxy-from-scratch on a candidate mixture.

The validation of the proxy itself is the expensive part: 96 sampled mixture ratios, each used to train a reference model on 50B tokens, gives the ground-truth ranking against which Spearman ρ is computed for merged-proxy scoring. The merged-proxy step is reported as roughly equivalent in cost to training a 0.013B model — so the search budget across the mixture simplex is dominated by the one-time component-model training, not by the mixture-search-loop itself.

The released artifact is the DeMix Corpora (22T tokens), a hierarchical multi-source pretraining corpus with the validated optimal mixture baked in.

Results

Proxy consistency (the load-bearing validation): against 96 ground-truth-trained reference models on sampled mixtures at 50B tokens, the merged-proxy ranking achieves a high Spearman correlation with reference rankings [§Proxy Consistency]. (Exact ρ not visible in the search-result snippets retrieved; see Table referenced under “Proxy Accuracy” in the paper.)
Cost: one merged-proxy benchmark run ≈ training a 0.013B-parameter model [Table 1], so evaluating thousands of mixture candidates is cheaper than training even a single small from-scratch proxy.
Downstream training (Qwen3-1.7B from scratch on 50B tokens): the DeMix-curated high-quality general subset beats all other general-corpus baselines on standard benchmarks [Fig. 3]; the full validated mixture is reported to “balance multi-domain capabilities” across general/math/code evaluation.
Mid-training multi-domain performance is reported in Table 9 for models trained on the final mixture for 50B tokens, against domain-specific data baselines.

Why it’s interesting

The premise — that the training operator T(·) is approximately linear in the data mixture, so weight-space merging of components substitutes for data-space mixing — is the same bet that MergeMix (arxiv 2601.17858, mid-training, 8B/16B with Spearman ρ>0.9 cross-scale) and Merge to Mix (arxiv 2505.16066, dataset mixing for fine-tuning) are making in adjacent regimes. Three independent papers in the same six-month window converging on “weight-merge as data-mix proxy” is enough that the team should have an opinion on whether the linearity hypothesis holds for Luma’s models.

Two concrete reasons for caution, per Kyle’s read:

Scale-transfer. The premise still requires component models to be trained at a usable scale — DeMix’s “at scale” is Qwen3-1.7B-class, not the target. For a frontier-scale run the team still extrapolates from smaller component runs to larger mixture runs, and the paper explicitly argues 100M → 1.7B does not transfer cleanly for tiny-scale mixture-proxies (RegMix-style). The merging trick doesn’t obviously make that scale-transfer problem easier — it sidesteps it for the within-scale search but reintroduces it for choosing the component-model scale.
Verification cost. The 96-reference-model validation is exactly the grid-search one would want to avoid; the paper validates that the proxy works at one scale, not that it generalizes across scales or model families. A Luma application would inherit an unverified ranking on a model family the paper hasn’t tested.

The 22T DeMix Corpora and the validated mixture are practically useful as drop-in pretraining data even if the search method itself doesn’t transfer. The orthogonal “LR is approximately independent of mix for non-dramatic mix changes” empirical observation (Kyle, validated at 1–2B scale) reduces the dimensionality of the mixture-search problem regardless of whether the merging proxy is the right tool — a useful prior to keep paired with this paper. Related: the closest existing wiki neighbour, Nemotron-CLIMB: CLustering-based Iterative Data Mixture Bootstrapping for Language Model Pre-training, makes the same iterative-predictor-loop bet but uses small proxy models rather than merged components — and inherits the proxy-transfer concern DeMix is trying to dodge.