Research map / Large language models / Reasoning and chain of thought

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

ICML 2026 · Large language models › Reasoning and chain of thought · approach: Routing, pruning and distillation for reasoning · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes DeMix, which trains per-domain component models at scale and uses weighted linear model merging as training-free proxies to search pre-training data mixtures, and releases the 22T-token DeMix Corpora.

Problem

Finding an effective data mixture for LLM pre-training balancing general and hard tasks (math, code) is hard: tiny-scale proxies are unreliable and large-scale proxy search is prohibitively expensive.

Approach

Decouple search from training: train a small set of component models on candidate datasets, then synthesize unlimited proxy models by weighted merging of their parameters (weights = mixture ratios), and use a LightGBM predictor with iterative sampling to find the best mixture.

Key results

• DeMix with 30B-token component models achieves macro-average Spearman rho 0.81 at 212B total budget, versus 0.53 for trained proxies at comparable cost; trained proxies need 1344B to reach 0.82.
• DeMix with 224 merged proxies achieves the best macro-average rank of 24.00 on Qwen3-1.7B, versus 27.67 for the best RegMix and CLIMB settings.
• On Qwen3-8B, DeMix with 448 proxies achieves the best average rank of 27.67, tied with the best CLIMB setting at 30.00 or better than RegMix (34.00).

Surprising findings

• Performance rank of DeMix declines when proxy count grows from 224 to 448, which the authors attribute to overfitting noise.
• Small 2B trained proxies perform worse than 8B-trained proxies under the same budget, suggesting tiny proxies are poor for math/code tasks.

Limitations

• Merged proxy consistency drops as token budget and parameter update magnitude grow (consistency 0.79 at 50B in the data-mix validation).
• Search restricted to ~7 coarse categories with non-general data under 50%; finer partitions increase search cost and reduce reliability.
• Relies on a linear-additivity approximation of weight deltas that holds only in small-update regimes.

Experimental setup

Datasets: FineWeb-Edu DCLM-baseline DOLMA-v1.7 Ultra-FineWeb Nemotron-CC FineMath MegaMath SwallowMath OpenCoder SwallowCode StarCoder

Benchmarks: ARC-E HellaSwag WinoGrande PIQA SIQA HumanEval MBPP GSM8K MATH

Metrics: Spearman's rho top-25% Spearman's rho capability recovery rate average rank

Baselines: RegMix CLIMB Uniform Heuristic

Builds on: RegMix CLIMB task arithmetic

Closest papers by meaning

• RegMix: Data Mixture as Regression for Language Model Pre-training (ICLR 2025 · similarity 0.91)
• Aioli: A Unified Optimization Framework for Language Model Data Mixing (ICLR 2025 · similarity 0.89)
• Data Mixing Laws: Optimizing Data Mixtures by Predicting Language Modeling Performance (ICLR 2025 · similarity 0.89)
• Data Mixing Optimization for Supervised Fine-Tuning of Large Language Models (ICML 2025 · similarity 0.89)
• Model Merging in Pre-training of Large Language Models (NeurIPS 2025 · similarity 0.88)
• MixMin: Finding Data Mixtures via Convex Minimization (ICML 2025 · similarity 0.88)