LinMix: Data Mixing as Linear Domain Attribution

AI summary

Overview

Research area: Machine learning, specifically language-model pretraining data curation ("data mixing"), data attribution, and optimal experimental design.

Technical level: Intermediate. The core idea is intuitive (sample extreme mixes, fit a straight line), but the paper includes formal results (linear-programming optimality, error bounds for Dirichlet sampling) and assumes familiarity with regression and simplex constraints.

Scope: The paper reframes regression-based data mixing as a domain attribution problem and introduces LinMix, a hyperparameter-free vertex-sampling scheme that makes an interpretable linear model viable for selecting pretraining data mixtures across 24 and 288 domains.

What This Paper Is About

Pretraining corpora are split into domains (topics, formats, code, math), and the fraction of the token budget given to each domain strongly affects downstream model quality. A popular approach trains many small proxy models on randomly sampled mixtures, fits a regression to predict performance from the mixture, and then optimizes that regression to pick a target mix. The authors show that the standard way of sampling those mixtures (Dirichlet sampling around the prior corpus distribution) is a poor experimental design that needs per-setup tuning and breaks down entirely at hundreds of domains, and they replace it with sampling at the vertices of the capped simplex so that a simple, interpretable linear model suffices.

Key Contributions

1. A reframing of data mixing as domain attribution. The paper formalizes regression-based data mixing as a domain-level datamodel, where the coefficients of a linear fit measure each domain's marginal contribution to the target metric, connecting data mixing to the data-attribution literature.

2. LinMix, a hyperparameter-free vertex sampler. Proxy mixes are drawn at the vertices of the capped simplex using random domain permutations that greedily fill each domain to its cap, paired with reverse-order permutations as antithetic counterparts. This decorrelates domain effects and needs no Dirichlet scale or temperature tuning.

3. Closed-form and provable mix optimization. With a linear surrogate, selecting the target mix becomes a linear program whose solution is itself a vertex: domains with positive estimated attribution are filled to their caps in decreasing order of coefficient. Proposition 1 and Corollary 1 extend this to an isoelastic objective with curvature r, giving a closed-form capped water-filling allocation.

4. Theory for why vertex sampling wins. Theorem 1 shows Dirichlet sampling inflates attribution error when its mixes concentrate around the prior or the prior becomes skewed; Theorem 2 establishes the optimality of vertex sampling in expectation for equal caps, with Proposition 3 covering the uncapped case.

Main Findings

• Dirichlet sampling fails to scale to fine-grained domains. At 288 domains, the baseline Dirichlet design does not converge with the default hyperparameters (τ = 2, γ_lo = 0.1, γ_hi = 5), nor even with γ_hi = 10, because sampled mixes violate the per-domain cap and are rejected. The authors hand-tune to γ_lo = 40 and γ_hi = 4000 by minimizing the condition number of the design matrix.

• No surrogate fits a Dirichlet design at 288 domains. Across 1024 × 60M proxy models, the best cross-validated Spearman correlation for Dirichlet sampling is ρ ≤ 0.39, whereas vertex sampling brings the linear model on par with the best nonlinear alternative. At 24 domains, every surrogate fits well under both samplers.

• Half the proxy budget is enough under vertex sampling. Half of the vertex swarm already fits better than the full Dirichlet swarm, because near-prior Dirichlet mixes vary HellaSwag BPB five times less than vertex mixes, leaving little signal to regress on.

• Linear attributions transfer across proxy scales. A linear model fit on a 30M-model swarm ranks the outcomes of 60M and 100M models at ρ ≥ 0.95 on 288 domains, while the Dirichlet fit does not transfer (ρ ≤ 0.42).

• 1B OLMo2 target results. Across 23 × 1B OLMo2 target models pretrained on 32B tokens, LinMix improves the 288-domain setting over the best nonlinear baseline: ARC accuracy .511 vs .449 (+13.8%), HellaSwag accuracy .615 vs .547 (+11.8%), HellaSwag + PIQA accuracy .690 vs .638 (+7.6%), and CPile BPB .834 vs .878. At 24 domains it is on par or better (ARC accuracy .510 vs .496, +2.8%). The abstract summarizes this as an 11% average accuracy gain over the best nonlinear baseline at 288 domains.

• 4B Qwen3.5 results on code and math. Across 10 × 4B Qwen3.5 target models pretrained on 80B tokens over 40 domains, LinMix reaches MBPP P@1 of .444 vs .304 for the best baseline (+46.1%) and GSM8K EM of .337 vs .294 (+14.4%), with gains carrying over to held-out HumanEval and MATH and to HellaSwag accuracy (.696 vs .667) and CPile BPB (.797 vs .801 for Olmix).

• Different targets produce visibly different mixes. RegMix yields similar mixes across HellaSwag, MBPP and GSM8K; Olmix differentiates less than LinMix, which downweights math and code domains for language modeling and upweights programming/software domains for coding and math/science domains for math.

• Utility curvature changes support, not the ranking. On 24 domains, using ln π for HellaSwag adds Social Life, Fashion & Beauty, Games and Industrial to the four domains at their caps (Health, Home & Hobbies, Food & Dining, Adult); using √π for CPile adds a sixth, Hardware, domain.

Methodology in Plain English

The workflow has three steps.

1. Choose which mixtures to test. Instead of drawing random mixtures clustered around the existing corpus proportions, LinMix draws a random ordering of the domains, then fills them one by one to their per-domain cap until the token budget runs out. The result is a mix where each domain is either at its maximum allowed share or absent, with at most one partially filled domain. Runs are paired: the next mixture uses the reversed ordering, so domains left empty by one fill are saturated first by its counterpart. This maximizes the spread of the design (the D-optimality criterion from classical experimental design) without any tunable parameters.

2. Fit a straight line. Each proxy model is trained on one mixture and evaluated, producing a score (e.g., bits-per-byte, negated so higher is better). A plain least-squares linear model maps mixture weights to that score; the coefficients, centered to zero mean, are the domain attributions.

3. Pick the mix. Maximizing the linear model over the capped simplex is a linear program solved by filling domains to their caps in decreasing order of coefficient. Optionally, the authors apply a concave transform (from linear through √π to ln π, indexed by curvature r) to encode diminishing returns from repeating a domain's data; this changes how the budget is spread among positively attributed domains but not which domains are included.

The empirical setup uses the 200B-token WebOrganizer corpus split into 2048-token sequences, with 24 coarse topic domains and 288 fine-grained topic–format domains covering 94% of the corpus. Proxy models are 30M and 60M OLMo 2 models (512 and 1024 runs respectively per sampler), trained past the 20 tokens-per-parameter Chinchilla ratio. A second setup extends the corpus with Nemotron-CC-Math and Stack-Edu (split into 15 programming languages) to a 416B-token, 40-domain corpus, with 512 × 100M proxy models and 4B Qwen3.5 targets trained on 80B tokens.

Why This Matters

Research impact. The paper argues that the sampling design, not the surrogate model, is the bottleneck in regression-based data mixing. That shifts attention from ever-more-flexible regressors to experimental design, and it links data mixing to the data-attribution and datamodel literatures. It also supplies negative results about Dirichlet sampling that explain why practitioners have had to hand-tune it, plus a proof that an interpretable linear model can be sufficient.

Real-world applications:

• Pretraining corpus curation for foundation models, where deciding how much CommonCrawl, code, math or academic text to include directly affects reasoning and coding ability.
• Compute-budget planning for data-mixing studies, since vertex sampling reaches a usable fit with half the proxy swarm that Dirichlet sampling needs.
• Domain-level auditing of a training corpus, because the linear coefficients give a directly readable ranking of which domains help or hurt a target capability.
• Targeted capability tuning, e.g. deliberately upweighting programming-language and math domains for code or math benchmarks rather than retraining from scratch on a new corpus.

Industry relevance. Teams that already run RegMix-style pipelines can swap the sampler and the surrogate without new infrastructure, and the method removes the per-setup tuning that makes scaling to hundreds of fine-grained domains impractical. The removal of hyperparameter search is the practical selling point; the paper does note that the cost of training a proxy swarm and the use of a single training run per target model are limitations, but it does not report wall-clock or dollar figures.

Future Directions

• Nonlinear models on vertex-sampled designs. The authors state that optimizing a vertex-sampled matrix with a nonlinear model could yield better target performance, and leave this to future work.
• Unequal caps. Theorem 2 establishes optimality of vertex sampling in expectation only for equal caps; the unequal-cap case is explicitly left open.
• Combining vertex sampling with adaptive search. LinMix fixes all proxy mixtures before observing outcomes; the paper suggests merging the design with adaptive or Bayesian allocation of new proxy runs to promising mixtures.
• Cheaper and more robust estimation. The paper reports that scaling the linear coefficients by their standard error could improve performance, and flags proxy-swarm cost and single-run-per-target-model evaluation as the main limitations to address.

Target Audience

Machine learning researchers and engineers working on pretraining data curation, data mixing, or data attribution; practitioners running proxy-model pipelines who need a method that scales to hundreds of domains without hyperparameter tuning; and readers interested in how classical optimal experimental design applies to modern language-model training decisions. Some comfort with linear regression, constrained optimization and the simplex is helpful.

Authors’ abstract

Language model performance depends strongly on the mixture of domains in the pretraining data corpus. Popular data mixing methods select these domain weights by training many small proxy models on random mixes and then fitting a regression model to predict performance from the mix. We recognize this practice as a domain attribution problem, where the regression measures each domain's contribution to the target metric. Through this lens, the standard Dirichlet sampling of proxy data mixes is a poor regression design that requires tuning and cannot support a linear fit for hundreds of domains. Building on classical optimal experimental design, we introduce LinMix, which samples mixes at the vertices of the capped simplex so that each domain is at its cap or absent. This hyperparameter-free vertex sampling decorrelates domain effects, enabling an interpretable linear model whose coefficients are the domain attributions. The optimal mix given by LinMix provably keeps only domains with positive attribution values and requires no search over candidate mixes. We further establish the optimality of vertex sampling, and prove that Dirichlet sampling inflates the attribution error when its mixes concentrate around the prior data mix. Pretraining experiments with 1B-parameter OLMo2 and 4B Qwen3.5 language models show that LinMix outperforms Dirichlet-based data mixing across reasoning, code, and math benchmarks, with relative gains of 46% on MBPP and 14% on GSM8K. At 288 fine-grained domains, only vertex sampling supports a linear fit, and LinMix improves benchmark accuracy by 11% on average over the best nonlinear baseline.