Pith. sign in

REVIEW 4 major objections 5 minor 2 cited by

Weighted averaging of separately trained component models preserves the benchmark ranking of models trained directly on the same data mixture, letting LLM pre-training data-mixture search run at a fraction of the compute.

Reviewed by Pith at T0; open to challenge. T0 means a machine referee read the full paper against a public rubric. the ladder, T0–T4 →

Challenge this review Re-run · record.json Download PDF Read on arXiv ↗

T0 review · deepseek-v4-flash

2026-08-03 05:54 UTC pith:37CAJHMW

load-bearing objection Clever proxy idea, but the final mixture is outside the representable subspace, so the central validation doesn't hold as written. the 4 major comments →

arxiv 2602.00747 v3 pith:37CAJHMW submitted 2026-01-31 cs.CL cs.AI

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

Shengrui Li , Fei Zhao , Kaiyan Zhao , Jieying Ye , Haifeng Liu , Fangcheng Shi , Zheyong Xie , Yao Hu
show 1 more author
Shaosheng Cao
This is my paper
classification cs.CL cs.AI
keywords data mixture optimizationmodel mergingLLM pre-trainingproxy modelsweighted parameter averagingbenchmark rankingpre-training corpusDeMix
verification ladder T0 review T1 audit T2 compute T3 formal T4 reserved

The pith

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

Choosing the right mix of pre-training data is one of the strongest levers on a language model's final abilities, but testing a mixture normally means training a whole model on it, which is too expensive for broad search. The paper proposes DeMix: train one component model per candidate dataset, then synthesize a proxy for any mixture by a weighted average of component weights. The central evidence is that these merged proxies preserve the ranking of reference models actually trained on those mixtures—macro-average Spearman rank correlation 0.81 at a total budget of about 212B tokens, versus 0.53 for 2B-token trained proxies at comparable cost. Because each proxy is a merge rather than a training run, the method can evaluate far more candidate mixtures and still find a better final mixture at lower compute. A 1.7B-parameter model trained on the selected mixture outperforms mixtures from trained-proxy search methods across general, math, and code benchmarks.

Core claim

The core discovery is that a weighted linear merge of component models—each trained once on one candidate data source mixed with 50% general data—behaves like a faithful proxy for a model trained directly on the corresponding weighted mixture. In the paper's evaluation, 96 random mixtures were scored both by merged proxies and by reference models trained on 50B tokens; the merged proxies reached macro-average Spearman ρ = 0.81 and top-25% ρ = 0.59 at a total search budget of 212B tokens, while 2B-token trained proxies under a comparable budget reached only 0.53 and 0.20. The paper also reports a capability recovery rate of up to 0.85, meaning the absolute benchmark scores of merged proxies a

What carries the argument

The load-bearing identity is approximate linear additivity of parameter updates: if training on a dataset changes model weights by a small delta (about 10% of total weight magnitude here), then the delta from training on a union of datasets is close to the sum of the individual deltas. Consequently, merging the individually trained component models with weights {α_i} approximates the parameters of a model trained on the same mixture {α_i}. Around this identity, DeMix wraps an iterative search: sample ratios uniformly, construct merged proxies, rank them on a benchmark suite, fit a regression predictor, and resample near the top-ranked ratios.

Load-bearing premise

The load-bearing premise is that a weighted average of separately trained component models approximates a model trained directly on that weighted mixture—so merged proxies rank mixtures the way real training would, even for domain-heavy mixtures far from the components' own 50%-general training distribution.

What would settle it

A direct test: sample several ratios with 70-90% domain weight, train 50B-token reference models on each, and compare merged-proxy rankings with the real trained-model rankings; if Spearman ρ in this region is much lower than the reported 0.81, the central claim fails where the final mixture lives.

Watch this falsifier — get emailed when new claim-graph text bears on it.

If this is right

  • Data-mixture search becomes essentially free after one training pass per candidate dataset: any number of sampled ratios can be evaluated by merging the same set of component models, so the number of search trials is no longer bounded by training budget.
  • Under a comparable token budget, merged proxies are far more reliable ranking signals than small trained proxies (macro ρ 0.81 vs 0.53), and the gap is especially stark on general benchmarks, where 2B-token proxies show near-zero correlation.
  • Increasing the number of merged proxies improves the quality of the final mixture up to a point: the paper observes rank improvement from 29.33 with 56 proxies to 24.00 with 224 proxies, and degradation at 448, suggesting overfitting to proxy noise.
  • Mixing each candidate dataset with at least 50% general data is necessary for proxy fidelity: dropping to 25% or 0% general data lowers Spearman ρ from about 0.79 to 0.67/0.65.
  • The released 22T-token corpus with validated mixtures can be used directly for pre-training, and a 1.7B model trained on it achieves better average benchmark rank than comparable mixed corpora.

Where Pith is reading between the lines

These are editorial extensions of the paper, not claims the author makes directly.

  • Editorial inference: the paper validates proxy fidelity inside the 50%-general subspace, but the final selected mixture (e.g., 21.8% general, 78.2% domain) sits far outside it, so the ranking transfer that matters for the final choice is an extrapolation, not a measured fact.
  • Editorial inference: the small-update additivity premise (δ≈10%) means merged proxies should become less faithful as components train longer or scale up; at frontier training scales, components may need to be re-trained more often, which would erode the claimed cost advantage.
  • Editorial inference: because the search objective is defined by a fixed benchmark suite, the chosen mixture is optimal for those benchmarks; a capability absent from the suite would not be optimized, and mixture choices might shift if new benchmarks are added.

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, simulated authors' rebuttal, and a circularity audit.

Referee Report

4 major / 5 minor

Summary. The paper proposes DeMix, a data-mixture search framework for LLM pre-training. Instead of training a new proxy model for every sampled mixture, DeMix trains seven component models, each on one candidate dataset mixed 50/50 with general data, and then synthesizes proxy models for arbitrary mixture weights by weighted linear merging. A LightGBM predictor is fit to benchmark rankings of these merged proxies and used, with iterative resampling, to select a final mixture, which is then validated by training a 1.7B model on 50B tokens. The paper reports that merged proxies achieve macro-average Spearman ρ=0.81 against 96 reference models at a total budget of ~212B tokens, and that the final DeMix mixture achieves a better macro-average rank (24.00) than RegMix and CLIMB baselines. The paper also releases DeMix Corpora, a 22T-token dataset with stage-wise validated mixtures.

Significance. If the proxy-fidelity claim were established, DeMix would be a meaningful advance: it decouples the cost of mixture search from the number of sampled mixtures, turning an expensive search into a fixed-cost component-training phase plus negligible merging/benchmarking. The paper is transparent about experimental settings, compares several merging methods, and releases code and data, which are concrete strengths. The central concern is that the proxy models and the final searched mixture live in different distributional subspaces; the paper's own appendix acknowledges that only mixtures with less than 50% non-general data are searchable, yet the headline final mixture is 78.2% non-general. This mismatch affects both the proxy-consistency validation and the attribution of the final mixture's performance to the search procedure. The issue is load-bearing but appears fixable within the manuscript's scope.

major comments (4)
  1. [Section 2.3, Eq. (5)-(6); Section 3; Appendix A.1] The derivation of the merged proxy assumes additivity of weight deltas for models trained on pure datasets: Δ(D_i ∪ D_j) ≈ Δ(D_i) + Δ(D_j), leading to M = Σα_iΘ_i approximating training on Σα_iD_i. However, each component model is trained on 0.5·D_i + 0.5·D_general (Section 3, β=0.5). Therefore M ≈ Θ_base + Δ(0.5·Σα_iD_i + 0.5·D_general), i.e., the proxy represents an affine-shifted mixture with at most 50% non-general data. This is not the mixture Σα_iD_i that the paper claims to optimize. Appendix A.1 concedes this: "Although we can only search within the subspace where non-general data accounts for less than 50%..." This concession should appear in the main text, and the derivation should be corrected or the search explicitly restricted to the affine subspace that merged proxies can represent.
  2. [Section 4.1; Section 3.3.1] The proxy-consistency evaluation samples 96 reference mixtures from the full 7-simplex. Under uniform sampling, the expected non-general share of a reference mixture is 6/7 ≈ 86%, while the merged proxies always contain at most 50% non-general data. Thus the reported ρ=0.81 measures rank correlation between two systematically different training distributions, not the fidelity of the mapping in Eq. (5)-(6) for the mixtures of interest. To support the claim that merging preserves ranking of 'models trained on real data mixtures', the evaluation should restrict the reference mixtures to the proxy-representable subspace (or otherwise correct for the affine shift), and report ρ separately within that subspace.
  3. [Table 3; Appendix C, Table 9] The final DeMix-224 mixture selected by the predictor has General=0.218 and non-general=0.782, i.e., 78.2% non-general data. This lies far outside the subspace of mixtures that merged proxies can represent (non-general ≤50%). The LightGBM predictor is trained on proxy scores from the ≤50% subspace and then used to select mixtures in the >50% region, but no out-of-subspace validation is provided. Consequently, the reported rank of 24.00 for the final mixture cannot be attributed to DeMix's search; it could reflect the extrapolation behavior of the regressor rather than the quality of the proxy signal. The authors should either validate extrapolation with reference models trained at high non-general shares, or change the search/proxy construction so that the final mixture lies inside the validated region.
  4. [Section 4.1, Capability Recovery] The Capability Recovery Rate compares average benchmark scores of the merged proxy and the reference model. Because the proxy's effective mixture is shifted toward general data (at least 50% general by construction), a high recovery rate may simply reflect that general benchmarks dominate the average, rather than that the proxy preserves the absolute performance of the reference mixture. At minimum, capability recovery should be reported per domain and conditioned on the actual mixture represented by the proxy, not on the nominal α.
minor comments (5)
  1. [Figure 1] The y-axis label "Proxy Accuracy (Spearman's)" is incomplete; it should read "Spearman's ρ".
  2. [Section 2.3] The paper states that δ is approximately 10% but does not show how this is measured or for which dataset/model. A brief empirical derivation or a reference to an appendix would help.
  3. [Appendix A.1] The concession about the searchable subspace is placed in the data-curation appendix. Since it directly limits the central claim, it should be moved or at least summarized in Section 2.3 or Section 3.
  4. [Table 8 / Section 5] The abstract and Section 5 describe "22T-token" DeMix Corpora, while Table 8 lists stage totals (14.4T, 5.8T, 2.1T) that sum to about 22T; this is consistent but should be stated explicitly in the main text to avoid confusion.
  5. [Table 2] For the DeMix rows, "Pre-Cost" is not clearly defined. The reader has to infer that it is the cost of training the base model; please define this column header in the table caption.

Circularity Check

0 steps flagged

No significant circularity; the core proxy-consistency claim is empirically validated against independently trained reference models, though the search subspace mismatch is a correctness concern, not a circularity.

full rationale

DeMix's core derivation chain is: (i) train component models on candidate datasets mixed 50/50 with general data; (ii) construct merged proxies by weighted linear combination of component weights (Eq. 6); (iii) validate that these proxies rank-order 96 separately trained 50B-token reference models (Spearman ρ, Table 2); (iv) fit a LightGBM predictor on the proxy benchmark scores to select a final mixture; (v) train an independent 50B-token model on that mixture and evaluate it relative to the same 96 reference models (Table 3). At no point is any parameter of the proxy or predictor fitted to the final model's benchmark scores or to the selected mixture's performance; the final evaluation is an external validation against independently trained reference models. The linearity assumption (Eqs. 4-5) is cited from prior external work (Wu et al. 2025; Lin et al. 2025b; Qin et al. 2022), not from the authors' own prior results, and is empirically tested here by the proxy-consistency experiments. The claimed main result (proxy ρ=0.81 at 212B total budget) is a direct measurement against independently trained reference models, not an artifact of self-citation or definition. The most serious concern — that the final DeMix-224 mixture (General=0.218, non-general=0.782 in Table 9) lies outside the subspace of mixtures representable by the merged proxies, since each component contains 50% general data so any merge has at least 50% general data (a limitation the paper itself concedes in Appendix A.1: 'we can only search within the subspace where non-general data accounts for less than 50%') — is a validity/generalization concern about extrapolation, not a circularity: the selected mixture is not defined in terms of the proxy scores, and the final model is independently trained. Accordingly, the circularity score remains low; the subspace mismatch should be treated as a correctness risk rather than circularity.

Axiom & Free-Parameter Ledger

5 free parameters · 4 axioms · 0 invented entities

The central methodological claim rests on two borrowed assumptions from task arithmetic (additivity of weight deltas and small updates) and on a subspace restriction induced by the 50%-general component training protocol. The paper does not introduce new theoretical entities; its contributions are empirical. The free parameters listed above are choices the authors make that materially affect proxy fidelity and the final mixture.

free parameters (5)
  • General-data ratio β in component training = 0.5
    Chosen to stabilize proxy fidelity; constrains search to non-general ≤ 50% of the effective mixture. Ablated in Table 5 (50% > 25% > 0%).
  • Component training token budget = 30B per component
    Selected after comparing 2B/10B/30B/50B in Table 2; 30B balances proxy accuracy and cost.
  • Number of candidate datasets N = 7
    Result of data curation; dimensionality of the mixture simplex; affects search difficulty.
  • Number of proxy evaluations = 112/224 (chosen)
    Table 3 shows rank improves from 56 to 224 proxies, declines at 448; the final recommendation uses 224. This is a hyperparameter of the search.
  • LightGBM hyperparameters = lr=0.02, 300 iterations
    Adopted from RegMix baseline; not ablated.
axioms (4)
  • domain assumption Weight-delta additivity: Δ(D_i ∪ D_j) ≈ Δ(D_i) + Δ(D_j) for small updates (Eq. 5).
    Borrowed from Qin et al. 2022, Wu et al. 2025, Lin et al. 2025b; the paper measures δ≈10%, which is asserted to be small but not rigorously justified.
  • domain assumption Small-update assumption δ≪1 (Eq. 4).
    Needed for linearity of weight deltas; measured at ~10%, which is borderline.
  • domain assumption Benchmark scores of merged models rank-order the quality of trained models on the same mixture.
    Tested empirically in Table 2, but the regression search extrapolates to mixtures beyond the 96 reference samples.
  • domain assumption The subspace where non-general data ≤ 50% contains a mixture competitive with the global optimum.
    Stated in Appendix A as a design choice; not empirically verified against mixtures with >50% domain data.

pith-pipeline@v1.3.0-alltime-deepseek · 19897 in / 19768 out tokens · 205408 ms · 2026-08-03T05:54:22.903058+00:00 · methodology

0 comments
read the original abstract

Determining an effective data mixture is a key factor in Large Language Model (LLM) pre-training, where models must balance general competence with proficiency on hard tasks such as math and code. However, identifying an optimal mixture remains an open challenge, as existing approaches either rely on unreliable tiny-scale proxy experiments or require prohibitively expensive large-scale exploration. To address this, we propose Decouple Searching from Training Mix (DeMix), a novel framework that leverages model merging to predict optimal data ratios. Instead of training proxy models for every sampled mixture, DeMix trains component models on candidate datasets at scale and derives data mixture proxies via weighted model merging. This paradigm decouples search from training costs, enabling evaluation of unlimited sampled mixtures without extra training burden and thus facilitating better mixture discovery through more search trials. Extensive experiments demonstrate that DeMix breaks the trade-off between sufficiency, accuracy and efficiency, obtaining the optimal mixture with higher benchmark performance at lower search cost. Additionally, we release the DeMix Corpora, a comprehensive 22T-token dataset comprising high-quality pre-training data with validated mixtures to facilitate open research. Our code and DeMix Corpora is available at https://github.com/Lucius-lsr/DeMix.

Figures

Figures reproduced from arXiv: 2602.00747 by Fangcheng Shi, Fei Zhao, Haifeng Liu, Jieying Ye, Kaiyan Zhao, Shaosheng Cao, Shengrui Li, Yao Hu, Zheyong Xie.

Figure 1
Figure 1. Figure 1: Methods such as RegMix and CLIMB require extensive proxies. Scaling each proxy leads to unaffordable overall budget. While our DeMix only require a few component models to merge unlimited training-free proxies. largely driven by the massive pre-training (Gemini Team et al., 2023; Achiam et al., 2023). Beyond scale alone, the composition of the pre-training corpus plays a critical role in shaping model capa… view at source ↗
Figure 2
Figure 2. Figure 2: Pipeline for DeMix. After (1) cleaning and categorizing massive data, (2) component models are trained on individual candidate datasets. Instead of large-scale training for every ratio, (3) weighted model merging serves as a computationally efficient proxy to estimate performance for various mixture ratios. Finally, (4) a predictor is trained on the benchmarked proxy models to regress the relationship betw… view at source ↗
Figure 3
Figure 3. Figure 3: General performance of high-quality general datasets. DOLMA (Soldaini et al., 2024), and Nemotron-Pretrain (Basant et al., 2025). In addition, multilingual (Messmer et al., 2025; De Gibert et al., 2024), code (Li et al., 2023), and mathematics (Allal et al., 2025a) corpora are not di￾rectly applicable for pre-training. As shown in [PITH_FULL_IMAGE:figures/full_fig_p007_3.png] view at source ↗
Figure 4
Figure 4. Figure 4: The data mixture constructed using our DeMix frame￾work. The three hierarchical levels from the inside out are domain, data category, and data origin. In addition, DeMix Corpora not only features large and high￾quality general corpus subset, but also adopts a validated optimal data mixture that best balances multi-domain capa￾bilities of pre-training. The final mixtures are illustrated in [PITH_FULL_IMAGE… view at source ↗
Figure 5
Figure 5. Figure 5: Data mixtures for the three stages of pre-training in the DeMix Corpora, with approximately 14T, 6T, and 2T tokens, respectively. The three hierarchical levels from the inside out are domain, data category, and data origin. based on quality scores. In contrast, during Stages 2 and 3, data quality directly influences the final training performance; thus, we employ the DeMix framework to optimize the data mi… view at source ↗

Review history (2 revisions) →

discussion (0)

Sign in with ORCID, Apple, or X to comment. Anyone can read and Pith papers without signing in.

Forward citations

Cited by 2 Pith papers

Reviewed papers in the Pith corpus that reference this work. Sorted by Pith novelty score.

  1. RegMix-D: Dynamic Data Mixing via Proxy Training Trajectories

    cs.CL 2026-06 unverdicted novelty 6.0

    RegMix-D fits regression models to proxy loss trajectories to produce dynamic data mixture schedules that outperform static RegMix and DoReMi on 25B-token Pile pretraining with a 1B model.

  2. Towards a Data-Parameter Correspondence for LLMs: A Preliminary Discussion

    cs.LG 2026-04 unverdicted novelty 4.0

    A data-parameter correspondence unifies data-centric and parameter-centric LLM optimizations as dual geometric operations on the statistical manifold via Fisher-Rao metric and Legendre duality.

Reference graph

Works this paper leans on

2 extracted references · 2 linked inside Pith · cited by 2 Pith papers

  1. [451]

    emnlp-main.451/

    URL https://aclanthology.org/2022. emnlp-main.451/. Sakaguchi, K., Bras, R. L., Bhagavatula, C., and Choi, Y . Winogrande: An adversarial winograd schema challenge at scale.Communications of the ACM, 64(9):99–106, 2021. Sap, M., Rashkin, H., Chen, D., LeBras, R., and Choi, Y . Socialiqa: Commonsense reasoning about social interac- tions.arXiv preprint arX...

    Pith/arXiv arXiv 2022

  2. [2017]

    ISBN 9781510860964

    Curran Associates Inc. ISBN 9781510860964. Kimi Team, Bai, Y ., Bao, Y ., Chen, G., Chen, J., Chen, N., Chen, R., Chen, Y ., Chen, Y ., Chen, Y ., et al. Kimi k2: Open agentic intelligence.arXiv preprint arXiv:2507.20534, 2025. Lee, K., Ippolito, D., Nystrom, A., Zhang, C., Eck, D., Callison-Burch, C., and Carlini, N. Deduplicating train- ing data makes l...

    Pith/arXiv arXiv 2025