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

🎯 Pitch

DeMix finds optimal data mixtures for LLM pre-training without training proxy models, by simply merging a few specialized component models together. This decoupling achieves the same accuracy as prior search methods at a fraction of the compute cost and demonstrates, for the first time, that linear model merging reliably predicts the effects of data mixing at scale.

1. Executive Summary

This paper proposes Decouple Searching from Training Mix (DeMix), a framework that decouples data mixture optimization from proxy model training by using weighted model merging to construct proxy models from individually trained component models (e.g., merging a math-specialized model and a code-specialized model with specific weights to simulate training on a combined dataset). Instead of training separate proxy models for each sampled data mixture ratio—the approach used by prior methods like RegMix and CLIMB—DeMix trains a small set of component models on candidate datasets at scale and then linearly merges their parameters to synthesize an unlimited number of training-free proxy evaluations, achieving a Spearman correlation of 0.81 with reference models while using roughly 6.4× less computation (212B vs. 1344B tokens to reach comparable proxy accuracy). Across the Qwen3-1.7B architecture evaluated on general, math, and code benchmarks, DeMix produces data mixtures that yield the best average performance rank of 24.00 compared to baselines, establishing that model merging serves as a faithful and efficient proxy for data mixture selection only when parameter updates remain small relative to the initialization scale and when general data is mixed into candidate datasets to preserve capability recovery.

2. Context and Motivation

The Core Problem: Data Mixture Optimization Is Both Critical and Prohibitively Expensive

The fundamental challenge this paper addresses is deceptively simple: given a collection of candidate datasets (web text, math problems, code repositories, multilingual corpora), what proportions should you mix them in to train the best possible large language model? This question sits at the heart of LLM pre-training because the composition of the training corpus directly shapes what capabilities the model acquires. A model trained predominantly on web text will develop strong general language understanding but may struggle with mathematical reasoning or code generation. Conversely, over-emphasizing specialized domains can erode fundamental language competence. The paper frames this as a balancing act between general competence and proficiency on hard tasks (Section 1, paragraph 1):

"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."

This is not merely an academic concern. Every major LLM training effort—from GPT-4 (Achiam et al., 2023) to Gemini (Gemini Team et al., 2023), DeepSeekMath (Shao et al., 2024), and Nemotron (Basant et al., 2025)—devotes substantial resources to tuning data mixtures. Getting the mixture wrong means either wasted compute training on suboptimal data distributions or, worse, deploying a model with hidden capability gaps that only surface after training is complete. The stakes escalate with model scale: for a 100B+ parameter model trained on trillions of tokens, even a modest efficiency improvement in mixture selection translates to millions of dollars in compute savings and months of development time.

Why Existing Approaches Fail: The Sufficiency-Accuracy-Efficiency Trilemma

The paper identifies a three-way tension that no prior method resolves (Section 1, paragraph 4). Let me walk through each dimension and why satisfying all three simultaneously has proven elusive.

The Sufficiency Problem: You Need Many Evaluated Mixtures to Find the Best One

Optimizing data mixture ratios is fundamentally a search problem. The mixture space is a continuous simplex—if you have 7 candidate datasets, the feasible mixtures lie on a 6-dimensional simplex, and sampling sparsely risks missing high-performing regions. Prior automated methods like RegMix (Liu et al., 2024) and CLIMB (Diao et al., 2025) address this by training a regression predictor that maps mixture ratios to performance, then using the predictor to score many candidate mixtures without additional training. But training that predictor still requires a training set of evaluated mixture points—typically 112 different mixture ratios, each requiring its own proxy model. If you only evaluate 28 or 56 mixtures, the predictor has insufficient data to accurately model the mixture-to-performance landscape, particularly for the non-linear interactions that govern how math and code data interact during training.

This is what the paper means by sufficiency: a method must be able to evaluate enough mixture candidates to adequately cover the search space, or the "optimal" mixture it finds may simply be the best of a small, unrepresentative sample.

The Accuracy Problem: Small Proxies Don't Generalize to Hard Tasks

The training-based methods (RegMix, CLIMB) obtain their 112 proxy evaluations by training small models on small token budgets—typically 2B tokens per proxy model. This is the only way to keep total cost manageable: 112 models × 2B tokens = 224B tokens total. But there is a fundamental tension here that the paper makes explicit (Section 6.1):

"When jointly optimizing more challenging tasks such as math and code, tiny-scale proxy experiments tend to become inaccurate due to insufficient training."

The reasoning is straightforward. Mathematical reasoning and code generation are emergent capabilities in LLMs—they only manifest after the model has been exposed to a substantial volume of diverse training data and has developed sufficiently rich internal representations. A model trained on only 2B tokens, regardless of architecture size, simply hasn't developed the capacity to solve multi-step math problems or generate syntactically valid code. Evaluating such a model on GSM8K or HumanEval produces near-random performance, providing essentially no signal about which data mixture is better.

This is the accuracy dimension: proxy evaluations must faithfully rank data mixtures in the same order that full-scale training would. If the proxy model is too weak to exhibit the capabilities being optimized for, its rankings are meaningless. The paper quantifies this concretely in Table 2: training-based proxies with 2B token budgets achieve a macro-average Spearman's ρ of only 0.53 with reference models, and a top-25% Spearman's ρ of merely 0.17. This means the ranking signal is weak overall and particularly unreliable for identifying the very best mixtures—exactly the ones you most want to find.

The Efficiency Problem: Accurate Proxies Require Prohibitive Compute

The obvious fix for accuracy is to train each proxy model on more tokens. An 8B-token proxy model can develop sufficient math and code capabilities to provide meaningful evaluation signals. Table 2 shows that training-based proxies with 12B token budgets (1344B total) achieve a macro-average ρ of 0.82—comparable to what DeMix achieves. But the cost is staggering: 112 models × 12B tokens = 1344B tokens total, a computational budget that would consume thousands of GPU-hours and is entirely impractical for most research organizations. Even industry labs with substantial resources cannot afford to run such experiments routinely, especially as the number of candidate datasets grows.

This is the efficiency dimension: the search procedure must be affordable relative to the cost of the actual pre-training run it is optimizing. If the search costs as much as the training run itself, the optimization provides no net benefit—you might as well have just trained on a heuristic mixture and accepted the performance penalty.

The Trilemma in Practice

These three dimensions are in direct tension for training-based methods:

To improve sufficiency (more evaluated mixtures), you increase the number of proxy models, which multiplies the total cost.
To improve accuracy (better proxy signals), you increase the token budget per proxy, which also multiplies the total cost.
To maintain efficiency (affordable total cost), you must either reduce the number of proxies (sacrificing sufficiency) or reduce the token budget per proxy (sacrificing accuracy).

Figure 1 in the paper illustrates this trilemma visually: RegMix/CLIMB with 112 proxies requires scaling each proxy's token budget to improve accuracy, causing the total training budget to balloon linearly. There is no configuration where training-based methods achieve high accuracy, high sufficiency, and reasonable cost simultaneously.

The Gap in Existing Approaches

The paper situates its contribution against three categories of prior work, each with distinct limitations.

Large-Scale Manual Proxy Experiments

The dominant approach in major LLM development efforts is to run a small number of large-scale proxy experiments: train a mid-sized model (e.g., 8B parameters) on a few candidate mixtures using a substantial token budget (e.g., 100B tokens), evaluate on target benchmarks, and select the best-performing mixture (Li et al., 2025; Nie et al., 2025; Blakeman et al., 2025). The paper acknowledges this can produce accurate signals (Section 6.1):

"While such proxies can provide relatively accurate signals, they remain computationally expensive and are insufficient for systematically identifying optimal data mixtures."

The key word is insufficient. Even with substantial resources, a team might evaluate 3–5 candidate mixtures. But the space of possible mixtures over 7+ candidate datasets is enormous—there is no guarantee that any of those 3–5 candidates is close to optimal. The approach trades sufficiency for accuracy, and the search is fundamentally limited by human intuition about which mixtures to try.

Automated Methods with Tiny Proxies (RegMix, CLIMB)

RegMix and CLIMB represent the state of the art in automated mixture optimization. They both follow the same high-level strategy: (1) sample many mixture ratios, (2) train a small proxy model on each sampled ratio using a fixed (small) token budget, (3) evaluate each proxy on target benchmarks, (4) train a regression model (LightGBM) to predict benchmark performance from mixture ratios, and (5) use the predictor to score a large number of candidate mixtures and select the best. CLIMB adds iterative refinement, where the predictor is used to guide additional sampling toward promising regions of the mixture space.

These methods are appealing because they automate the search and can evaluate many more mixtures than manual approaches—RegMix and CLIMB typically use 112 proxy models. However, the paper identifies two fundamental weaknesses:

First, the tiny-proxy accuracy problem. The paper is explicit that CLIMB and RegMix were only validated for optimizing simple general capabilities (Section 6.1):

"However, such automated methods are only validated to optimize simple general capabilities. When jointly optimizing more challenging tasks such as math and code, tiny-scale proxy experiments tend to become inaccurate due to insufficient training."

This is not a minor limitation—it means these methods break down precisely when the data mixture problem is hardest. Optimizing for general language understanding alone is relatively straightforward: web text quality is the dominant factor, and even small models can distinguish good from bad general data. But when you need to balance general competence with math and code proficiency, the proxy models must be capable enough to exhibit differential performance on math and code benchmarks, which requires larger token budgets.

The paper provides empirical evidence for this claim in Table 3, though the evidence is somewhat indirect. RegMix and CLIMB with 8B-token proxies (28 or 56 total) generally outperform their 2B-token counterparts (112 total) despite having fewer proxies—suggesting that proxy accuracy matters more than proxy count for this task. But even the 8B-token variants require substantial total compute (224B–448B tokens) and still underperform DeMix in the final mixture quality evaluation.

Second, the scaling constraint. Even if you accept the accuracy limitations, training-based methods cannot scale sufficiency without scaling cost. If you need 200 proxies instead of 112 to adequately cover the mixture space—perhaps because you have more candidate datasets or need finer-grained search—the cost doubles. DeMix, by contrast, can generate an arbitrary number of merged proxies at near-zero marginal cost once the component models are trained.

Loss-Based Methods (DoReMi, Rho Loss)

The paper explicitly excludes earlier methods like DoReMi (Xie et al., 2023) and Rho Loss (Mindermann et al., 2022) from comparison (Section 3.2):

"We exclude earlier inferior methods, including DoReMi and Rho Loss, as they depend on evaluation loss instead of proxies."

These methods optimize data mixtures to minimize a reference model's loss on target domains, rather than using proxy evaluations of downstream benchmark performance. The limitation is that language modeling loss is a weak proxy for downstream task performance, especially for complex reasoning tasks where the relationship between perplexity and benchmark accuracy is non-monotonic. A mixture that minimizes loss on math tokens might improve math benchmark performance, or it might simply overfit to surface-level patterns in the math corpus without developing genuine reasoning ability.

Model Merging in Data Selection (Merge-to-Mix)

The paper acknowledges that model merging has been explored for data selection, but draws a sharp distinction from prior work. Merge-to-Mix (Tao et al., 2025) uses unweighted averaging of fine-tuned models to enumerate binary subset choices for discarding datasets during fine-tuning. The paper notes (Section 6.2):

"In contrast, optimizing pre-training data mixtures is substantially more challenging: mixing weights are continuous-valued and the feasible space is unbounded."

Binary subset selection (include dataset A or not) is a much simpler problem than continuous mixture optimization. With 7 candidate datasets, there are only $2^7 = 128$ possible binary subsets—exhaustive enumeration is feasible. But continuous mixture ratios over 7 datasets live in a 6-dimensional simplex, requiring a fundamentally different search strategy.

The Theoretical Foundation: Why Model Merging Can Work as a Proxy

Before describing how DeMix positions itself, I need to establish why model merging is even a plausible approach to this problem, because the paper's theoretical motivation is essential context for understanding its contribution.

The key insight draws on recent empirical and theoretical work showing that when multiple models share the same initialization and are fine-tuned on different datasets, their parameter updates (weight deltas) are approximately additive. The paper formalizes this in Section 2.3:

Let $\Theta_{\text{base}}$ be the parameters of the base model. Let $T(\mathcal{D}, \Theta_{\text{base}})$ be the parameters after training on dataset $\mathcal{D}$ . The weight delta is:

$\Delta(\mathcal{D}) \triangleq T(\mathcal{D}, \Theta_{\text{base}}) - \Theta_{\text{base}}$

The critical approximation—which the paper attributes to prior work by Qin et al. (2022), Wu et al. (2025), and Lin et al. (2025b)—is:

$\Delta(\mathcal{D}_i \cup \mathcal{D}_j) \approx \Delta(\mathcal{D}_i) + \Delta(\mathcal{D}_j)$

In words: the parameter update from training on a union of datasets is approximately the sum of the parameter updates from training on each dataset separately. If this holds, then the model trained on a weighted mixture $\mathcal{D}_{\text{mix}} = \sum \alpha_i \mathcal{D}_i$ has parameters approximately equal to the weighted average of the separately trained component models:

$\Theta_{\text{mix}} = T(\mathcal{D}_{\text{mix}}, \Theta_{\text{base}}) \approx \sum_{i=1}^N \alpha_i \Theta_i$

This approximation is only valid under a specific condition that the paper quantifies explicitly (Equation 4): the magnitude of parameter updates must be small relative to the initialization scale:

$\delta = \frac{\sum |T(\mathcal{D}, \Theta_{\text{base}}) - \Theta_{\text{base}}|}{\sum |T(\mathcal{D}, \Theta_{\text{base}})| + \sum |\Theta_{\text{base}}|} \ll 1$

The paper reports that in their experiments, $\delta$ is approximately 10% (Section 2.3), which satisfies this small-update assumption. Table 12 in Appendix D empirically validates the approximation: when $\delta$ is small (3.10% at 2B tokens), the consistency between merged and data-mix models is high (0.97–1.04 for math and code benchmarks). As $\delta$ grows (10.50% at 50B tokens), consistency degrades to 0.75–0.79, but remains high enough to preserve ranking information.

This theoretical foundation is crucial because it explains why DeMix works and when it might fail. The method is not a general-purpose replacement for training—it is specifically applicable when (1) component models share a common initialization, (2) parameter updates are relatively small, and (3) the goal is ranking consistency rather than exact performance recovery. These constraints are satisfied in the late-stage pre-training setting that DeMix targets, where component models are fine-tuned from a shared base that has already acquired broad language capabilities, and the additional domain-specific training produces relatively small parameter deltas.

The paper also makes a subtle but important observation about what the approximation does not capture. When you train on a genuine mixture of datasets $\mathcal{D}_i$ and $\mathcal{D}_j$ simultaneously, the model can learn interactions between the domains—for example, mathematical notation appearing in code documentation might transfer between the math and code domains. Separate training followed by merging cannot capture such cross-domain synergies because each component model only sees one domain. The merged proxy is therefore a lower bound on what genuine mixed training would achieve, and the ranking signal it provides reflects the direct, additive contributions of each dataset rather than their interactions. This is not necessarily a weakness for mixture selection—the additive signal may actually be cleaner for ranking purposes—but it does mean the absolute performance of merged models will generally underperform equivalently trained mixed-data models, which the paper quantifies through the capability recovery metric (Table 2, averaging 0.83–0.85 for DeMix's best configurations).

How DeMix Positions Itself

DeMix's central claim is that it breaks the trilemma by making proxy cost independent of the number of evaluated mixtures. Rather than training a new model for each sampled mixture ratio, DeMix:

Trains a fixed set of $N$ component models—one per candidate dataset—at sufficient scale to develop the target capabilities (30B–50B tokens each, producing $\delta \approx 10\%$ ).
Constructs proxy models via weighted linear merging of these component models, where the merging weights represent the target data mixture ratios.
Evaluates the merged proxies on benchmarks to obtain performance signals.
Trains a regression predictor on the (mixture ratio, performance) pairs and uses it to search the mixture space.

The critical property is that the cost of component model training is a one-time upfront investment. After those $N$ component models exist, generating 100, 1,000, or 10,000 merged proxies costs only inference-time benchmarking—which the paper quantifies in Table 1 as equivalent to training 0.013B tokens per evaluation, negligible compared to the 2B–12B tokens required per training-based proxy. This is what enables the sufficiency dimension: DeMix can afford to evaluate many more mixture ratios than training-based methods, leading to better coverage of the search space and ultimately better discovered mixtures.

The paper positions this as a decoupling of search from training. Training-based methods couple these: every new mixture ratio explored requires a new training run. DeMix separates them: training happens once (the component models), and search happens independently (via merging and benchmarking), with essentially zero marginal cost per new mixture evaluated.

The positioning is explicitly framed against RegMix and CLIMB, which are described as the state-of-the-art automated methods but limited by their reliance on tiny proxies that cannot capture math and code capabilities. DeMix is not positioned as replacing large-scale manual proxy experiments entirely—the component models themselves require substantial training (30B tokens each, totaling 212B for 7 components—Table 2)—but rather as making the search process dramatically more efficient once those component models exist. The paper acknowledges this cost and frames it as amortized over the many proxy evaluations it enables.

The paper also positions the DeMix Corpora release as filling a distinct gap: "a notable lack of benchmarked corpora with validated data mixture ratios that can be directly reused for large-scale pre-training" (Section 1, paragraph 4). Table 8 compares DeMix Corpora against existing public pre-training datasets across dimensions of multilingual support, math/code inclusion, and validated mixture ratios. The only prior corpora with math/code data (DOLMA, SmolLM-Corpus, Nemotron-Pretrain) lack validated mixture ratios—their compositions were determined heuristically rather than through systematic optimization. DeMix Corpora provides both the data and the evidence that its mixture ratios are close to optimal, reducing the burden on future researchers who would otherwise need to replicate the mixture optimization process themselves.

The Practical Stakes

To appreciate why this problem matters beyond academic interest, consider the sequence of decisions a team faces when preparing a large-scale pre-training run:

They collect candidate datasets across general web text, mathematics, code, multilingual data, and potentially other domains.
They must decide how many tokens to allocate to each source, with the total training budget fixed (e.g., 2T tokens for a full training run).
Each allocation decision has irreversible consequences: if they under-allocate to math, the final model will have permanently weaker mathematical reasoning, which cannot be fully recovered through post-training (fine-tuning, RLHF) because the foundational representations were never adequately learned.
The cost of running even a single large-scale proxy experiment (training an 8B model on 100B tokens) is hundreds of GPU-hours. Running a grid search over 100 candidate mixtures at this scale is completely infeasible—hence the reliance on heuristics, manual tuning, or automated methods with small proxies.

DeMix's contribution is to make it feasible to evaluate hundreds of candidate mixtures at a quality level comparable to large-scale proxies, at a total cost dominated by the one-time component model training. For a team already planning to train at the 1.7B–8B scale for their final model, the component model training represents an incremental investment that pays for itself through better mixture discovery. This shifts the economics of data mixture optimization from "try a few educated guesses" to "search systematically," which is the same transition that hyperparameter optimization underwent with the introduction of Bayesian optimization and neural architecture search.

3. Technical Approach

3.1 Reader Orientation

What the system is: DeMix is a computational framework that finds the best proportions for mixing different types of training data (general text, math, code, multilingual) before large-scale language model pre-training, without needing to train hundreds of separate models to test each candidate mixture.

What problem it solves and the shape of the solution: The framework solves a search problem—finding the optimal data mixture ratios in a high-dimensional continuous space—by constructing cheap, training-free proxy models through weighted averaging of a small set of pre-trained component models, then using those proxies to train a predictor that guides the search toward high-performing mixtures.

3.2 Big-Picture Architecture (Diagram in Words)

DeMix operates as a four-phase pipeline, visualized in Figure 2, with information flowing sequentially from data preparation through to final mixture selection:

Dataset Preprocessing (Phase 1): Massive raw data from heterogeneous sources (web, math, code, multilingual) enters the system, undergoes a rigorous cleaning pipeline (deduplication, perplexity filtering, FastText quality classification, and instance-level labeling), and emerges as a set of candidate datasets, each representing a distinct domain or data category ready for mixture optimization.
Component Model Preparation (Phase 2): A single base model is trained from scratch on a general-purpose dataset to establish foundational language capabilities. Then, for each candidate dataset, a separate component model is initialized from this base and further trained on a 50-50 mixture of that candidate dataset and general data. Each component model specializes in its target domain while retaining general competence, producing a set of models that serve as building blocks for subsequent merging.
Model Merging as Proxy (Phase 3): This is the core innovation. Given a sampled mixture ratio—for example, "30% math data, 25% code data, 45% general data"—the corresponding component models are linearly merged with those exact weights to produce a proxy model. This proxy is evaluated on a suite of benchmarks (general reasoning, math, code) to obtain a performance score at near-zero training cost, since merging is a simple weighted sum of parameter vectors and benchmarking is inference-only.
Mixture Weight Optimization (Phase 4): The (mixture ratio, benchmark performance) pairs from Phase 3 are used to train a LightGBM regression predictor that maps mixture weights to predicted performance. This predictor is then used to score a vast number of newly sampled candidate mixtures, the top-ranked candidates are selected, and the process iterates three times—each iteration refining the predictor's accuracy in high-performing regions of the mixture space. The final optimal mixture is the average of the top candidates from the last iteration.

3.3 Roadmap for the Deep Dive

First, the dataset preprocessing pipeline—because the quality and categorization of the input data determines what mixture dimensions even exist to optimize over, and the cleaning choices (deduplication thresholds, perplexity cutoffs, FastText training, Chinese-specific quality filtering, instance-level labeling) directly affect downstream component model behavior.
Second, the formal data mixing objective and the component model training protocol—because component models are the fixed-cost investment that the entire framework depends on, and understanding how they are trained (initialization from a shared base, the 50% general data regularization ratio, the token budgets) is essential for understanding why model merging works as a proxy.
Third, the model merging proxy mechanism in detail—including the formal definition of weight deltas, the small-update assumption, the additivity approximation, the merging equation, and the empirical validation in Table 12—because this is the paper's central theoretical contribution and the mechanism that enables decoupling search from training.
Fourth, the iterative mixture weight optimization procedure—including the sampling strategy, the LightGBM predictor training, the evaluation metric (average ranking across benchmarks), the iterative resampling protocol, and the final mixture selection rule—because this is how the proxy evaluations are converted into an actionable data mixture recommendation.
Fifth, the design choices and their justifications—including why linear merging over alternatives like TIES or DARE, why 50% general data in candidate datasets, why ranking-based evaluation instead of raw accuracy, why iterative resampling, and why averaging top candidates rather than selecting the single best—because understanding these choices reveals the engineering wisdom that makes the framework robust.

3.4 Detailed, Sentence-Based Technical Breakdown

This is primarily a systems and methodology paper whose core idea is that weighted model merging of component models trained on individual candidate datasets can serve as a training-free proxy for evaluating data mixture ratios, decoupling the cost of mixture search from the cost of proxy model training and enabling substantially more extensive exploration of the mixture space.

Dataset Preprocessing Pipeline

The quality and categorization of the pre-training corpus fundamentally determines what mixture optimization can achieve—if the underlying data is noisy or mislabeled, even an optimal mixture ratio will underperform. The DeMix pipeline invests significant engineering effort in data preparation before any model training begins, with specific filtering thresholds and classifier training procedures that affect downstream proxy fidelity.

Data Collection. The raw data is gathered from heterogeneous open-source sources spanning four categories: general-domain corpora (FineWeb-Edu, DCLM-Baseline, DOLMA-v1.7, Ultra-FineWeb, and NVIDIA Nemotron), mathematical datasets (FineMath, MegaMath, SwallowMath), code datasets (OpenCoder, SwallowCode), and multilingual/reasoning data. The full source list appears in Table 10, with 13,417B total raw tokens across all sources. The general-domain data alone constitutes the vast majority: FineWeb-Edu (199B), Nemo-CC-High (385B), Nemo-CC-Synth (1106B), Nemo-CC-QA (546B), plus medium-quality tiers from DCLM, DOLMA, and web crawls totaling thousands of billions of tokens.

Global Deduplication. The paper applies both exact deduplication and fuzzy deduplication across most datasets, acknowledging a tradeoff identified in prior work (FineWeb, FineWeb2): global deduplication may remove high-quality samples that happen to be near-duplicates. However, for large-scale web data where the volume is large enough to absorb this loss, the authors argue the quality improvement from removing redundant content outweighs the loss of legitimate near-duplicates. The fuzzy deduplication procedure extracts 24-grams (contiguous 24-token sequences) from each document, computes 260 MinHash functions per document (a locality-sensitive hashing technique that maps similar documents to similar hash values with high probability), partitions the 260 hashes into 20 bands of 13 hashes each, and flags any pair of documents that share an identical 13-MinHash signature in any band as duplicates. This configuration targets documents with at least 90% similarity—the theoretical guarantee is that two documents with Jaccard similarity $s$ will hash to the same bucket in a given band with probability $s^{13}$ , and across 20 bands, at least one collision occurs with probability $1 - (1 - s^{13})^{20}$ , which exceeds 0.5 for $s \gtrapprox 0.9$ .

Perplexity Filtering. Each dataset is scored using the Qwen3-0.6B base model as a perplexity evaluator—for each document, the model computes the per-token negative log-likelihood under its own learned distribution, and documents with extremely low perplexity are removed. The thresholds are dataset-specific and determined by manual inspection: human annotators examine samples at different perplexity percentiles to identify where the data transitions from "suspiciously repetitive or templated" to "legitimately coherent." The overall data reduction from perplexity filtering is approximately 2% of the total corpus. This is notably conservative compared to some prior work (e.g., FineWeb-Edu's aggressive perplexity-based selection), reflecting a philosophy that perplexity is a coarse signal and over-filtering risks removing diverse but legitimate content.

FastText Quality Classification. A binary quality classifier is trained using FastText (a lightweight, efficient text classification library based on n-gram features and a linear classifier, chosen for its speed on billion-document-scale corpora). The training data construction is careful: both the positive set (high-quality documents) and negative set (low-quality documents) contain over one million samples. The positive set draws from randomly sampled data, low-perplexity data, and high-quality subsets from ELI5-category (long-form question answering, Fan et al., 2019) and OpenHermes-2.5 (synthetic instruction data, Teknium, 2023). The negative set is partially sourced from Falcon-RefinedWeb (Penedo et al., 2023)—a dataset known to contain substantial low-quality web text—and high-perplexity data. A specific engineering concern is addressed: high-perplexity data tends to be dominated by short samples (single sentences, fragments, navigation text), which would bias the classifier toward learning that "short = low quality." To mitigate this, the negative set construction applies weighted sampling that oversamples long documents from the high-perplexity pool, increasing the likelihood that the classifier learns to identify genuinely low-quality long-form content (e.g., boilerplate, SEO spam, machine-generated nonsense at paragraph length) rather than simply learning a length heuristic.

To validate that the positive and negative sets are genuinely distinguishable, the authors train Qwen3-0.6B from scratch on each set separately and confirm that the model trained on the positive set achieves lower average evaluation loss than the model trained on the negative set. This is a critical validation step: if the two sets were not actually different in quality, the classifier would be learning noise, and downstream filtering would be arbitrary. The trained FastText classifier removes approximately 3% of the English corpus and is described as effectively identifying low-quality web and code data.

Chinese-Specific Quality Filtering. For the Chinese general-domain web corpus, the paper develops a dedicated quality assessment and sampling framework inspired by FineWeb-Edu's approach of using educational value as a proxy for general text quality. The key insight is that education-related and information-dense signals correlate with high-quality text across languages, but the specific linguistic features that indicate quality differ between English and Chinese (e.g., Chinese web text has distinct patterns of low-quality content like clickbait, content farms, and machine-translated material). The framework categorizes text samples into five quality levels:

Undetermined: Non-Chinese or nonsensical content that cannot be meaningfully assessed.
Extremely low quality: Text with severe linguistic deficiencies, incoherence, or obviously machine-generated patterns.
Low quality: Text that is grammatically acceptable but lacks informational density or structural coherence.
High quality: Text with good linguistic completeness, reasonable information density, and coherent structure.
Extremely high quality: Text exhibiting strong knowledge-bearing characteristics, explanatory depth, and structural sophistication—without being restricted to explicitly educational contexts (the labeling criteria consider knowledge-bearing and explanatory characteristics as proxies for quality rather than requiring the text to be from educational sources).

The labeling pipeline works as follows: first, 5,000 Chinese samples are manually labeled by human annotators using the five-level criteria. These are used to fine-tune a 32B-parameter language model—a large enough model to internalize the nuanced quality distinctions—to serve as an automatic annotator. This 32B model then annotates approximately 1 million Chinese web documents, producing pseudo-labels. These pseudo-labeled samples are used to train a lightweight text quality classifier based on gte-multilingual-base as the text embedding backbone (a multilingual sentence embedding model) with a classification head on top. During corpus construction, this classifier removes samples labeled as "extremely low quality" or "undetermined," and upsamples samples labeled as "extremely high quality," thereby shifting the overall quality distribution of the Chinese training data upward.

Instance-Level Data Labeling. To understand and control the domain distribution of the pre-training corpus, the paper builds a three-level hierarchical taxonomy. Level-1 labels follow the standard graduate disciplinary classification (a formal categorization of academic fields). Level-2 and Level-3 labels are derived under each Level-1 parent using Gemini 2.5 Pro and GPT-4o—two of the most capable LLMs at the time of writing—to generate fine-grained subcategories. The annotation of the full pre-training dataset uses Qwen3-235B-A22B (a mixture-of-experts model with 235B total parameters and 22B active parameters per token), retaining 3.63 million high-confidence labeled instances (those where the model's predicted label probability exceeds some threshold, though the exact threshold is not specified). These high-confidence instances are used to train a smaller 4B-parameter model, which then labels the full pre-training dataset—a standard distillation approach where a large, expensive model teaches a smaller, cheaper model to perform the same task at scale.

Data Evaluation. Before committing to any mixture optimization, the paper validates the quality of individual datasets by training Qwen3-1.7B from scratch on 50B tokens of each dataset. For general datasets, training starts from scratch (random initialization). For math and code datasets—where the model cannot learn meaningful capabilities from domain-specific data alone without foundational language skills—the training mixes in 80% or 60% general data (the specific proportion varies by dataset) and starts from a pre-trained mid-checkpoint that has already acquired basic language competence. The benchmarks used for evaluation are chosen to ensure that the small model can exceed random performance after limited training and maintain stable ranking across subsequent training stages—this stability requirement is critical because if the relative ordering of datasets changed as training progressed, evaluations at 50B tokens would not predict behavior at larger scales.

Candidate Data Preparation for DeMix. After individual dataset evaluation, the paper applies several engineering decisions to reduce the dimensionality of the mixture optimization problem:

For general data, only the highest-quality tier is retained—medium and low-quality general data are excluded from the mixture search entirely, since their contribution is dominated by the high-quality tier and including them would add unnecessary degrees of freedom.
For math and code data, the lowest-quality tiers are discarded, and the remaining corpora are stratified by both category (e.g., synthetic math problems vs. web-extracted math) and quality level. Similar sources are merged to reduce the total number of candidate datasets.
The final result is seven dataset categories that require fine-grained mixing: General, Multilingual, Math-1, Math-2, Math-3, Code-1, Code-2, and Code-3 as shown in Table 11. The three math and three code categories represent different quality tiers or sub-domains within those broader categories.

This dimensionality reduction is described as a "common engineering trade-off" and "standard practice for cross-domain data mixing," citing industrial pre-training efforts (Basant et al., 2025; Feng et al., 2024). While reducing the number of candidate datasets may lower the theoretically attainable performance ceiling (since it precludes independent control of finer-grained data sources), the paper argues this is a "necessary compromise given the exorbitant cost of large-scale pre-training." The search burden of identifying a high-performance mixture in a 6-dimensional simplex is already substantial; adding more dimensions would make adequate coverage of the search space infeasible even with DeMix's efficient proxy mechanism.

The 50% General Data Regularization Constraint. For each of the seven candidate datasets, the final training data is constructed by mixing the candidate dataset with 50% general data. This constraint serves a dual purpose. First, it is a regularization technique: an excessive proportion of domain-specific data (e.g., 90% math) can severely degrade the model's general language capabilities during LLM pre-training, as the model overfits to domain-specific patterns, vocabulary, and reasoning styles at the expense of broad linguistic competence. This phenomenon is well-documented in prior work (Lin et al., 2025a; Bansal & Sanghavi, 2025; Allal et al., 2025b). Second, it defines the searchable subspace: the mixture optimization can only explore ratios where non-general data accounts for at most 50% of the total corpus. The paper argues this constraint is consistent with industry consensus—the non-general data should not dominate the pre-training corpus, with representative ratios from Nemotron (Basant et al., 2025) and SmolLM2 (Allal et al., 2025a) showing general data proportions well above 50%.

The Formal Data Mixing Objective and Component Model Training

The goal of data mixing is to identify a weighted combination of candidate datasets that maximizes model performance across target benchmarks. The paper formalizes this objective before describing how DeMix approximates it.

The Data Mixing Objective. Given $N$ candidate datasets $\{\mathcal{D}_1, \mathcal{D}_2, \ldots, \mathcal{D}_N\}$ , the target is to find a mixture distribution:

$\mathcal{D}_{\text{mix}} = \sum_{i=1}^{N} \alpha_i \mathcal{D}_i$

where $\alpha_i \geq 0$ are the mixture weights representing the sampling probability for each dataset, and the constraint $\sum_{i=1}^N \alpha_i = 1$ ensures the weights form a valid probability distribution.

What this means operationally: During training, when the data loader samples the next batch, it first selects dataset $\mathcal{D}_i$ with probability $\alpha_i$ , then draws a random example from that dataset. The weights $\alpha_i$ thus control the expected proportion of training tokens coming from each source. The optimization problem is to find the $\alpha_i$ vector that maximizes the trained model's average performance across a suite of evaluation benchmarks spanning general language understanding, mathematical reasoning, and code generation.

Why this formulation matters: The mixture weights are continuous-valued and live in a $(N-1)$ -dimensional simplex (the set of all probability distributions over $N$ categories). For $N = 7$ candidate datasets, this is a 6-dimensional continuous space—far too large for grid search or random sampling with any reasonable density. This is why automated search with a regression predictor is necessary, and why the fidelity of the proxy evaluations that train that predictor is so critical.

Component Model Training Protocol. The component models are the trained models that will later be merged to create proxies. Their training follows a specific protocol designed to ensure they can serve as effective building blocks.

Step 1: Base Model Training. A shared base model $\Theta_{\text{base}}$ is trained from scratch on 50 billion tokens of general-purpose data. This model uses the Qwen3-1.7B architecture (1.7 billion parameters, the specific architecture details are from Yang et al., 2025) and is trained with: a global batch size of 512 sequences, a sequence length of 8192 tokens ( $512 \times 8192 = 4,194,304$ tokens per optimizer step), an initial learning rate of $3 \times 10^{-4}$ decayed via a cosine schedule that retains a minimum of 20% of the initial learning rate (so the learning rate decays from $3 \times 10^{-4}$ to $6 \times 10^{-5}$ over the course of training). The base model provides foundational language capabilities—understanding syntax, semantics, basic reasoning, and factual knowledge from the general corpus—that all component models will share.

Step 2: Component Model Training. For each of the $N = 7$ candidate datasets $\mathcal{D}_i$ , a separate component model $\Theta_i$ is initialized from the shared base model $\Theta_{\text{base}}$ and further trained on a mixture consisting of 50% general data and 50% the candidate dataset $\mathcal{D}_i$ . The training uses the same hyperparameters as the base model training (batch size 512, sequence length 8192, learning rate $3 \times 10^{-4}$ with cosine decay to 20%, global batch size 512, sequence length 8192). The amount of additional training varies across experiments: the paper evaluates component models trained on 2B, 10B, 30B, and 50B tokens (Table 2), with the 30B-token configuration identified as the sweet spot for balancing proxy accuracy against training cost.

The purpose of the 50% general data admixture is to prevent the component models from catastrophically forgetting their general language capabilities during domain-specific training. A model trained exclusively on mathematical content would rapidly lose its ability to process natural language instructions, generate coherent text, or reason about non-mathematical concepts—making it useless as a building block for merging, since the merged model would inherit this degradation in proportion to the math component's weight. By maintaining 50% general data during component training, each component model stays anchored to the base model's general capabilities while still developing domain-specific specialization. This is the same regularization principle that constrains the final mixture search space to at most 50% non-general data.

The paper reports the magnitude of parameter updates during component model training through the metric $\delta$ (Equation 4):

$\delta = \frac{\sum |T(\mathcal{D}, \Theta_{\text{base}}) - \Theta_{\text{base}}|}{\sum |T(\mathcal{D}, \Theta_{\text{base}})| + \sum |\Theta_{\text{base}}|} \ll 1$

where the sums are over all parameter elements (absolute values of weights), $T(\mathcal{D}, \Theta_{\text{base}})$ is the component model parameters after training on dataset $\mathcal{D}$ , and the numerator is the total absolute parameter change while the denominator is approximately twice the total absolute parameter magnitude (since the trained parameters and base parameters are similar in scale). As reported in Appendix D, Table 12, $\delta$ grows with training budget: approximately 3.10% at 2B tokens, 6.90% at 10B tokens, 10.10% at 30B tokens, and 10.50% at 50B tokens. The paper's experiments primarily use the 30B-token configuration ( $\delta \approx 10.1\%$ ), which the authors claim satisfies the small-update assumption necessary for the model merging approximation to hold.

Model Merging as Proxy: The Core Mechanism

This is the central technical contribution of the paper. The mechanism translates the abstract data mixing objective into a concrete computational procedure that evaluates any mixture ratio at near-zero marginal training cost.

Weight Deltas and the Additivity Approximation. The formal foundation begins with defining the weight delta—the vector difference between a trained model's parameters and its initialization. For any dataset $\mathcal{D}$ and base model $\Theta_{\text{base}}$ :

$\Delta(\mathcal{D}) \triangleq T(\mathcal{D}, \Theta_{\text{base}}) - \Theta_{\text{base}}$

where $T(\mathcal{D}, \Theta_{\text{base}})$ is the training operator that returns the model parameters after training on dataset $\mathcal{D}$ starting from initialization $\Theta_{\text{base}}$ , and the subtraction is element-wise across all parameter tensors. For a component model trained on candidate dataset $\mathcal{D}_i$ , its parameters are $\Theta_i = \Theta_{\text{base}} + \Delta(\mathcal{D}_i)$ —the base parameters plus the domain-specific update.

The critical approximation that makes DeMix possible is the additivity of weight deltas (Equation 5):

$\Delta(\mathcal{D}_i \cup \mathcal{D}_j) \approx \Delta(\mathcal{D}_i) + \Delta(\mathcal{D}_j)$

What this equation claims: If you train a model on the union of datasets $i$ and $j$ , the resulting parameter changes (relative to the base model) are approximately equal to the sum of the parameter changes you would get from training on each dataset separately. In other words, the effect of training on multiple datasets simultaneously is roughly the sum of the individual effects.

What this enables: If the approximation holds, then the model trained on a weighted mixture of datasets $\mathcal{D}_{\text{mix}} = \sum_{i=1}^N \alpha_i \mathcal{D}_i$ has parameters approximately equal to the weighted average of the individually trained component models:

$\Theta_{\text{mix}} = T(\mathcal{D}_{\text{mix}}, \Theta_{\text{base}}) \approx \Theta_{\text{base}} + \sum_{i=1}^N \alpha_i \Delta(\mathcal{D}_i) = \sum_{i=1}^N \alpha_i \Theta_i$

The last equality uses the identity $\Theta_i = \Theta_{\text{base}} + \Delta(\mathcal{D}_i)$ and the constraint $\sum \alpha_i = 1$ to simplify: $\Theta_{\text{base}} + \sum \alpha_i \Delta(\mathcal{D}_i) = \Theta_{\text{base}} + \sum \alpha_i(\Theta_i - \Theta_{\text{base}}) = (1 - \sum \alpha_i)\Theta_{\text{base}} + \sum \alpha_i \Theta_i = \sum \alpha_i \Theta_i$ .

Why this form is powerful: It means that to evaluate what a model trained on mixture $\alpha$ would look like, you don't need to actually train it—you just take a weighted average of the already-trained component models. The cost of evaluating a new mixture ratio is reduced from "train a model on tens of billions of tokens" to "perform one weighted sum of parameter vectors and run inference on benchmarks." This is what the paper means by "decoupling search from training."

The Small-Update Condition (Equation 4). The additivity approximation is only valid when parameter updates are small relative to the initialization scale. The paper formalizes this through $\delta$ , defined above, and states that $\delta \ll 1$ is the necessary condition. In their experiments, $\delta \approx 10\%$ for the 30B-token component models, which the authors consider sufficient based on prior empirical work (Qin et al., 2022; Wu et al., 2025; Lin et al., 2025b) showing that weight delta additivity holds for update magnitudes in this range. Intuitively, when updates are small, the training dynamics are approximately linear—the gradient of the loss with respect to parameters doesn't change much during training, so training on dataset A then dataset B (or vice versa) produces similar parameter trajectories to training on both simultaneously. When updates are large, non-linear interactions between datasets become significant, and the approximation breaks down.

The Merging Equation (Equation 6). For any sampled mixture ratio $\{\alpha_i^j\}_{i=1}^N$ (where $j$ indexes the sampled mixture, and $\alpha_i^j$ is the weight for dataset $i$ in mixture $j$ ), the merged proxy model $M_{\text{mix}}^j$ is constructed as:

$M_{\text{mix}}^j = \sum_{i=1}^{N} \alpha_i^j \Theta_i$

where $\Theta_i$ is the component model for candidate dataset $\mathcal{D}_i$ , and the sum is element-wise over all parameter tensors (weights, biases, layer normalization parameters, embedding matrices—every trainable parameter in the model). The constraint $\sum_i \alpha_i^j = 1$ ensures the merged model remains in the convex hull of the component models, which is important because interpolating between models tends to produce more coherent behavior than extrapolating outside the convex hull (a phenomenon well-documented in the model merging literature, e.g., Wortsman et al., 2022a; Ilharco et al., 2022).

What this operation physically computes: Each parameter in the merged model is a weighted average of the corresponding parameters in the component models. For example, the weight matrix of the first attention layer in the merged model is $\alpha_1^j$ times the weight matrix from component model 1, plus $\alpha_2^j$ times the weight matrix from component model 2, and so on. This is a purely arithmetic operation—no gradients, no forward passes, no optimization—and can be performed in seconds on a single GPU for a 1.7B parameter model.

What this enables downstream: The merged model $M_{\text{mix}}^j$ serves as a training-free proxy for what a model trained from scratch on the actual data mixture $\mathcal{D}_{\text{mix}}^j = \sum_i \alpha_i^j \mathcal{D}_i$ would look like. The proxy can be evaluated on any benchmark by running standard inference, producing scores that approximate—with some degradation—what the real trained model would achieve. The quality of this approximation is measured by the proxy accuracy metrics (Spearman's $\rho$ , capability recovery) in Table 2.

Empirical Validation of the Merging Approximation (Appendix D, Table 12). The paper provides direct evidence for the additivity approximation by comparing merged proxy models against models trained on genuine mixed data. Using one math dataset $\mathcal{D}_1$ and one code dataset $\mathcal{D}_2$ , they consider several mixture proportions (20%/80%, 40%/60%, 60%/40%, 80%/20%) and compare two approaches:

Model merge: Train two component models separately on $\mathcal{D}_1$ and $\mathcal{D}_2$ , then merge their parameters according to the target mixture ratio.
Data mix: Train a single model directly on the mixed dataset $\mathcal{D}_1 + \mathcal{D}_2$ with the target mixture ratio.

The consistency score is defined as the ratio of the merged model's benchmark score to the data-mix model's benchmark score, averaged across all tested mixture proportions. A consistency score of 1.0 would mean perfect match; values above 1.0 indicate the merged model outperforms the data-mix model (which can happen due to the regularization effect of separate training), and values below 1.0 indicate the merged model underperforms (likely because it misses cross-domain synergies present in joint training).

The results in Table 12 show:

Token Budget	$\delta$	Math Consistency	Code Consistency
2B	3.10%	1.04	0.97
10B	6.90%	0.96	0.81
30B	10.10%	0.82	0.75
50B	10.50%	0.79	0.75

As $\delta$ grows, consistency degrades—the merged proxy becomes a less accurate approximation of the data-mix model. However, even at $\delta \approx 10.5\%$ (50B tokens), the consistency remains at 0.75–0.79, meaning the merged model retains about three-quarters of the relative performance signal of the genuine mixed-data model. For the 30B-token configuration ( $\delta \approx 10.1\%$ ) that the paper selects as the primary operating point, consistency is 0.82 for math and 0.75 for code. This degradation is acceptable because the goal is ranking mixtures, not exactly recovering their absolute performance—as long as the ranking of merged proxies correlates strongly with the ranking of real trained models (which the Spearman's $\rho$ of 0.81 confirms), the exact absolute performance level is less important.

Why this validation is necessary: If the additivity approximation did not hold in practice—if merged models produced essentially random rankings relative to real trained models—then DeMix would be useless regardless of its computational efficiency. The validation in Table 12 provides the empirical grounding that the theoretical approximation (Equation 5) translates into practically useful proxy signals.

The Cost Structure. The key efficiency property of DeMix is visible in Table 1 and Table 2. The pre-cost of training component models is fixed: for the 30B-token configuration with 7 component models, the total training budget is $7 \times 30\text{B} + 50\text{B} \text{ (base model)} = 260\text{B}$ tokens (though Table 2 reports 212B for DeMix with 30B components—the discrepancy likely reflects that the base model cost is amortized or reported differently). The marginal cost per additional proxy is only the benchmarking cost, which Table 1 quantifies as equivalent to training 0.013B tokens. So evaluating 112 proxies costs $7 \times 30\text{B} + 112 \times 0.013\text{B} \approx 211\text{B}$ total tokens, while evaluating 1,000 proxies would cost $7 \times 30\text{B} + 1,000 \times 0.013\text{B} \approx 223\text{B}$ —a negligible increase. In contrast, training-based methods scale linearly: 1,000 proxies at 2B tokens each would cost 2,000B tokens, making such extensive exploration completely infeasible.

The Merging Method Choice. Table 4 compares six merging methods on proxy accuracy, using component models trained at 30B, 40B, and 50B tokens (the reported $\rho$ and capability recovery are macro-averages across these three configurations and 96 mixture ratios). The methods compared are:

Linear (Wortsman et al., 2022a): Simple weighted averaging of all parameters as in Equation 6. No hyperparameters, no pruning, no post-processing.
Multi-SLERP (Goddard et al., 2024): Spherical linear interpolation extended to multiple models, which interpolates along the geodesic on the hypersphere rather than the straight line in Euclidean space. This preserves the angular relationship between parameter vectors.
Breadcrumbs (Davari & Belilovsky, 2024): Uses sparse masks to select which parameters to merge, reducing interference between models.
DARE (Yu et al., 2024): Randomly drops a fraction of weight delta elements and rescales the remaining ones, which has a hyperparameter (the drop rate) that requires tuning.
DELLA (Deep et al., 2024): Magnitude-based sampling of weight deltas, which also has hyperparameters controlling the sampling strategy.
TIES (Yadav et al., 2023): Trim, Elect Sign, and Merge—a three-step procedure that prunes small-magnitude changes, resolves sign conflicts between models, and then merges only the agreed-upon directions. This has hyperparameters for the trimming threshold.

The results show Linear merging achieves the highest capability recovery (0.845) and the second-highest macro-average $\rho$ (0.787, essentially tied with Multi-SLERP at 0.785 and DELLA at 0.784). The paper's choice of Linear merging is justified on three grounds: (1) it is hyperparameter-free, eliminating the need to tune pruning thresholds or drop rates that could affect downstream mixture optimization, (2) it is computationally trivial—a single weighted sum per parameter—making it suitable for generating hundreds of proxies, and (3) its performance is empirically competitive with or superior to more complex methods. The slight edge of Linear over Multi-SLERP in capability recovery suggests that for this specific setting (component models with small weight deltas from a shared base), the straight-line interpolation in Euclidean space actually preserves functional behavior better than spherical interpolation.

Iterative Mixture Weight Optimization

With the proxy mechanism established, the remaining challenge is to use merged proxy evaluations to find the optimal mixture ratio. This is a regression-guided search problem: a predictor is trained to map mixture ratios to performance, and that predictor is used to explore the mixture space efficiently.

The Evaluation Metric. Rather than using raw benchmark accuracy as the optimization target—which would require the predictor to model absolute performance levels that vary dramatically across benchmarks (GSM8K accuracy ranges from ~5% to ~25% in these experiments, while HellaSwag ranges from ~50% to ~55%)—DeMix uses average ranking across benchmarks. For each merged proxy model $M_{\text{mix}}^j$ , it is evaluated on a suite of nine benchmarks:

General: ARC-E (Clark et al., 2018), HellaSwag (Zellers et al., 2019), WinoGrande (Sakaguchi et al., 2021), PIQA (Bisk et al., 2020), SIQA (Sap et al., 2019)
Code: HumanEval (Chen et al., 2021), MBPP (Austin et al., 2021)
Math: GSM8K (Cobbe et al., 2021b), MATH (Hendrycks et al., 2021)

For each benchmark, the model receives a score (accuracy for classification tasks, pass@1 for code generation). Within each benchmark, the 96 reference models that were trained on real data mixtures at 50B tokens are ranked, and the merged proxy's score is mapped to its rank in this reference distribution. The final metric $r_j$ for mixture $j$ is the macro-average of its ranks across the general, code, and math benchmark groups—that is, the average of the average rank on general benchmarks, the average rank on code benchmarks, and the average rank on math benchmarks. This domain-level macro-averaging ensures that no single domain dominates the optimization: a mixture that excels at math but collapses on general language would receive a poor average rank despite strong math performance.

Why ranking over raw scores: Rankings are scale-invariant and robust to the systematic performance degradation introduced by model merging. Even if merged proxies consistently underperform real trained models by some margin (as the capability recovery of ~0.83 indicates), their relative ordering should be preserved if the merging approximation is faithful. Rankings also naturally handle the different dynamic ranges of different benchmarks—the difference between 50% and 55% on HellaSwag might represent a similar "improvement" as the difference between 5% and 15% on MATH in terms of what it says about the model's capabilities, and ranking captures this by normalizing to the reference distribution.

Iterative Predictor-Guided Search. The optimization proceeds through three iterations of the following cycle, with the number of sampled mixtures decreasing at each iteration:

Iteration 1 (64 samples):

Random sampling: 64 mixture weight vectors $\{\alpha_i^j\}$ are sampled uniformly from the 6-dimensional simplex—each vector is a point in the simplex where all 7 weights are non-negative and sum to 1. Uniform sampling from a simplex is achieved by drawing 7 independent exponential random variables and normalizing them to sum to 1 (this produces a Dirichlet(1,1,...,1) distribution, which is uniform over the simplex).
Proxy construction: For each sampled ratio $j$ , a merged proxy $M_{\text{mix}}^j$ is constructed via Equation 6 and evaluated on the benchmark suite to obtain its average rank $r_j$ .
Predictor training: A LightGBM regression model $f: \mathbb{R}^7 \to \mathbb{R}$ is trained on the 64 pairs $(\alpha^j, r_j)$ to predict ranking score from mixture weights. LightGBM (Ke et al., 2017) is a gradient-boosted decision tree algorithm chosen for its efficiency with small-to-medium datasets and its ability to capture non-linear interactions between features without extensive hyperparameter tuning. The hyperparameters are: learning rate 0.02, number of iterations (boosting rounds) 300. The input features are the 7 mixture weights; the target is the average rank.
Candidate scoring: The trained predictor $f$ is applied to a large set of newly sampled mixture ratios (the paper does not specify exactly how many, but "a large number" consistent with the goal of thorough exploration). The predictor scores each candidate, and the top-ranked candidates (those with the lowest predicted rank, since lower rank means better performance) are selected.

Iteration 2 (32 samples): The top-ranked candidates from Iteration 1 are used as the sampling pool. 32 new mixture ratios are drawn—likely by perturbing the top candidates or sampling from their vicinity, though the exact perturbation strategy is not detailed. These 32 ratios go through the same proxy construction and evaluation process, yielding 32 new (ratio, rank) pairs. The predictor is retrained on the combined dataset of $64 + 32 = 96$ pairs, incorporating both the initial random exploration and the focused samples from promising regions.

Iteration 3 (16 samples): The same procedure repeats: the top-ranked candidates from Iteration 2 inform sampling of 16 new ratios, which are evaluated and added to the training set, yielding a final predictor trained on $64 + 32 + 16 = 112$ total proxy evaluations.

The Logic of Iterative Refinement. The decreasing sample sizes (64 → 32 → 16) reflect a shift from exploration to exploitation. The first iteration casts a wide net, randomly sampling across the entire simplex to give the predictor a rough global picture of the mixture-to-performance landscape. The second and third iterations focus sampling on the high-performing regions identified by the predictor, improving its accuracy where it matters most—near the optimum. This is directly inspired by CLIMB's iterative sampling strategy (Diao et al., 2025), but applied to merged proxies rather than training-based proxies.

Why iterative refinement helps: The mixture-to-performance function is likely non-linear (the contribution of math data may depend on how much code data is present, due to cross-domain transfer) and may have multiple local optima. A predictor trained only on random samples may be inaccurate near the global optimum if that region was sparsely sampled in the initial random draw. By iteratively focusing samples on high-predicted-performance regions and retraining the predictor with the new data, the predictor's accuracy in the region that matters—the top of the ranking—improves with each iteration. This is the same principle behind Bayesian optimization with expected improvement acquisition functions, though DeMix uses a simpler resampling strategy guided by the current predictor.

Final Mixture Selection. After the third iteration, the trained predictor $f$ is applied to score a large number of newly sampled mixture ratios. The top 128 candidates by predicted rank are selected, and the final optimal mixture ratio is computed as the arithmetic mean of these 128 candidates (element-wise averaging of the 7-dimensional weight vectors, then renormalizing to sum to 1).

Why average over top candidates rather than selecting the single best: Individual mixture ratios may have noisy predictor scores, especially if they lie in regions of the simplex that were not densely sampled during the iterative process. Averaging over the top 128 candidates provides a form of bootstrap aggregation: the mean of multiple high-scoring candidates is more stable and less sensitive to predictor errors than any single candidate. This is also a conservative choice—if there is a broad plateau of near-optimal mixtures, the average will lie near the center of that plateau, which is a safer bet than gambling on a potentially spurious peak at the edge. If the true optimum is a sharp peak, averaging will miss it, but sharp peaks in mixture space are unlikely given the smooth, gradual way that changing data proportions typically affects model capabilities.

The Total Proxy Budget and Its Allocation. The total number of proxy evaluations is 112 (64 + 32 + 16), matching the default used in CLIMB's experiments. However, the paper also experiments with other proxy counts in Table 3: 56 proxies (presumably allocated as 32 + 16 + 8), 224 proxies (128 + 64 + 32), and 448 proxies (256 + 128 + 64). The results show that scaling from 56 to 112 to 224 proxies improves the final mixture quality (rank improves from 29.33 to 25.67 to 24.00), but further scaling to 448 proxies degrades performance (rank drops to 27.67). The paper attributes this to overfitting: with 448 proxy evaluations, the LightGBM predictor with fixed hyperparameters (learning rate 0.02, 300 iterations) may start fitting noise in the proxy evaluations rather than the true mixture-to-performance signal. This highlights that while DeMix removes the training cost barrier to scaling proxy count, the finite sample size for predictor training still imposes limits.

The Computational Cost Breakdown (Table 1). A single benchmarking run—evaluating one merged proxy on all 9 benchmarks—consumes 0.3 GPU-hours on H800 GPUs. The paper maps this to an equivalent training cost: 0.3 GPU-hours is roughly the cost of training on 0.013B tokens (13 million tokens) with the Qwen3-1.7B architecture. So 112 proxy evaluations cost the equivalent of $112 \times 0.013\text{B} \approx 1.5\text{B}$ training tokens in benchmarking—negligible compared to the 210B+ tokens invested in component model training. The component model training cost of approximately 211B tokens (for 7 components at 30B each, with the base model cost presumably included in this figure based on Table 2's reporting) is the dominant cost, but it is a one-time investment that enables unlimited subsequent proxy evaluations.

Design Choices and Their Justifications

Each design decision in the DeMix pipeline represents a tradeoff, and understanding why specific choices were made reveals the engineering principles that make the system work.

Why 50% General Data in Candidate Datasets (Table 5). The ablation in Table 5 is decisive: reducing the general data proportion from 50% to 25% causes the Spearman's $\rho$ to drop from 0.787 to 0.667, and capability recovery falls from 0.845 to 0.796. Eliminating general data entirely (0%) further reduces $\rho$ to 0.652 and capability recovery to 0.795. The mechanism is straightforward: without sufficient general data during component training, the component models lose their general language capabilities, becoming narrowly specialized. When these narrow models are merged, the resulting proxy inherits their domain-specific strengths but also their general-capability deficiencies, in proportion to the mixture weights. Since the benchmarks include general language tasks, this deficiency depresses scores in ways that are not purely a function of the mixture ratio—the proxy is penalized for the component training artifact rather than the mixture effect. The 50% proportion is likely chosen as the maximum that still preserves meaningful domain specialization: above 50%, the general data overwhelms the domain signal; below, the capability loss corrupts the proxy signal.

Why Linear Merging Over Alternatives (Table 4). As discussed above, Linear merging achieves the best capability recovery (0.845) and competitive $\rho$ (0.787) while being hyperparameter-free. The hyperparameter-free property is important for two reasons. First, it eliminates a source of variance: if the merging method had tunable hyperparameters, the optimal settings might depend on the specific mixture ratio being evaluated, creating a circular problem where you need to know the answer to set up the evaluation. Second, it makes the method reproducible and transferable—other researchers can apply DeMix to their own candidate datasets without needing to tune merging hyperparameters.

The superiority of Linear over more sophisticated methods in this setting is explained by the small-update regime. Methods like TIES and DARE are designed to resolve interference when merging models that have undergone large, divergent updates (e.g., fine-tuning on entirely different tasks from a pre-trained checkpoint). In DeMix's setting, the updates are small ( $\delta \approx 10\%$ ) and the component models are trained on related tasks (all are language modeling, just with different data distributions). In this regime, the simple arithmetic mean already produces coherent models, and the additional pruning or sign-resolution steps in TIES/DARE may actually remove useful signal.

Why Iterative Resampling (Section 2.4, Step 4). The iterative process serves two functions simultaneously. First, it biases the training data for the predictor toward the high-performing region of mixture space, improving predictor accuracy near the optimum (exploitation). Second, by retaining the random samples from early iterations, it prevents the predictor from becoming overconfident in a locally optimal region that was discovered early but is not actually global (exploration). The decreasing sample sizes (64, 32, 16) reflect a heuristic for balancing these: the first iteration provides broad coverage, and subsequent iterations refine without wasting too many samples on regions already known to be suboptimal.

The paper does not compare against a non-iterative baseline (e.g., 112 random samples evaluated once and a single predictor trained), so the marginal benefit of iterative refinement over simple random sampling is not directly quantified. However, given that DeMix's proxy construction is nearly free, the cost of iteration is only the additional benchmarking—negligible relative to component training—so the iterative approach is a "free lunch" in terms of total budget.

Why Averaging Top 128 Candidates Rather Than Selecting the Best Single Candidate. As discussed, this is a stability choice. With 7-dimensional mixture ratios and only 112 training points for the predictor, the predictor's estimates for any individual candidate have non-trivial variance. Taking the centroid of many top candidates provides a mixture that is likely near-optimal even if no individual candidate in the top-128 set is exactly optimal. This is analogous to the practice in hyperparameter optimization of selecting the best configuration from a set of candidates rather than trusting the surrogate model's single-point prediction.

The choice of 128 candidates (rather than, say, 10 or 1,000) is not ablated, but it reflects a balance: too few candidates and the averaging doesn't provide much stabilization; too many and the average includes candidates from suboptimal regions, pulling the final mixture away from the optimum.

Why Difficulty Bins Aren't Used Here (Contrast with the Reference Example). Unlike the reference example paper (which optimized test-time compute allocation conditioned on question difficulty), DeMix optimizes a single global data mixture ratio. This is the appropriate formulation for pre-training data mixing because the mixture determines the entire training data distribution—it cannot be conditioned on individual training examples in the same way that inference strategies can be conditioned on individual prompts. The search is over a single fixed ratio that applies to the entire training corpus.

Why OpenCompass for Evaluation (Section 3). The paper uses OpenCompass (OpenCompass Contributors, 2023) as the evaluation framework. This is a standardized benchmark suite widely used in the Chinese LLM community, providing consistent implementations of the evaluation protocols for each benchmark (prompt templates, few-shot settings, answer extraction, grading). Using a standardized framework improves reproducibility and makes DeMix's reported scores directly comparable to other work using the same evaluation infrastructure.

Why the Specific Model Architecture (Qwen3-1.7B). The Qwen3 architecture (Yang et al., 2025) is a modern decoder-only transformer that is representative of current LLM design. The 1.7B parameter scale is chosen as a practical balance: large enough to exhibit non-trivial math and code capabilities after 50B tokens of training (enabling meaningful benchmark signals), but small enough that training multiple component models at 30B–50B tokens each remains feasible within academic or small-industry compute budgets. The choice also enables the scalability experiment on Qwen3-8B (Section 4.3.3, Table 6), which tests whether mixtures optimized at 1.7B scale transfer to larger models—an important practical question since industrial pre-training runs are typically at much larger scales than the proxy experiments.

Why Three-Stage Pre-Training for the DeMix Corpora (Appendix A.2). The DeMix Corpora is organized into three stages with different data mixtures, reflecting the industry practice of dynamic data scheduling during pre-training (Feng et al., 2024; Basant et al., 2025). In Stage 1 (~14T tokens), the focus is on data diversity and broad general knowledge acquisition, with a simpler mixture optimization based on quality upsampling/downsampling rather than the full DeMix framework. In Stages 2 and 3 (~6T and ~2T tokens respectively), the proportion of high-quality math and code data increases substantially, and DeMix is applied to optimize the mixture ratios. This staged approach acknowledges that the optimal data mixture is not static throughout pre-training—early training benefits from diverse exposure to build foundational representations, while later training benefits from focused high-quality data to refine specific capabilities. The DeMix framework is specifically applied in the later stages where data quality critically impacts final performance and where the mixture optimization problem is most challenging due to the inclusion of math and code domains alongside general data.

4. Key Insights and Innovations

Innovation 1: The Decoupling Paradigm — Why "Search" and "Training" Should Be Separate Resources

The paper's most fundamental conceptual move is not any specific algorithm or architecture, but rather the redefinition of data mixture optimization as two separable problems: model training (which produces capability) and mixture evaluation (which consumes capability to produce ranking signals). Prior to DeMix, the field implicitly treated these as a single coupled process—to evaluate a new mixture ratio, you trained a new model (or at least continued training an existing one). RegMix, CLIMB, and the manual proxy approach used by industrial labs all operate within this paradigm: every step of mixture exploration requires a commensurate step of model training expenditure.

DeMix argues that this coupling is artificial and wasteful. Training produces model parameters; mixture evaluation requires only that those parameters reflect the target data distribution faithfully enough to rank mixtures correctly. If you can construct parameters that approximate the result of training on a given mixture without actually performing that training, you break the coupling entirely. The paper identifies a specific condition—small parameter updates from a shared initialization—under which weighted linear merging satisfies this requirement, but the deeper insight is the architectural separation itself.

This is a fundamental reframing, not an incremental improvement. It's analogous to the shift in optimization from "evaluate the objective by running the full simulation" to "build a surrogate model and evaluate that instead." Just as Bayesian optimization decouples function evaluation from the expensive process that generates function values, DeMix decouples mixture evaluation from the expensive process (training) that generates model capabilities. The conceptual payoff is that mixture search becomes essentially free at the margin—you can evaluate 1,000 candidate mixtures for roughly the same cost as 100, once the component models exist. Table 2 makes this concrete: DeMix with 30B-token components achieves a Spearman's ρ of 0.81 using 212B total tokens, while training-based methods need 1,344B tokens (6.4× more) to reach comparable proxy accuracy (0.82). But more revealing than the 6.4× multiplier is what happens at the low-budget extreme: training-based proxies at 224B total tokens achieve a ρ of only 0.53. DeMix doesn't just scale better—it fundamentally alters the shape of the cost-accuracy curve, making high accuracy achievable at budgets where training-based methods are still producing near-random rankings.

The decoupling also unlocks a practical capability that training-based methods structurally cannot provide: post-hoc re-analysis. Once the component models are trained, you can evaluate any mixture ratio you can imagine—not just the ones you thought to sample during the initial search. If a new benchmark becomes important after the mixture is selected, or if you want to understand the sensitivity of performance to specific dataset proportions, you can generate merged proxies for an entire grid of ratios and analyze the landscape retrospectively. This turns mixture optimization from a one-shot decision into an interactive exploration, which is impossible when each evaluation requires weeks of training.

Innovation 2: The Small-Update Regime as an Enabling Condition, Not a Limitation

The paper's second distinctive contribution is the identification, quantification, and operationalization of the small-update regime as the condition that makes the decoupling possible. This is significant not because the observation that weight deltas are approximately additive when small is novel—Qin et al. (2022), Wu et al. (2025), and Lin et al. (2025b) had already established this—but because DeMix shows that this regime is practically achievable and useful in the specific context that matters most for data mixture optimization: late-stage pre-training where math and code capabilities are being developed.

The diagnostic move is the paper's operationalization of the update magnitude through δ (Equation 4) and its empirical mapping to proxy fidelity (Table 12). At δ ≈ 3% (2B token component training), the proxy is nearly perfect—math consistency of 1.04 suggests the merged model actually slightly outperforms the data-mix model, likely due to the regularizing effect of separate training. At δ ≈ 10% (30B–50B token training, which the paper selects as the operating point), consistency degrades to 0.75–0.82, which is sufficient for ranking but indicates meaningful deviation from the true data-mix parameters. At even larger δ, the approximation would presumably break down entirely, though the paper does not test this boundary.

What makes this an insight rather than just an empirical observation is the counterintuitive implication for proxy design. The natural instinct in proxy-based optimization is to make proxies as accurate as possible—train them on more tokens, use larger models, make them as close to the real thing as you can afford. DeMix reveals a tension: making component models more accurate (by training them longer, increasing δ) actually degrades the merging approximation, because larger updates violate the linearity assumption. The optimal operating point is not "as much training as possible" but rather "enough training to develop the target capabilities while keeping updates small enough for linear merging to remain faithful." This is a genuinely novel design principle that has no analogue in training-based proxy methods, where more training is monotonically better (if you can afford it).

The paper doesn't fully explore this tension—it selects 30B tokens based on the efficiency-accuracy tradeoff in Table 2, but doesn't systematically vary δ to find the precise optimum. However, the very existence of this tension, and the paper's explicit quantification of it, represents a new way of thinking about proxy design. It suggests that future work on model merging for data selection should treat δ as a hyperparameter to be optimized, balancing component model capability against merging fidelity, rather than simply maximizing component model performance.

Innovation 3: Domain-Balanced Evaluation as a Constraint on Mixture Optimization

DeMix's third conceptual contribution is subtler but equally important: the recognition that multi-domain mixture optimization requires a ranking-based evaluation metric that prevents any single domain from dominating the objective. This sounds like an implementation detail, but it reflects a deeper insight about what "optimal" means when you're optimizing for capabilities that lie on fundamentally different scales.

Consider the alternative: if you simply maximize average accuracy across all benchmarks, a mixture that produces 55% on HellaSwag and 5% on MATH would score 30% on average. A mixture that produces 50% on HellaSwag and 15% on MATH would score 32.5%. Both are plausible outcomes, but the former model is essentially incapable of mathematical reasoning while the latter has merely slightly degraded language understanding. Which is "better" depends entirely on your deployment priorities, and there is no objective scale on which a 5-point HellaSwag drop equals a 10-point MATH gain. The raw-accuracy optimization implicitly imposes an exchange rate determined by the variance and dynamic range of each benchmark, which is an artifact of benchmark design rather than a meaningful preference.

DeMix's solution—macro-averaged ranking within each domain (general, code, math), then averaging those domain-level ranks—is a specific implementation of a broader principle: optimize for balanced capability improvement rather than aggregate score maximization. By converting raw scores to ranks within a reference distribution of 96 trained models, the metric normalizes away scale differences between benchmarks. By macro-averaging across domains before computing the final rank, it ensures that improving math from terrible to mediocre is treated as comparable to improving general language from good to slightly better—both are equally weighted in the final objective, regardless of the absolute score differences involved.

This is a design choice with philosophical implications. It embeds a normative preference into the optimization: DeMix assumes that balanced multi-domain capability is inherently valuable, and it optimizes accordingly. A team that genuinely preferred a pure language model with zero math capability would want a different objective. But for the use case the paper targets—general-purpose LLMs that must perform adequately across all domains—this balanced objective is precisely aligned with practical requirements, and the paper's explicit formulation of it as a ranking-based, domain-macro-averaged metric makes the preference transparent and reproducible.

The evidence that this matters is in Table 3. The uniform mixture (equal proportions of all candidate datasets) achieves rank 36.67 averaged across domains—broken down, it's strong on general benchmarks (rank 9) but terrible on math (rank 57). DeMix's optimized mixture (224 proxies) improves the average rank to 24.00 by dramatically improving math (rank 14) and code (rank 12) while accepting a moderate decline in general performance (rank 46). The tradeoff is explicit and quantifiable: DeMix sacrifices some general language capability to gain much more math and code capability, exactly as the balanced objective intends. Whether this tradeoff is "correct" depends on the application, but the framework makes the tradeoff visible and controllable, which is a substantial advance over methods that collapse everything into a single loss number.

Innovation 4: Verifier Over-Optimization as a First-Class Phenomenon in Test-Time Scaling

Note: This section header appears to be an artifact from the reference example. Based on the DeMix paper's actual content, there is no fourth fundamental innovation of comparable depth to the three above. The paper's remaining contributions—the DeMix Corpora release, the iterative predictor-guided search with LightGBM, the specific engineering choices (50% general data, Linear merging, etc.)—are important for reproducibility and practical adoption but are incremental refinements rather than conceptual innovations. The three innovations above capture the paper's distinctive intellectual contribution: the decoupling paradigm, the operationalization of the small-update regime as a design constraint, and the balanced multi-domain evaluation principle. I will not fabricate a fourth innovation to pad the count.

5. Experimental Analysis

Evaluation Methodology

Dataset. All mixture optimization experiments use data drawn from the DeMix Corpora, a 22T-token dataset constructed from heterogeneous open-source sources (web, math, code, multilingual) and cleaned through a pipeline of global deduplication, perplexity filtering, FastText quality classification, Chinese-specific quality filtering, and instance-level labeling (Appendix A, Table 10). The candidate datasets for mixture optimization consist of seven categories: General, Multilingual, and three quality-stratified tiers each for Math and Code (Table 11). For evaluation of final mixtures, models are trained on 50B tokens sampled according to the optimized mixture ratios and evaluated on nine downstream benchmarks drawn from standard LLM evaluation suites.
Base model(s). The primary experiments use the Qwen3-1.7B architecture (Yang et al., 2025), a 1.7B-parameter decoder-only transformer. The paper argues this scale is "large enough to exhibit non-trivial math and code capabilities after sufficient training" while remaining feasible for training multiple component models. The scalability experiment (Section 4.3.3) additionally uses Qwen3-8B to test whether mixtures optimized at 1.7B scale transfer to larger models. All component models are initialized from a shared base model trained from scratch on 50B tokens of general-purpose data.
Metrics. Two categories of metrics are used. Proxy Consistency metrics measure how well merged proxy models approximate real trained models: (a) Spearman's ρ — the rank correlation between proxy benchmark scores and reference model scores, computed across 96 randomly sampled mixture ratios, reported as macro-average across general, code, and math benchmarks plus a top-25% variant restricted to the highest-performing reference models; (b) Capability Recovery — the ratio of the average benchmark score of the merged proxy to that of the corresponding reference model, measuring absolute performance retention. Mixture Quality metrics evaluate the downstream model trained on the final optimized mixture: the model is ranked against the 96 reference models on each benchmark, and the macro-average of these ranks across general, code, and math domains is reported as the final rank metric (lower rank is better). All evaluations use the OpenCompass framework (OpenCompass Contributors, 2023).
Baselines. Four categories of baselines are compared. Uniform: equal allocation across all seven candidate datasets, serving as a naive non-optimized reference. Heuristic: manually tuned mixtures (two variants in Table 3, three variants in Table 11), representing the common practice of expert-guided ratio selection. RegMix (Liu et al., 2024): samples mixture ratios, trains small proxy models on each at fixed token budgets, evaluates them, and fits a LightGBM predictor to map ratios to performance—the non-iterative predecessor to CLIMB. CLIMB (Diao et al., 2025): extends RegMix with iterative resampling, where the predictor guides additional proxy training toward promising mixture regions. For both RegMix and CLIMB, the paper experiments with two proxy configurations: 112 proxies at 2B tokens each (total budget 224B) and 28 or 56 proxies at 8B tokens each (total budget 224B or 448B). Methods like DoReMi and Rho Loss are excluded because they optimize evaluation loss rather than downstream benchmark performance (Section 3.2).
Generation budget / compute accounting. Compute is measured in training tokens (billions of tokens processed during model training), with benchmarking cost converted to equivalent training tokens using GPU-hour parity. Table 1 establishes the conversion: one benchmarking run (evaluating one model on all nine benchmarks) costs 0.3 H800 GPU-hours, equivalent to training on 0.013B tokens. For DeMix, the total budget sums the component model training cost (e.g., 7 components × 30B tokens = 210B) plus the base model training (50B) plus benchmarking cost (number of proxies × 0.013B). For training-based baselines, the total budget is simply the number of proxies multiplied by the token budget per proxy. All comparisons are made at equal or comparable total token budgets, with DeMix's efficiency advantage coming from the fact that its marginal cost per additional proxy is near-zero (~0.013B tokens for benchmarking) versus 2–12B tokens for each additional training-based proxy.
Cross-validation / statistical protocol. For proxy accuracy measurement, 96 mixture ratios are randomly sampled from the simplex, and 96 corresponding reference models are trained on 50B tokens each to serve as the ground-truth performance standard. Spearman's ρ and capability recovery are computed by comparing merged proxy performance against these reference models. For the scalability experiment (Table 6), the same mixture ratios obtained from the 1.7B experiments are applied directly to Qwen3-8B training—no cross-validation across scales is performed; the transfer is evaluated directly. The paper does not report confidence intervals or standard errors for any metric, nor does it describe a formal statistical test for comparing methods. The two-fold cross-validation protocol mentioned in the reference example paper is not used here; instead, the iterative predictor-guided search serves as the mechanism for avoiding overfitting to specific sampled mixtures, with the averaging of top-128 candidates providing additional stabilization.

Main Quantitative Results

Proxy Consistency: DeMix versus Training-Based Proxies

The central quantitative claim of the paper is that merged proxies achieve substantially higher ranking consistency with reference models than training-based proxies at equivalent compute budgets, and reach comparable consistency at dramatically lower budgets. Table 2 provides the evidence, comparing DeMix at four component model training scales (2B, 10B, 30B, 50B tokens per component × 7 components) against training-based proxies at four per-proxy budgets (2B, 4B, 8B, 12B tokens × 112 proxies).

At the lowest budget tier, the contrast is stark. Training-based proxies at 2B tokens each (224B total, 112 proxies) achieve a macro-average Spearman's ρ of only 0.53 and a top-25% ρ of merely 0.17. The top-25% ρ is particularly revealing: it measures ranking consistency specifically among the best mixtures—exactly the ones you most need to identify accurately. A value of 0.17 indicates that the 2B-token proxies are essentially useless for distinguishing top-tier mixtures from each other or from mediocre ones. In contrast, DeMix with only 2B-token component models (15B total training cost) already achieves macro-average ρ of 0.55 and top-25% ρ of 0.27—slightly better at 15× lower cost. This configuration is not practically useful (the absolute ρ is still low), but it demonstrates that even minimal component training produces more informative proxies than training-based methods at comparable budgets.

At moderate budgets, DeMix pulls decisively ahead. With 10B-token components (71B total budget), DeMix achieves macro-average ρ of 0.60 and top-25% ρ of 0.41. Training-based proxies need 8B tokens each (896B total) to reach macro-average ρ of 0.71 and top-25% ρ of 0.47—roughly comparable overall but at 12.6× higher cost. For the top-25% metric specifically, DeMix at 71B (0.41) actually outperforms training-based at 448B (0.38).

At the primary operating point (30B-token components, 212B total budget), DeMix achieves macro-average ρ of 0.81 and top-25% ρ of 0.59. Training-based proxies require 12B tokens each (1,344B total) to reach macro-average ρ of 0.82 and top-25% ρ of 0.57—essentially identical accuracy at 6.4× higher cost (1,344B vs. 212B). The paper reports: "DeMix achieves substantially higher proxy accuracy than training-based proxy models under the same limited budget, and reach comparable accuracy with approximately 6× less computation budget (200 v.s. 1200)." The 6× figure specifically refers to the ~200B tokens for DeMix versus ~1,200B tokens for training-based proxies to reach ρ ≈ 0.80–0.82.

At higher budgets, DeMix's returns diminish. Scaling component models to 50B tokens (351B total) yields macro-average ρ of 0.80—actually slightly lower than the 30B configuration (0.81)—while top-25% ρ drops to 0.50 (from 0.59). This is consistent with the theoretical prediction: larger component training budgets increase δ (the parameter update magnitude reported in Appendix D, Table 12 as growing from 10.10% at 30B to 10.50% at 50B), which degrades the linearity assumption underlying the merging approximation. The paper does not test beyond 50B, so the exact δ at which the approximation becomes unusable is not characterized, but the 50B result already shows the beginning of degradation.

Capability recovery follows a similar pattern. DeMix at 30B recovers 83% of reference model performance on average (0.83), while training-based proxies at 12B recover 87% (0.87). This is the one metric where training-based proxies maintain an edge even at high budgets—directly trained models, even with smaller per-proxy budgets, preserve more absolute performance than merged proxies. But for ranking, which is what matters for mixture selection, the correlation metrics are the relevant ones, and DeMix matches or exceeds training-based methods there.

Domain-specific breakdown. Table 2 reports Spearman's ρ separately for general, code, and math benchmarks. The patterns reveal domain-specific difficulty in proxy construction. For training-based proxies at 2B tokens, the general benchmark ρ is near-zero (0.12), code ρ is moderate (0.60), and math ρ is strong (0.85). This counterintuitive pattern—math ranking being most accurate at the smallest scale—likely reflects that even 2B-token proxies can capture gross differences in math capability (does the model get 2% or 8% on GSM8K?), while general benchmark differences are subtler (does the model get 52% or 54% on HellaSwag?) and require more training to resolve. DeMix at 30B components achieves ρ of 0.64 (general), 0.81 (code), and 0.98 (math), showing that the merging approximation preserves ranking information best for math and code—likely because domain-specific capabilities are more directly tied to the presence of domain-specific data, making the linear additivity assumption more accurate for those domains than for general language where cross-domain interactions are more complex.

The finer partition experiment (Table 7) replicates the analysis with 15 candidate datasets instead of 7, testing whether DeMix's proxy accuracy holds under more granular data partitioning. The results are qualitatively consistent: training-based proxies at 12B per proxy (1,344B total) achieve macro-average ρ of 0.84 and top-25% ρ of 0.58; DeMix at 30B components (451B total) achieves 0.83 and 0.59 respectively—again nearly identical accuracy at roughly 3× lower cost. The consistency of this finding across both 7-category and 15-category partitions suggests that DeMix's advantage is not an artifact of the specific coarse categorization used in the main experiments, though the absolute ρ values are slightly lower across the board for the 15-category case, consistent with the increased difficulty of the higher-dimensional search problem.

Mixture Quality: Downstream Model Performance

The proxy consistency results establish that DeMix produces faithful ranking signals, but the ultimate test is whether those signals lead to better data mixtures when used for actual model training. Table 3 reports the benchmark performance and ranks of models trained on 50B tokens using the mixtures discovered by each method.

The headline result: DeMix with 224 merged proxies achieves the best average rank of 24.00 (lower is better, representing better performance relative to the 96 reference models), outperforming all baseline configurations. The rank is decomposed into domain-specific ranks: 46 (general), 12 (code), 14 (math), with the macro-average of 24.00.

Comparison against training-based methods at similar budgets:

DeMix (224 proxies, 212B total budget): rank 24.00
RegMix with 112 × 2B proxies (224B total): rank 38.00
CLIMB with 112 × 2B proxies (224B total): rank 34.67
RegMix with 28 × 8B proxies (224B total): rank 35.00
CLIMB with 28 × 8B proxies (224B total): rank 30.00

At the 224B budget tier, DeMix substantially outperforms all alternatives. The gap is particularly large against the 2B-token proxy configurations, consistent with the poor proxy accuracy (ρ ≈ 0.53) of those configurations—garbage-in, garbage-out. The 8B-token proxy configurations perform better but still trail DeMix, suggesting that even when per-proxy accuracy is improved, the limited number of proxies (28 vs. 112–224) constrains the search coverage.

Scaling the training budget for baselines: When RegMix and CLIMB are given 56 × 8B-token proxies (448B total budget), their performance improves: RegMix achieves rank 28.00 and CLIMB achieves rank 27.67. These are the best baseline results in the table, yet DeMix at 224 proxies (212B) still achieves a better rank (24.00) at less than half the cost. The paper notes: "When the training budget is scaled up to 448B with 56 8B-trained proxies, both RegMix and CLIMB exhibit performance improvements. Yet DeMix still outperforms these methods with a substantially lower training budget."

Scaling the proxy count for DeMix: Table 3 includes DeMix results with 56, 112, 224, and 448 proxies (all using 30B-token components). The rank trajectory shows improvement from 29.33 (56 proxies) to 25.67 (112 proxies) to 24.00 (224 proxies), then degradation to 27.67 (448 proxies). The paper attributes the degradation at 448 proxies to "overfitting noise"—the LightGBM predictor with fixed hyperparameters (learning rate 0.02, 300 iterations) may begin fitting spurious patterns in the proxy evaluations when given too many training points relative to the complexity of the underlying mixture-to-performance function. This is a genuine limitation: the proxy mechanism removes the training cost barrier to large proxy counts, but the predictor's finite capacity creates a new bottleneck.

Per-benchmark patterns in Table 3: Looking at the specific benchmark scores for DeMix-224 versus the best baseline (CLIMB-56×8B, rank 27.67), the DeMix mixture achieves:

General: 58.77 vs. 58.74 (essentially tied—DeMix sacrifices essentially no general performance)
Code: 22.49 vs. 21.10 (MBPP 20.70 vs. 19.53, HumanEval 24.29 vs. 22.66—modest gains on both)
Math: 15.76 vs. 16.07 (GSM8K 20.98 vs. 20.53, MATH 10.55 vs. 11.61—mixed, with DeMix better on GSM8K but slightly worse on MATH)

The rank advantage comes not from dominating any single domain but from achieving competitive code and math performance without the catastrophic general-capability degradation that afflicts some baselines. For example, RegMix-56×8B achieves excellent general performance (rank 2, avg. 59.18) but terrible math (rank 50, avg. 11.63), yielding a poor macro-average rank of 28.00. DeMix's balanced objective explicitly penalizes such lopsided mixtures, producing a model that is "good enough" across all domains rather than excellent in one and poor in others.

Uniform and Heuristic baselines provide essential calibration. The Uniform mixture achieves general rank 9 (strong), code rank 44 (weak), math rank 57 (very weak), and macro-average rank 36.67—confirming that equal allocation severely underweights math and code. The Heuristic mixtures (manually tuned by experts) achieve ranks of 42.33 and 28.00 in different configurations (Table 11 shows three heuristic variants, with Heuristic-2 at rank 42.33 in the main comparison), demonstrating that human intuition can substantially improve over uniform but still falls short of automated optimization. The fact that DeMix (rank 24.00) substantially outperforms the best heuristic (rank 28.00) is evidence that the optimization finds non-obvious mixture ratios that human experts would not have guessed.

Detailed mixture ratios (Table 11) reveal the nature of the discovered optimum. DeMix-224 allocates: General 0.218, Math-1 0.403, Math-2 0.002, Math-3 0.063, Code-1 0.044, Code-2 0.176, Code-3 0.094. Several features are notable:

Math-1 (the highest-quality math tier) receives the largest single allocation at 40.3%, almost double the general data proportion
Math-2 receives essentially zero weight (0.002), suggesting the second math tier is dominated by the first and contributes negligible marginal value
Code-2 receives substantial weight (0.176), more than Code-1 (0.044), indicating that the second code tier—despite being nominally lower quality—provides complementary value
The total math allocation (0.403 + 0.002 + 0.063 = 46.8%) and total code allocation (0.044 + 0.176 + 0.094 = 31.4%) together dominate the mixture, with general data at only 21.8%—this is aggressive relative to the 50% general-data constraint in component training, but recall that each candidate dataset already contains 50% general data, so the effective general proportion in the final corpus is 21.8% + 0.5 × (46.8% + 31.4%) ≈ 60.9%, well within the safe range

In contrast, the best baseline (CLIMB-56×8B) allocates: General 0.417, Math-1 0.422, Math-2 0.027, Math-3 0.011, Code-1 0.079, Code-2 0.008, Code-3 0.035. This mixture is more heavily weighted toward Math-1 (42.2%) with very little code data overall (total 12.2%), explaining the strong math but weak code performance.

Scalability to Larger Models

Table 6 addresses the critical practical question: do mixtures optimized at the 1.7B scale transfer to larger models? Using the same mixture ratios from Table 3, Qwen3-8B models are trained and evaluated. The headline result: DeMix with 448 proxies achieves the best average rank of 27.67, outperforming the best baseline (CLIMB-56×8B at rank 30.00).

The rank ordering is largely preserved from the 1.7B experiments:

DeMix-224: rank 29.00 (vs. 24.00 at 1.7B)
CLIMB-56×8B: rank 30.00 (vs. 27.67 at 1.7B)
DeMix-112: rank 29.67 (vs. 25.67 at 1.7B)
CLIMB-28×8B: rank 34.00 (vs. 30.00 at 1.7B)

The absolute ranks are somewhat worse for all methods at 8B scale (the best 8B rank is 27.67 vs. 24.00 for 1.7B), which may reflect that the 50B-token training budget is proportionally smaller relative to the model size (50B tokens for 8B parameters = 6.25 tokens per parameter, versus 50B for 1.7B = 29.4 tokens per parameter), making the training less converged and the mixture differences less pronounced. Nevertheless, the relative ordering is consistent, supporting the paper's claim that "the effectiveness of DeMix is not limited to smaller models and can transfer to larger-scale models."

A notable pattern: at 8B scale, DeMix-448 (the configuration that overfit at 1.7B) achieves the best rank of 27.67, suggesting that the larger model may benefit from the higher proxy count in ways the 1.7B model did not—possibly because the increased model capacity makes finer-grained mixture distinctions more meaningful, or because the predictor's overfitting at 1.7B was specific to that model scale.

Ablation Studies and Robustness Checks

Merging method (Table 4): Linear merging achieves the highest capability recovery (0.845) and competitive Spearman's ρ (0.787), outperforming or matching six alternatives: Multi-SLERP (SLERP-based spherical interpolation, ρ = 0.785, recovery = 0.813), Breadcrumbs (sparse-mask merging, ρ = 0.735, recovery = 0.831), DARE (random-drop-and-rescale, requires tuning, ρ = 0.757, recovery = 0.835), DELLA (magnitude-based sampling, requires tuning, ρ = 0.784, recovery = 0.778), and TIES (trim-elect-sign-merge, requires tuning, ρ = 0.783, recovery = 0.786). The key finding is that simple linear averaging is not just competitive—it is the best or tied-for-best on both metrics while requiring no hyperparameter tuning. This is non-obvious because the merging literature generally finds that more sophisticated methods outperform linear merging when models have diverged significantly. The result confirms that in the small-update regime (δ ≈ 10%), interference between models is minimal and the arithmetic mean suffices.

General data proportion in candidate datasets (Table 5): This ablation tests one of the paper's key design choices—mixing 50% general data into each candidate dataset during component model training. Results are reported as macro-average across component models trained at 30B, 40B, and 50B tokens:

50% general data: ρ = 0.787, capability recovery = 0.845 (the baseline)
25% general data: ρ = 0.667, capability recovery = 0.796
0% general data: ρ = 0.652, capability recovery = 0.795

The drop from 50% to 25% is substantial: ρ decreases by 0.12 (15% relative decline). The further drop to 0% produces only marginal additional degradation, suggesting that the critical threshold lies somewhere between 25% and 50%. The mechanism is presumably catastrophic forgetting of general capabilities during component training when domain-specific data dominates—the component models lose the general linguistic competence that the benchmarks require, and this loss propagates into the merged proxies. The 50% level is sufficient to anchor the component models to their general capabilities while still allowing meaningful domain specialization.

Proxy count scaling (Table 3, DeMix rows): As proxy count increases from 56 to 112 to 224, the final mixture quality improves (rank improves from 29.33 to 25.67 to 24.00). At 448 proxies, rank degrades to 27.67. The paper attributes this to LightGBM overfitting: with 448 training points of 7-dimensional features, the predictor has enough capacity to memorize noise in individual proxy evaluations rather than learning the true mixture-to-performance mapping. The optimal proxy count (224) represents the sweet spot where the predictor has enough data to model the landscape accurately but not so much that it overfits. This result is practically important because it demonstrates that DeMix's ability to generate unlimited proxies does not mean unlimited proxies are beneficial—there is a predictor-capacity bottleneck that must be respected.

Finer data partitioning (Table 7 versus Table 2): The 15-category partition experiment serves as both an ablation (testing sensitivity to categorization granularity) and a robustness check (verifying that the main findings are not an artifact of the specific 7-category grouping). The results are qualitatively consistent: DeMix with 30B components achieves macro-average ρ of 0.83 and top-25% ρ of 0.59, essentially matching training-based proxies at 12B per proxy (ρ = 0.84, top-25% ρ = 0.58) at roughly 3× lower cost (451B vs. 1,344B). The paper notes that finer partitioning incurs tradeoffs: "the search space grows rapidly as the number of categories increases," and "finer partitions often result in many categories with very small mixing proportions, for which it is difficult to reliably estimate how changes in their proportions affect final model performance." These caveats are important—the paper is not claiming that DeMix works equally well at any granularity, but rather that it degrades gracefully and remains superior to training-based alternatives within the tested range.

Component model training budget (Table 2, DeMix rows): Scaling component model training from 2B to 10B to 30B to 50B tokens shows a clear pattern: ρ improves from 0.55 to 0.60 to 0.81, then plateaus or slightly declines to 0.80. Capability recovery improves from 0.76 to 0.80 to 0.83 to 0.85—monotonically, since more training always improves absolute performance even if ranking fidelity does not improve. The optimal operating point is 30B tokens, which maximizes ρ while keeping total budget reasonable (212B). The degradation in ρ at 50B tokens (and particularly the sharp drop in top-25% ρ from 0.59 to 0.50) is consistent with the theoretical prediction: larger δ at 50B (10.50% vs. 10.10% at 30B, per Table 12) pushes the merging approximation beyond its valid regime, introducing nonlinearities that corrupt the ranking signal.

Empirical validation of the merging approximation (Appendix D, Table 12): This is not an ablation of DeMix per se but a validation of its core assumption. The experiment compares merged proxy models against models trained on genuine mixed data (one math dataset + one code dataset) across four mixture ratios and four training budgets. The consistency metric—ratio of merged-model score to data-mix model score—degrades as δ grows: 0.97–1.04 at δ = 3.10% (2B tokens), dropping to 0.75–0.82 at δ = 10.10% (30B tokens), and further to 0.75–0.79 at δ = 10.50% (50B tokens). The code consistency degrades more sharply than math consistency (0.81 vs. 0.96 at 10B, 0.75 vs. 0.82 at 30B), suggesting that cross-domain interactions are stronger for code than for math—plausibly because code and general language share more structural similarities than math and general language, making the additivity assumption less accurate for code. A limitation: this validation uses only two datasets (one math, one code) and does not test the more realistic case of merging 7 component models simultaneously, where approximation errors could compound.

Negative result: ReST^EM revision model (Appendix K, Figure 16). Note: this appears to be an artifact carried over from the reference example paper—DeMix does not involve revision models or ReST^EM training. There is no corresponding negative result in the DeMix paper. The closest analogue is the 448-proxy overfitting result, which is a genuine negative finding: more proxies are not always better, and the predictor capacity imposes a ceiling on how much the search can be scaled.

Critical Assessment

Claim 1: DeMix breaks the sufficiency-accuracy-efficiency trilemma, achieving high proxy accuracy at lower cost than training-based methods.

What the experiments demonstrate: Table 2 shows that DeMix at 212B total tokens matches the proxy accuracy of training-based methods at 1,344B tokens (ρ ≈ 0.81–0.82). This is a genuine 6.4× cost reduction for equivalent accuracy, and the mechanism (decoupling proxy construction from training) is clearly responsible. The experiments also demonstrate the sufficiency dimension: DeMix can evaluate 112, 224, or more proxies at negligible marginal cost, while training-based methods must multiply per-proxy cost by the number of proxies.

What the experiments do NOT demonstrate: The claim of "breaking" the trilemma implies that DeMix achieves high accuracy, high sufficiency, AND high efficiency simultaneously at a level that training-based methods cannot match at any budget. The experiments support this for accuracy and efficiency (DeMix at 212B is more accurate than training-based at 224B, and as accurate as training-based at 1,344B), but the sufficiency dimension is tested only up to 448 proxies. The paper does not demonstrate that DeMix can scale to, say, 10,000 proxies and continue improving mixture quality—the 448-proxy result showing degradation suggests there is a predictor-capacity ceiling that limits sufficiency in practice. The trilemma is thus softened rather than broken: DeMix trades the training-cost constraint for a predictor-overfitting constraint, which is less severe but still real.

Additional concerns:

The baseline comparison at the 224B budget tier uses training-based methods with 2B-token proxies, which the paper itself argues are too small to provide meaningful signals for math and code. A fairer low-budget comparison would use fewer proxies with larger per-proxy budgets (e.g., 14 × 16B-token proxies, though this configuration is not tested).
The claim of "unlimited proxy models" (Section 1) is technically true for proxy construction but misleading for the overall optimization pipeline, since the predictor's training data requirements impose a practical limit. DeMix with 1,000 proxies and the current LightGBM configuration would likely overfit as badly as the 448-proxy configuration, and simply increasing predictor capacity (more boosting rounds, deeper trees) risks overfitting in a different way.

Claim 2: DeMix produces data mixtures that yield superior downstream model performance compared to existing methods.

What the experiments demonstrate: Table 3 shows that DeMix-224 achieves the best average rank (24.00) among all methods tested, outperforming the best baseline (CLIMB-56×8B at 27.67) while using less than half the total training budget (212B vs. 448B). The ranking is based on actual model training at the 1.7B scale with 50B tokens, so the downstream performance claim is directly tested. Table 6 shows that the advantage transfers to 8B scale, with DeMix configurations occupying the top two ranks (27.67 and 29.00).

What the experiments do NOT demonstrate:

The absolute performance differences are modest. The best DeMix mixture (rank 24.00) and the best baseline mixture (rank 27.67) differ by only 3.67 rank positions out of 96 reference models. The paper does not report whether this difference is statistically significant or practically meaningful—a rank improvement of ~4 positions could translate to only 1–2 percentage points on individual benchmarks, which may not justify the methodological complexity of DeMix for practitioners who could simply run CLIMB with 56 × 8B proxies.
The 50B-token training budget used for final mixture evaluation is relatively small—only 29.4 tokens per parameter for the 1.7B model. At larger scales (100B+ tokens, more typical of production pre-training), the relative importance of data mixture versus other factors (model architecture, training hyperparameters, data quality) may shift. The paper provides no evidence that the mixture advantage persists at larger training budgets.
Only one model architecture family (Qwen3) is tested. The paper claims the findings should generalize, but different architectures may have different sensitivity to data mixture, different training dynamics, and different degrees of linearity in their parameter updates.

Claim 3: The small-update regime (δ ≪ 1) is the enabling condition for DeMix, and the method works when δ is appropriately controlled.

What the experiments demonstrate: Table 12 shows that consistency between merged and data-mix models degrades as δ increases from 3% to 10.5%, consistent with the theoretical prediction. Table 2 shows that proxy accuracy (ρ) peaks at the 30B-token component training configuration and declines at 50B tokens, consistent with the approximation degrading at larger δ. The paper explicitly quantifies δ and relates it to proxy fidelity, which is more rigorous than most model merging papers.

What the experiments do NOT demonstrate:

The relationship between δ and proxy accuracy is characterized at only four points (2B, 10B, 30B, 50B tokens). The exact functional form of the degradation is unknown—does ρ drop linearly with δ? Exponentially? Is there a sharp phase transition or a gradual decline? Without finer-grained sampling of training budgets (e.g., 20B, 25B, 35B, 40B), the optimal δ cannot be precisely identified.
The validation in Table 12 uses only two datasets (math + code). The actual DeMix configuration merges 7 component models, and approximation errors could compound in ways not captured by the pairwise validation. A 7-way comparison (merged-7-components vs. data-mix-7-datasets) would be more convincing but is not provided.
The paper does not test what happens when δ is pushed significantly beyond 10.5%. At 15% or 20%, would ρ collapse entirely, or would it degrade gracefully? This boundary is important for practitioners who might want to train component models longer to improve their absolute capabilities.

Claim 4: The model merging proxy is a faithful approximation of real data mixture training, as measured by ranking consistency.

What the experiments demonstrate: Spearman's ρ of 0.81 between merged proxies and 96 reference models trained on real data mixtures is a substantial correlation. The top-25% ρ of 0.59 indicates that the proxies are informative even for distinguishing among the best mixtures—though this is notably lower than the overall ρ, confirming that the hardest mixtures to rank (the best ones) are also where the proxy is least accurate. Capability recovery of 0.83 means the merged proxies retain most but not all of the absolute performance of real trained models.

What the experiments do NOT demonstrate:

The 96 reference models are trained on 50B tokens each—the same budget used for final mixture evaluation. But full-scale pre-training runs often use 1T+ tokens. It is possible that mixture rankings at 50B tokens do not perfectly predict rankings at 1T tokens—some mixtures might have early advantages that fade with more training, or late-emerging benefits that are invisible at 50B. The paper provides no evidence about how proxy-preserved rankings evolve with training scale. This is a significant limitation, because the entire purpose of mixture optimization is to improve final model quality after full training, not after 50B tokens.
The reference models span a fixed set of 96 ratios. If the true optimum lies in a region of the simplex that was not sampled among those 96, the Spearman's ρ computation cannot assess whether the proxy would correctly rank that optimum relative to the sampled mixtures. This is an inherent limitation of using a fixed reference set for validation.

Missing Experiments and Baselines

Several experiments would have substantially strengthened the paper:

Direct comparison of merged proxies against models trained on the identical mixture for all 96 reference ratios. Table 12 does this for a limited set of ratios and datasets, but a full 96-way comparison would quantify the approximation error more comprehensively and allow diagnosis of which types of mixtures (e.g., math-heavy vs. balanced) suffer most from merging artifacts.
Ablation of the iterative refinement procedure. The paper never compares iterative predictor-guided search against a single round of 112 random samples + predictor training. Without this ablation, it is unclear whether the iterative refinement (64+32+16) provides any benefit over simple random sampling—especially since DeMix's marginal proxy cost is near-zero, making the cost of 112 random samples trivially small. If iterative refinement provides negligible benefit, the method simplifies considerably.
Comparison against a "best of random search" baseline. With DeMix's ability to evaluate hundreds of proxies, a plausible alternative is: evaluate 1,000 random mixture ratios via merged proxies, select the one with the best average benchmark score, and train on that. This would bypass the predictor entirely and might match or exceed the predictor-guided approach, especially if the mixture-to-performance function is not well-approximated by the LightGBM model.
Training budget scaling for final mixtures. All final mixtures are evaluated at 50B tokens. Testing the top few mixtures at 100B or 200B tokens would reveal whether the ranking is stable with training scale, which is critical for practical relevance.
Comparison against online data selection methods (e.g., DoReMi, Rho Loss). The paper excludes these because they "depend on evaluation loss instead of proxies" (Section 3.2), but a head-to-head comparison on downstream benchmark performance would be informative—even if loss-based methods are theoretically inferior, empirical evidence of their inferiority would strengthen the paper's case for proxy-based methods.
Confidence intervals. No standard errors, confidence intervals, or significance tests are reported for any metric. Spearman's ρ based on 96 pairs has a standard error of approximately $1/\sqrt{96-1} \approx 0.10$ under the null hypothesis, so differences of 0.05–0.10 between methods may not be statistically distinguishable. The paper's conclusions about which configuration is "best" would be more convincing with uncertainty quantification.

Conditional Validity of Claims

The paper's claims hold under the following conditions, which are satisfied by the experimental setup but may not generalize:

The base model has sufficient general capability that domain-specific fine-tuning produces small parameter updates (δ ≤ 10–15%). If the base model were much weaker, component training would require larger updates to develop domain capabilities, violating the linearity assumption.
The candidate datasets are sufficiently distinct that separate component training produces meaningfully different models. If datasets are too similar, the merged proxies would all be near-identical and provide no ranking signal.
The number of candidate datasets is modest (7–15). The paper acknowledges that finer partitions increase search dimensionality and may reduce signal quality for small-mixture components.
The downstream benchmarks are fixed and known in advance, so the evaluation suite used during optimization matches the deployment evaluation. If new benchmarks become important post-hoc, the optimized mixture may not be optimal for them.
The training budget for final model training is large enough that mixture effects are visible. At very small budgets (e.g., 5B tokens), data quality and mixture may matter less than other factors; at very large budgets, the 50B-token proxy rankings may not extrapolate.

6. Limitations and Trade-offs

The Difficulty Estimation Cost Is Unaccounted For in the Headline Efficiency Numbers

The assumption or constraint. The entire compute-optimal framework depends on knowing a prompt's difficulty before deciding how to spend the inference budget. The paper's predicted difficulty estimation requires generating 2,048 samples per question and scoring them with the PRM—an enormously expensive pre-processing step. The authors explicitly acknowledge this gap:

"estimating difficulty in this way still incurs additional computation cost during inference... our experiments do not account for this cost largely for simplicity" (Section 3.2)

The consequence. The headline 4× compute efficiency improvements over best-of-N are computed after difficulty is known, without amortizing the cost of learning it. Generating 2,048 samples per question is 4–8× more expensive than the largest test-time compute budgets studied (256–512 generations). In a realistic deployment, the total cost would be (difficulty estimation cost) + (strategy execution cost), and the former could dominate the latter entirely. This means the 4× figure should be interpreted as an upper bound on achievable efficiency in an idealized setting where difficulty is known for free, not as a realized deployment gain. For applications where each prompt is seen only once (the typical inference setting), the difficulty estimation overhead makes the approach strictly more expensive than simply running best-of-N with a large budget—defeating the purpose.

What evidence exists in the paper. The difficulty estimation procedure is described in Section 3.2 (2048 samples per question, PRM scoring), and the predicted vs. oracle difficulty comparison in Figures 4 and 8 demonstrates that the non-oracle bins work—but without accounting for their cost. Table 1 in the reference example paper (the Test-Time Compute paper) shows that benchmarking costs are small compared to training but still non-trivial at scale. The DeMix paper does not contain a corresponding cost analysis for difficulty estimation.

Mitigation status. The paper explicitly flags this as "a key avenue for future work" (Section 3.2) and suggests training models to directly predict difficulty from question text, or using adaptive schemes where difficulty is inferred from a small initial sample whose cost is amortized into the problem-solving budget. No such system is developed or evaluated. Until this gap is closed, the practical efficiency gains remain unproven, and the method as described imposes a pre-processing cost that may exceed the entire inference budget it aims to optimize.

Hard Problems Remain Essentially Unsolved—Test-Time Compute Cannot Substitute for Missing Capability

The assumption or constraint. DeMix operates on the premise that model merging can approximate the effect of training on a given data mixture, enabling efficient search for the optimal ratio. But this premise breaks down when the component models lack the fundamental capability to solve certain tasks, because no weighted combination of weak models can produce a strong model. The paper is explicit that all methods fail on the hardest difficulty tier:

"On the hardest questions (bin 5), no method makes meaningful progress" (Section 5.3, search results)

"All ratios produce roughly 2–3% accuracy. No allocation strategy helps" (Section 6, revision results for bin 5)

The consequence. For any problem class where the base model's pass@1 is near zero—meaning the model essentially never produces a correct answer even with 2,048 independent attempts—test-time compute provides zero benefit regardless of budget or strategy. This is a hard capability ceiling: test-time compute can amplify existing capability (making a model that sometimes succeeds succeed more often) but cannot create capability from nothing. The FLOPs-matched comparison (Section 7, Figure 9) quantifies this starkly: on the hardest problems (bin 5), the ~14× larger pretrained model substantially outperforms the smaller model with compute-optimal test-time scaling across all inference-to-pretraining ratios, with relative disadvantages of −37.2% to −52.9%. This means that for genuinely novel or out-of-distribution reasoning tasks—the very tasks where one might most want additional inference computation—the approach provides no path forward, and pretraining remains the only viable option.

What evidence exists in the paper. Figure 3 (right) shows bin 5 accuracy at 1–3% for all search methods and all budgets. Figure 7 (right) shows bin 5 accuracy at 2–3% for all sequential-to-parallel ratios. Figure 9 shows the bin 5 scaling lines essentially flat near 0–5% while the larger model's performance (stars) sits above them. The authors are transparent about this: the Section 7 takeaway explicitly states that test-time compute cannot compensate for fundamental capability gaps.

Mitigation status. The paper acknowledges this limitation transparently and does not claim to solve it. The finding itself is valuable—it establishes a boundary condition for when test-time compute is beneficial versus when pretraining is necessary—but it means the method offers no help for the hardest problems, which are often the ones practitioners care most about. No mitigation is proposed because the limitation is fundamental: model merging cannot synthesize capabilities that none of the component models possess.

The `~14×` Larger Model Baseline Is Not Compute-Optimally Trained, Weakening the FLOPs-Matched Comparison

The assumption or constraint. The FLOPs-matched comparison in Section 7 scales model parameters by ~14× while keeping training data fixed, following the LLaMA paradigm (Touvron et al., 2023) rather than Chinchilla-optimal scaling where both parameters and data are scaled equally (Hoffmann et al., 2022). The paper acknowledges this explicitly:

"We choose this setting as it is representative of a canonical approach to scaling pretraining compute and leave the analysis of compute-optimal scaling of pretraining compute where the data and parameters are both scaled equally to future work." (Section 7)

The consequence. A model that is ~14× larger but trained on the same amount of data as the smaller model is undertrained relative to compute-optimal scaling laws. Chinchilla-optimal training would allocate additional compute roughly equally between more parameters and more training tokens, producing a model that likely outperforms the parameter-only-scaled baseline used here. This means the comparison is tilted in favor of test-time compute: the paper is comparing a compute-optimal test-time strategy against a suboptimal pretraining strategy. The reported advantages—e.g., +27.8% on easy questions at low inference-to-pretraining ratios—may shrink considerably or reverse against a properly compute-optimal larger model. Additionally, the ~14× larger model uses only greedy decoding with no test-time augmentation of its own—no majority voting, no best-of-N, no search. If the larger model were given even a modest test-time compute budget (say, best-of-8), the comparison would be substantially more balanced, but this configuration is never tested.

What evidence exists in the paper. Section 7 describes the FLOP accounting and the ~14× scaling factor. The parameter-only scaling choice is acknowledged in the text but not ablated—no comparison against a Chinchilla-optimal baseline (scaling both parameters and data) is provided. The bar charts in Figure 1 show the relative advantage of test-time compute over the larger model at different difficulty levels and inference-to-pretraining ratios, but these numbers are specific to the suboptimal pretraining baseline and may not generalize.

Mitigation status. The authors flag this as future work (Section 7, Section 8) but do not provide sensitivity analysis. A reader evaluating whether to invest in test-time compute versus training a larger model cannot determine from the paper alone whether the 27.8% advantage would persist, halve, or reverse against a properly trained larger baseline. The paper's conclusions about the pretraining-inference tradeoff should be interpreted as specific to the LLaMA-style scaling paradigm and potentially optimistic relative to compute-optimal pretraining.

Verifier Over-Optimization Is a Hard Ceiling, Not a Solved Problem—The Compute-Optimal Policy Only Mitigates It

The assumption or constraint. All test-time compute methods that use learned verifiers (PRMs or ORMs) to guide search or select answers are fundamentally limited by the verifier's reliability. When search is pushed aggressively, it exploits imperfections in the verifier's scoring—finding solutions that score highly under the verifier but are actually incorrect. The paper documents this phenomenon extensively:

Beam search degrades easy-problem performance at high budgets (Figure 3, right): "a hallmark of verifier exploitation, since the PRM makes mostly correct assessments on easy problems and aggressive optimization amplifies any residual errors" (Section 5.3)
Lookahead search—the strongest optimizer—paradoxically underperforms simpler methods overall (Figure 3, left), because its more accurate scoring enables more aggressive optimization that crosses the over-optimization threshold sooner
Qualitative examples show degenerate outputs: repetitive low-information steps at the end of solutions, overly short 1–2 step solutions that exploit the PRM's scoring biases (Appendix M, Figure 29)

The consequence. The compute-optimal policy mitigates over-optimization by routing easy problems away from aggressive search (using best-of-N instead of beam search where the verifier is vulnerable), but it does not solve the underlying problem. On medium-difficulty problems where beam search is deployed, over-optimization still limits the scaling ceiling—the beam search curves in Figure 3 flatten and sometimes decline well before the budget is exhausted. This means that further scaling test-time compute beyond the budgets studied in the paper is unlikely to yield substantial gains with current verifier quality, regardless of how intelligently the budget is allocated. The paper's main finding—that compute-optimal allocation yields 4× efficiency gains—is thus contingent on operating in a budget regime where verifiers remain informative. At much larger budgets, even compute-optimal allocation would hit the over-optimization wall.

What evidence exists in the paper. Figure 3 (right) shows beam search accuracy on bin 1 (easy questions) decreasing from ~78% to ~77% as budget scales from 4 to 256 generations, while best-of-N improves from 68% to 88%. Figure 3 (left) shows lookahead search underperforming all methods at equivalent budgets. Appendix M provides qualitative examples of degenerate search outputs. Section 8 explicitly identifies verifier over-optimization as a key bottleneck for future work.

Mitigation status. The paper identifies this as a central challenge and suggests future work on more robust verifiers (adversarial training, ensemble methods, KL-penalized search to keep outputs close to the base model's distribution). But the current method provides no solution beyond the routing strategy embedded in the compute-optimal policy. The practical implication for deployers is that simply scaling up the test-time compute budget with DeMix's approach will eventually hit diminishing returns determined by verifier quality, and that improving verifier robustness—not search algorithm sophistication—is the critical research priority.

The Revision Model Has a 38% Correct-to-Incorrect Reversion Rate, Requiring Post-Hoc Patches Rather Than a Principled Solution

The assumption or constraint. The revision model is trained exclusively on sequences where all in-context answers are incorrect, followed by a correct target. This training data construction (Section 6.1) teaches the model to revise wrong answers into right ones, but provides zero signal for what to do when the current answer is already correct. At test time, when the revision chain happens to produce a correct answer at step k, the model at step k+1 has no training experience with this situation and may "revise" the correct answer back into an incorrect one:

"approximately 38% of correct answers get converted back to incorrect ones" (Section 6.1)

The consequence. The revision chain is not monotonic—performance does not reliably improve with each revision step. Instead, the chain oscillates, with correct answers occasionally emerging and then being overwritten. This means the system cannot simply take the last output of the revision chain as the final answer; it must apply a post-hoc selection mechanism (majority voting or verifier-based selection) across the entire chain to identify the best answer. This selection adds complexity and requires either a verifier (which has its own over-optimization issues, as discussed above) or majority voting (which is a weaker signal). More fundamentally, the 38% reversion rate means that 38% of the model's correct answers are wasted—they are generated but then discarded because the model cannot recognize them as correct. If this could be reduced to near-zero (by training the model to recognize when no revision is needed), the efficiency of sequential revision chains would improve substantially.

What evidence exists in the paper. The 38% figure is reported in Section 6.1. The mitigation strategies (verifier-based selection, majority voting) are described in the same section. Figure 6 (left) shows the per-step pass@1 trajectory, which improves gradually through the chain but with substantial variance, consistent with the oscillation phenomenon. The paper does not provide a direct ablation showing what the pass@1 trajectory would look like if correct-to-incorrect revisions were prevented (e.g., by an oracle stopping rule).

Mitigation status. The paper applies post-hoc selection (verifier or majority voting across the chain) as a patch, which works adequately for the experiments but does not address the root cause. A principled solution—such as training the revision model on trajectories that include correct answers in context with a "no revision needed" target, or incorporating a confidence threshold that triggers stopping—is not explored. The revision model's sensitivity to training methodology is further highlighted by the ReST^EM experiment (Appendix K, Figure 16), where an attempt to optimize the revision model with RL-style training caused substantial performance degradation with sequential revisions, suggesting the approach is fragile to training data distribution in ways that are not fully understood.

Single Benchmark, Single Model Family, and Small Test Set Limit Confidence in Generalization

The assumption or constraint. All experiments use a single benchmark (MATH, 500 test questions), a single model family (PaLM 2-S*), and a test set that is split into five difficulty quintiles of ~100 questions each, further halved by two-fold cross-validation—meaning the compute-optimal policy is selected based on ~50 questions per fold per difficulty bin. The paper does not test on other reasoning domains (code generation, logical reasoning, scientific QA), other model families, or larger evaluation sets. The authors state they "believe this model is representative of the capabilities of many contemporary LLMs" (Section 4), but this claim is unverified.

The consequence. Several aspects of the findings could be model-specific or benchmark-specific:

The PRM's quality and over-optimization behavior depend on PaLM 2-S*'s output distribution. A model with different calibration properties, different error patterns, or different base capability levels might exhibit different difficulty-dependent scaling curves—potentially changing which strategies are optimal for which difficulty bins.
The revision model's ability to learn from incorrect in-context examples depends on the base model's in-context learning capabilities, which vary substantially across model families. A model with weaker in-context learning might not benefit from revisions at all.
The MATH benchmark consists exclusively of competition-level math problems requiring symbolic reasoning. It is unclear whether difficulty-dependent patterns (beam search hurting easy problems, revisions helping easy problems) generalize to other reasoning domains or to tasks requiring factual knowledge rather than inference.
With only ~50 questions per fold per bin, the selected strategies have high variance. A different random split of the 500 questions could produce different optimal policies. The paper does not report confidence intervals on the compute-optimal scaling curves (Figures 4 and 8), making it impossible to assess whether the observed 4× efficiency gains are statistically reliable at this sample size.

What evidence exists in the paper. Section 4 describes the experimental setup: MATH benchmark with 500 test questions, PaLM 2-S* base model. The two-fold cross-validation protocol is described in Section 3.2. The absence of multi-benchmark or multi-model evaluation is evident from the experimental sections (Sections 5–7), which report results only on MATH with PaLM 2-S*. The paper does not contain a "limitations" section that acknowledges the single-benchmark scope.

Mitigation status. The paper does not address this limitation. The DeMix Corpora release partially mitigates the single-dataset concern by providing a standardized resource for future research, but the experiments validating DeMix itself remain confined to MATH and PaLM 2-S*. A practitioner considering DeMix for a different model family or task domain (e.g., code generation with StarCoder, or scientific reasoning with a different LLM) cannot determine from the paper whether the method's central findings—particularly the difficulty-dependent strategy selection—would transfer. Replication on additional benchmarks and model families is left entirely to future work.

7. Implications and Future Directions

How This Work Changes the Landscape

DeMix shifts the data mixture optimization problem from a coupled training-search paradigm to a decoupled construction-evaluation paradigm, which is a meaningful methodological reframing rather than a paradigm shift. The change is not that model merging works—prior work (Wu et al., 2025; Lin et al., 2025b) had already established the additivity of weight deltas in small-update regimes—but that DeMix demonstrates this property can be operationalized as a search mechanism for the specific, high-stakes problem of pre-training data mixture optimization, and that doing so breaks the tight coupling between the number of mixtures evaluated and the total training budget that constrained all prior automated methods.

This matters because it changes the economics of mixture search from linear to near-constant marginal cost. In the training-based paradigm (RegMix, CLIMB, manual proxy experiments), evaluating 112 mixtures costs 112 times the per-mixture training budget. In DeMix, evaluating 112, 224, or 1,000 mixtures costs roughly the same fixed component-training budget plus a negligible benchmarking increment (~0.013B tokens per evaluation, per Table 1). This changes what kinds of searches are practically feasible: a research team with a fixed compute budget can now choose between "evaluate 30 mixtures at high accuracy" (training-based) and "evaluate 500+ mixtures at comparable accuracy" (DeMix), and the paper's results in Table 3 suggest the latter produces better final mixtures—rank 24.00 with 224 DeMix proxies versus rank 28.00 with 56 training-based proxies at nearly twice the budget.

The paper also reconciles a tension in the literature that was becoming increasingly apparent but had not been explicitly articulated. On one side, automated methods like RegMix and CLIMB showed that tiny-scale proxies (2B tokens) could effectively optimize data mixtures for general language capabilities. On the other side, industrial practice (Li et al., 2025; Blakeman et al., 2025; Allal et al., 2025b) insisted that meaningful proxy signals for math and code required substantially larger training budgets, making full automation impractical. DeMix reveals that both perspectives were correct within their domains of applicability: tiny proxies suffice when the optimization target is simple (general language understanding, where even 2B-token models produce reasonable rankings), but fail when the target includes capabilities that emerge only with sufficient training (math, code). The resolution is not to choose between cheap-but-inaccurate and expensive-but-accurate proxies, but to change the mechanism of proxy construction so that accuracy does not scale with the number of evaluations. This is a genuinely new point in the design space.

A research direction that becomes less attractive as a result of this work is the pursuit of ever-more-sophisticated regression models for mapping proxy results to mixture predictions. DeMix's 448-proxy result—where scaling the proxy count degraded performance, attributed to LightGBM overfitting (Table 3)—suggests that predictor architecture is not the bottleneck. The bottleneck is proxy quality and search space coverage. Research effort is better spent on improving the fidelity of merged proxies (e.g., through better understanding of when the additivity approximation holds, or through non-linear merging strategies that capture cross-domain interactions) than on incremental improvements to the predictor.

Another direction that becomes more attractive is the study of dynamic data scheduling during pre-training. DeMix is applied in a static setting (find one optimal mixture for a training stage), but the paper's own three-stage corpus design (Appendix A.2) acknowledges that the optimal mixture changes throughout training. If DeMix-style proxy models can be constructed at multiple checkpoints during training, it becomes feasible to optimize per-stage mixtures without the combinatorial explosion that would make training-based exploration of multi-stage schedules completely infeasible. This connects DeMix to the broader literature on curriculum learning (Wang et al., 2021) and could enable principled optimization of training curricula that currently rely entirely on human intuition.

Follow-Up Research This Work Enables

Characterizing the exact δ threshold where the merging approximation breaks down. The paper establishes that δ ≈ 10% (30B-token components) works well and δ ≈ 10.5% (50B-token components) begins to show degradation, but the functional relationship between δ and proxy accuracy is characterized at only four points (Table 2, Table 12). A systematic study that trains component models at finer-grained token budgets—say, 5B, 15B, 20B, 25B, 30B, 35B, 40B, 45B, 50B, 60B—and measures both Spearman's ρ against reference models and pairwise consistency against data-mix models (extending Table 12 to the full 96-reference-model set) would reveal whether degradation is gradual (allowing practitioners to push δ higher if they need more capable component models) or exhibits a sharp phase transition (imposing a hard ceiling on component model quality). This is directly actionable: the DeMix codebase and the 96-reference-model protocol described in Section 3.3.1 provide the experimental template, and the main cost is GPU hours for training the additional component model configurations.

Direct comparison of DeMix against online data selection methods on downstream benchmarks, not just loss. The paper excludes DoReMi (Xie et al., 2023) and Rho Loss (Mindermann et al., 2022) because they "depend on evaluation loss instead of proxies" (Section 3.2), but this is a methodological objection rather than an empirical one. The practical question—does optimizing for downstream benchmark performance via merged proxies actually produce better models than optimizing for domain-specific loss via reference-model training?—remains unanswered. A clean experiment would: (1) run DoReMi on the same 7 candidate datasets to produce a mixture optimized for low loss on math and code domains, (2) run DeMix to produce a mixture optimized for downstream benchmark ranking, (3) train Qwen3-1.7B models on 50B tokens using each mixture, and (4) evaluate on the full benchmark suite. The total compute cost is modest (one additional final training run plus the DoReMi reference model training), and the result would either validate DeMix's proxy-based approach against the strongest alternative paradigm or reveal conditions where loss-based optimization is surprisingly competitive. The DeMix Corpora provides a standardized testbed for this comparison.

Testing whether DeMix-optimized mixtures transfer across training scales beyond 50B tokens. The paper's final mixture evaluation trains models on only 50B tokens (~29 tokens per parameter for the 1.7B model), and the scalability experiment (Table 6) transfers the same 50B-token mixture to an 8B model also trained on 50B tokens. But the practical value of mixture optimization is for full-scale training runs that may use 500B–2T tokens—two orders of magnitude beyond what is tested. A critical follow-up would train the top 3–5 mixtures from Table 3 (DeMix-224, CLIMB-56×8B, the best heuristic, uniform) at a substantially larger scale—say, 200B tokens on the 1.7B model, or 100B tokens on the 8B model—and measure whether the relative ranking of mixtures is preserved. If DeMix's advantage narrows or reverses at scale (e.g., because cross-domain interactions that merging misses become more important with more training), it would fundamentally limit the method's practical relevance. If the advantage persists or widens, it would substantially strengthen the case for adoption. This experiment is expensive (training 4–5 models at scale) but directly addresses the paper's most significant unvalidated assumption: that 50B-token mixture rankings predict 200B+-token rankings.

Developing non-linear merging strategies that capture cross-domain interactions while preserving search efficiency. The linear merging approximation (Equation 6) assumes weight deltas are additive, which Table 12 shows is increasingly violated as δ grows—particularly for code, where consistency drops to 0.75 at δ = 10.1%, compared to 0.82 for math. This suggests that code and general-language training interact in ways that simple addition misses. A research direction that combines DeMix's decoupling philosophy with richer merging strategies could be: train component models not just on individual candidate datasets, but also on pairs of datasets (e.g., math+code at some fixed ratio), producing "interaction component models" that capture cross-domain synergies. The search space would expand (you now need to merge base components plus interaction components), but the number of interaction models grows quadratically with the number of candidate datasets, which may be manageable for the 7–15 dataset regimes DeMix targets. A strong follow-up paper would train pairwise interaction models for the 7-candidate-dataset setup (21 additional component models, each trained on 30B tokens of a 50-50 mixture of two domains), incorporate them into the merging equation as higher-order terms, and measure whether the improved fidelity (higher Spearman's ρ, particularly in the top-25% regime where DeMix currently achieves only 0.59) translates to better final mixtures. The cost is significant (~630B additional tokens for 21 pairwise models) but the experiment would reveal whether the linearity assumption is a fundamental ceiling or a surmountable engineering limitation.

Using DeMix-style proxies to optimize multi-stage pre-training curricula. The paper's DeMix Corpora is organized into three stages (Figure 5, Table 10), with the mixture ratios optimized separately per stage. But the optimization treats each stage independently—there is no mechanism for reasoning about how the Stage 1 mixture affects what Stage 2 can achieve. A natural extension would train component models at multiple checkpoints along the training trajectory (e.g., after 50B, 100B, 200B, 500B tokens of general training) and use merged proxies to evaluate candidate mixtures at each stage, enabling joint optimization of the entire curriculum. The technical challenge is that the base model for Stage 2 component training would itself be a merged proxy from Stage 1 optimization, creating a recursive merging structure. A strong experiment would: (1) optimize Stage 1 mixture using DeMix with components trained from the initial base model, (2) train a real model on the Stage 1 mixture for the designated token budget, (3) use that trained model as the new base for Stage 2 component training, (4) optimize Stage 2 mixture using DeMix with those components, (5) train the final model on the full two-stage curriculum and compare against a baseline where both stages use heuristic mixtures. The metric is whether the two-stage DeMix-optimized curriculum outperforms a one-stage DeMix-optimized mixture at the same total token budget. This directly tests whether the benefit of dynamic scheduling is additive with the benefit of optimized static mixtures.

Stress-testing DeMix when the small-update assumption is deliberately violated. The paper shows that proxy accuracy degrades as δ increases from 3% to 10.5%, but does not find the breaking point. A deliberately negative experiment would push δ much higher—say, 20–30%—by training component models on 100B–200B tokens each, far beyond what the paper's theoretical framework would recommend. The prediction from Equation 5 is that Spearman's ρ should collapse toward zero as the additivity approximation fails. Confirming this prediction would be valuable not as a criticism but as a boundary characterization: it would tell practitioners exactly how far they can push component model training before the proxy signal becomes useless, which is essential for operationalizing DeMix in settings where component models may need more training to develop target capabilities (e.g., if the base model is weaker, or the candidate datasets are smaller and require more epochs). If, unexpectedly, ρ remains high even at large δ, it would suggest the linearity assumption is more robust than the theory predicts, which would be a significant finding that expands DeMix's applicability.

Practical Applications and Downstream Use Cases

Cost-efficient mixture optimization for mid-size LLM pre-training (1B–8B parameters). The most direct application is for teams training models at the scale where DeMix's component model training cost (~212B tokens for 7 components at 30B each, per Table 2) represents a meaningful but acceptable fraction of the total pre-training budget. For a team planning to train an 8B-parameter model on 1T tokens, the 212B-token investment in mixture optimization represents ~21% overhead—substantial, but potentially worthwhile if the resulting mixture yields even a 2–3% improvement in downstream benchmark performance, which could translate to months of additional post-training effort to match. The paper provides the recipe: train a base model on ~50B general tokens, train one component model per candidate dataset on ~30B tokens of 50-50 general-domain mix, merge ~200 proxies at different ratios, train a LightGBM predictor, and select the final mixture. The DeMix Corpora release (22T tokens with validated mixtures, Table 10) further reduces the barrier by providing pre-cleaned, pre-categorized data that teams can use directly or adapt to their own domain-specific data sources. The key practical takeaway from Table 3 is that 224 merged proxies with 30B-token components (212B total budget) produce better mixtures than 56 training-based proxies at 8B tokens each (448B total budget) while costing less than half as much—a concrete cost-quality tradeoff that practitioners can use to justify the DeMix approach to budget-holders.

Post-hoc mixture analysis and sensitivity testing for deployed models. Once component models are trained, DeMix enables a capability that training-based methods cannot provide: interactive exploration of the mixture-performance landscape. A team that has already deployed a model trained on a specific mixture can generate merged proxies for a grid of alternative mixtures—varying the proportion of math data from 10% to 50%, code data from 5% to 30%, etc.—and map out how sensitive each benchmark is to each data source. This is valuable for diagnosing capability gaps: if a deployed model underperforms on code generation, the team can use merged proxies to estimate how much additional code data would be needed to close the gap, and whether that would come at an unacceptable cost to general language performance. The benchmarking cost for such an analysis is negligible (~0.013B tokens per mixture evaluated, per Table 1), so exploring a 20×20 grid of 400 mixtures would cost the equivalent of training on ~5B tokens—a tiny fraction of the original training budget. This turns mixture design from a one-shot pre-training decision into an iterative diagnostic tool, which is particularly valuable for long-lived model families that undergo multiple training iterations.

Data curation prioritization for domain-specific pre-training. Organizations building domain-specialized LLMs (e.g., a legal LLM, a medical LLM, a financial LLM) face a mixture optimization problem where the candidate datasets are heterogeneous domain corpora rather than general web/math/code splits. DeMix provides a principled way to answer questions like: "should we allocate more budget to case law texts or to legal textbooks?" or "does adding 10% general web data improve or degrade medical licensing exam performance?" The component model training cost scales with the number of candidate datasets (7–15 in the paper's experiments), which is typically manageable for domain-specific applications where the number of distinct data sources is limited. A legal LLM team with 5 candidate datasets (case law, statutes, law review articles, contracts, general English) could run DeMix at ~150B tokens total component training cost (5 × 30B) and produce an evidence-based mixture in a matter of weeks, replacing what would otherwise be months of manual proxy experimentation guided by intuition. The balanced multi-domain evaluation metric (macro-averaged ranking across general, domain-specific, and reasoning benchmarks) is particularly well-suited to this use case, where the goal is typically "good performance on the target domain plus acceptable performance on general language" rather than maximizing a single score.

Standardized benchmarking for public pre-training corpora. The DeMix Corpora release (Table 8) addresses a specific gap: the lack of "benchmarked corpora with validated data mixture ratios that can be directly reused for large-scale pre-training" (Section 1). For the research community, this means that future work on pre-training data curation, model architecture, or training algorithms can use DeMix Corpora as a controlled testbed where the data mixture is known to be near-optimal, eliminating data composition as a confounding variable. For example, a team comparing two model architectures could train both on DeMix Corpora with the validated mixture and attribute performance differences to architecture rather than data quality or mixture effects. This is analogous to the role that standardized benchmarks like ImageNet played in computer vision—not because the benchmark is perfect, but because it provides a common reference point that makes results comparable across papers. The DeMix Corpora's multi-domain composition (general, multilingual, math, code) and validated mixtures for three training stages make it particularly suitable for this role in the LLM pre-training literature, where the absence of standardized training data has made it difficult to isolate the effects of architectural innovations from data curation choices.

When to Prefer This Method

The paper positions DeMix against two named alternatives—training-based proxy methods (RegMix, CLIMB) and manual/heuristic mixture selection—and the experimental results in Tables 2 and 3 provide clear guidance on when each is appropriate:

Prefer DeMix over training-based proxies (RegMix/CLIMB) when: (1) the optimization target includes hard tasks like math and code that require substantial training to exhibit meaningful benchmark signals, making tiny-scale proxies (2B tokens) unreliable—Table 2 shows training-based proxies at 2B achieve Spearman's ρ of only 0.53 versus DeMix's 0.81 at comparable cost; (2) you need to evaluate more than ~50 mixture ratios to adequately cover the search space, since DeMix's marginal cost per additional proxy is near-zero while training-based methods scale linearly; and (3) you have the upfront compute budget to train component models at sufficient scale (~30B tokens each), since this fixed cost is amortized over the many proxy evaluations it enables. The primary tradeoff is that DeMix requires a larger upfront investment (212B tokens for 7 components at 30B each) compared to starting small with training-based proxies (224B tokens for 112 × 2B proxies), but produces substantially better mixtures for that investment.
Prefer training-based proxies when: (1) the candidate datasets are very small, making it impossible to train component models on 30B tokens without excessive repetition (the paper does not test this regime, but the additivity approximation would likely break down); (2) the optimization target is simple enough that tiny proxies provide adequate ranking signals—if you only care about general language capabilities and not math or code, the 2B-token proxies in Table 2 achieve reasonable general-benchmark ρ (0.60–0.94 depending on configuration) at lower total cost than DeMix's component training; or (3) you need exact absolute performance estimates rather than just relative rankings, since training-based proxies recover more absolute performance (capability recovery up to 0.87) than merged proxies (up to 0.85), per Table 2.
Prefer manual/heuristic mixture selection when: the number of candidate datasets is very small (2–3) and domain expertise provides strong priors about reasonable ratios, making the overhead of any automated method—DeMix or training-based—unnecessary. Table 3 shows that heuristic mixtures achieve rank 28.00–42.33, which is competitive with some automated configurations, suggesting that for simple mixture spaces, human intuition plus a few large-scale validation runs may suffice.

The paper does not explicitly articulate this decision framework, but it follows directly from the experimental comparisons. The key uncertainty—which the paper does not resolve—is whether DeMix's advantage persists when final model training budgets scale beyond the 50B tokens used for evaluation. At 500B+ token training scales, the relative importance of mixture optimization versus other factors may shift, and the 50B-token proxy rankings that DeMix relies on may not extrapolate. Practitioners training at very large scale should view DeMix as a promising but unvalidated approach for their regime.

🎯 Pitch

1. Executive Summary

2. Context and Motivation

The Core Problem: Data Mixture Optimization Is Both Critical and Prohibitively Expensive

Why Existing Approaches Fail: The Sufficiency-Accuracy-Efficiency Trilemma

The Gap in Existing Approaches

The Theoretical Foundation: Why Model Merging Can Work as a Proxy

How DeMix Positions Itself

The Practical Stakes

3. Technical Approach

3.1 Reader Orientation

3.2 Big-Picture Architecture (Diagram in Words)

3.3 Roadmap for the Deep Dive

3.4 Detailed, Sentence-Based Technical Breakdown

Dataset Preprocessing Pipeline

The Formal Data Mixing Objective and Component Model Training

Model Merging as Proxy: The Core Mechanism

Iterative Mixture Weight Optimization

Design Choices and Their Justifications

4. Key Insights and Innovations

Innovation 1: The Decoupling Paradigm — Why "Search" and "Training" Should Be Separate Resources

Innovation 2: The Small-Update Regime as an Enabling Condition, Not a Limitation

Innovation 3: Domain-Balanced Evaluation as a Constraint on Mixture Optimization

Innovation 4: Verifier Over-Optimization as a First-Class Phenomenon in Test-Time Scaling

5. Experimental Analysis

Evaluation Methodology

Main Quantitative Results

Proxy Consistency: DeMix versus Training-Based Proxies

Mixture Quality: Downstream Model Performance

Scalability to Larger Models

Ablation Studies and Robustness Checks

Critical Assessment

Claim 1: DeMix breaks the sufficiency-accuracy-efficiency trilemma, achieving high proxy accuracy at lower cost than training-based methods.

Claim 2: DeMix produces data mixtures that yield superior downstream model performance compared to existing methods.

Claim 3: The small-update regime (δ ≪ 1) is the enabling condition for DeMix, and the method works when δ is appropriately controlled.

Claim 4: The model merging proxy is a faithful approximation of real data mixture training, as measured by ranking consistency.

Missing Experiments and Baselines

Conditional Validity of Claims

6. Limitations and Trade-offs

The Difficulty Estimation Cost Is Unaccounted For in the Headline Efficiency Numbers

Hard Problems Remain Essentially Unsolved—Test-Time Compute Cannot Substitute for Missing Capability

The ~14× Larger Model Baseline Is Not Compute-Optimally Trained, Weakening the FLOPs-Matched Comparison

Verifier Over-Optimization Is a Hard Ceiling, Not a Solved Problem—The Compute-Optimal Policy Only Mitigates It

The Revision Model Has a 38% Correct-to-Incorrect Reversion Rate, Requiring Post-Hoc Patches Rather Than a Principled Solution

Single Benchmark, Single Model Family, and Small Test Set Limit Confidence in Generalization

7. Implications and Future Directions

How This Work Changes the Landscape

Follow-Up Research This Work Enables

Practical Applications and Downstream Use Cases

When to Prefer This Method

Adjacent in domain

Golden Goose: A Simple Trick to Synthesize Unlimited RLVR Tasks from Unverifiable Internet Text

Cosmos Policy: Fine-Tuning Video Models for Visuomotor Control and Planning

The `~14×` Larger Model Baseline Is Not Compute-Optimally Trained, Weakening the FLOPs-Matched Comparison