PaperlayerView exact revision on arXiv

Source-linked AI summary

LinMix: Data Mixing as Linear Domain Attribution

Hubert Baniecki, Przemyslaw Biecek, Martin Pawelczyk

arXiv:2610.37390v1cs.LGcs.CLstat.ML

TL;DR

Regression-based data mixing must choose how to sample proxy mixes, and existing Dirichlet designs need tuning and struggle to support linear fits at hundreds of fine-grained domains. LinMix samples capped-simplex vertices to estimate interpretable domain attributions and optimize target mixes. Experiments report stronger performance than Dirichlet-based approaches across reasoning, code, and math benchmarks, with linear fits supported at 288 domains.

  • Problem

    Regression-based data mixing lacks a well-understood proxy-mix design and popular Dirichlet methods require heuristic tuning that becomes difficult at hundreds of fine-grained domains.

    Source: §1 · p. 2

  • Method

    LinMix samples proxy mixes at capped-simplex vertices and fits an interpretable linear model whose coefficients estimate domain contributions for target-mix optimization.

    Source: §1 · p. 2, §3 · p. 3

  • Results

    LinMix outperforms Dirichlet-based mixing across reported reasoning, code, and math benchmarks, and at 288 domains vertex sampling supports a linear fit while Dirichlet does not fit well even when tuned.

    Source: §1 · p. 2, §4.2 · p. 5, §4.3 · p. 6

  • Takeaways & Limitations

    The paper establishes vertex sampling as theoretically optimal and shows that Dirichlet sampling inflates attribution error when mixes concentrate around the prior.

    Source: §1 · p. 2, §5 · p. 7

    The approach requires training a proxy swarm, and experiments use a single random repetition for each target model.

    Source: F Limitations · p. 30

Abstract

from arXiv · show

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.

1 Introduction

Data mixing chooses domain weights for pretraining, but regression-based methods leave the proxy-mix distribution and scalability to many fine-grained domains unresolved. LinMix frames mix selection as domain attribution and introduces vertex sampling, with experiments and theory supporting the method.

1 Introduction

  • Regression-based data mixing trains proxy models on sampled mixes, then optimizes a fitted regression model to predict performance from the mix.The paper frames this procedure as a domain-attribution problem.Source: §1 · pp. 1–2
  • The proxy-mix distribution is a regression-design choice: sampled mixes determine how much a fixed proxy-run budget reveals about each domain’s effect.The authors also identify scaling to hundreds of fine-grained domains as a challenge for popular methods, whose hyperparameters require per-setup heuristic tuning.Source: §1 · p. 2

Domain attribution

  • A linear fit to proxy runs yields a domain-level datamodel whose coefficients measure each domain’s marginal contribution to the target outcome.The paper relates this framing to data attribution, where model behavior is attributed to training data.Source: §1 · p. 2

Contributions

  • LinMix samples proxy mixes at capped-simplex vertices, enabling an interpretable linear model to select an optimal pretraining mix.The paper reports experiments with 1B-parameter OLMo2 and 4B Qwen3.5 models and proves vertex sampling optimal.Source: §1 · p. 2

2 Background

Regression-based data mixing uses proxy-model outcomes to fit a performance predictor, then optimizes that predictor over a capped simplex. Existing approaches commonly sample proxy mixes from a tempered-prior Dirichlet distribution.

Background

  • The corpus prior sets each domain’s natural proportion, while data mixing seeks target weights that maximize downstream performance.A target-scale evaluation of even one candidate mix would require a full pretraining run, motivating the use of smaller proxy models.Source: §2 · p. 2

Sampling Π

  • REGMIX and OLMIX populate proxy design matrices by sampling mixes from a tempered version of the prior using a Dirichlet distribution.The Dirichlet scale γ controls concentration: small values produce sparse mixes, while large values produce tight perturbations around the tempered prior.Source: §2 · pp. 2–3

Optimizing πtarget

  • Per-domain caps restrict target weights to limit repeated use of domain data; the cap is set as a multiple of the prior share, subject to a maximum of one.The cited approaches use different choices for the up-sampling factor κ.Source: §2 · p. 3
  • Each proxy model is trained on a candidate domain mix and evaluated; regression predicts outcomes from the mixes, and optimization selects a target mix within per-domain caps.The candidate mixes form the rows of design matrix Π.Source: §2 · pp. 2–3
  • REGMIX approximately optimizes a gradient-boosted predictor over sampled candidate mixes, whereas OLMIX uses an exponential mixing law that is convex in the mix.The excerpt states that OLMIX’s convex predictor can be minimized exactly over the capped simplex.Source: §2 · p. 3
  • The design matrix determines how well regression coefficients can be identified from the limited proxy runs, but mix sampling has received little attention.This paper addresses the design of Π.Source: §2 · p. 3

3 Data Mixing as Linear Domain Attribution

LinMix treats data mixing as domain attribution and changes the proxy-run design to support an interpretable linear fit. It samples capped-simplex vertices and uses the resulting attributions to optimize the target mix.

3 Data Mixing as Linear Domain Attribution

  • LinMix changes the proxy design to maximize information about domain effects, allowing a linear fit where Dirichlet-sampled runs cannot support one, especially at hundreds of domains.The method samples vertices of the capped simplex and reports that nonlinear models offer no better fit after this design change.Source: §3 · p. 3

Sampling Π

  • Vertex sampling aims to maximize det(Π⊤Π) by drawing proxy mixes at capped-simplex vertices, without sampling hyperparameters.For each run, domains are greedily filled to their caps in a uniformly random order until the unit budget is exhausted.Source: §3 · p. 3
  • Each vertex mix has at most one interior coordinate; pairing runs with reversed domain orders creates largely disjoint supports and keeps Π⊤Π well conditioned.The paired fills contribute near-orthogonal directions to the design matrix.Source: §3 · p. 3

Optimizing πtarget

  • Least squares estimates domain attributions from proxy outcomes, with covariance governed by (Π⊤Π)−1; coefficients are centered so each measures value relative to the average.For metrics where lower is better, outcomes are negated so positive attribution indicates above-average value.Source: §3 · pp. 3–4
  • The linear target-mix objective is optimized by filling domains to their caps in decreasing order of estimated attribution until the budget is exhausted.The solution is itself a vertex of the capped simplex.Source: §3 · p. 4

Transforming the utility of πk

  • The utility of a domain weight can be transformed to represent diminishing returns from repeating that domain’s data.A logarithmic utility makes each additional sequence worth less than the previous one; the linear and logarithmic cases are framed as isoelastic-objective extremes.Source: §3 · p. 4

Domain support and uniqueness

  • For positive-curvature objectives, the optimal mix assigns zero weight to domains with nonpositive estimated attribution, provided positive-attribution domains’ caps can absorb the budget.The optimal mix is unique under the stated condition for r in (0, 1].Source: §3 · p. 4

Closed form of the optimal mix

  • The optimal mix has a closed form on the positive-attribution support; as r approaches zero, it becomes a greedy fill ordered by decreasing attribution.In that limit, optimal mixes saturate top-domain caps, so the predicted mix interpolates rather than extrapolates from the sampled vertices.Source: §3 · p. 4

4 Experiments

Experiments compare proxy-fit quality and target-model performance across domain granularities and model scales. LinMix’s vertex sampling avoids Dirichlet tuning difficulties and supports linear fits, while benchmark comparisons favor LinMix in the reported settings.

4 Experiments

  • The experiments use WebOrganizer at 24 topic domains and 288 topic–format domains, the latter covering 94% of the corpus.The 1B OLMo2 setup uses 30M and 60M proxies and sets the up-sampling cap so no sequence repeats at target scale.Source: §4 · p. 4
  • Targets include ARC, HellaSwag, PIQA, and CPile, with proxy BPB used to optimize mixes and held-out target-model performance reported.ARC uses its test set; HellaSwag and PIQA use validation sets because their test labels are withheld.Source: §4 · p. 5

4.1 Dirichlet sampling of proxy data mixes requires tuning hyperparameters

  • At 288 domains, Dirichlet sampling fails to converge with default hyperparameters, prompting a condition-number-based choice of γlo = 40 and γhi = 4000.The selected mixes are much closer to the prior; the passage says related work gives no guidance for tuning these hyperparameters.Source: §4.1 · p. 5, Figure 2 · p. 5

4.2 Vertex sampling outperforms Dirichlet, enabling a linear model fit

  • At 288 domains, no surrogate fits Dirichlet-sampled runs well even after tuning (ρ ≤0.39 at m = 1024), while vertex sampling lets a linear model match the best nonlinear alternative.Half the vertex swarm fits better than the full Dirichlet swarm; near-prior Dirichlet mixes vary HellaSwag BPB five times less than vertex mixes.Source: §4.2 · p. 5

4.3 LINMIX outperforms data mixing baselines on a 1B target model benchmark

  • LinMix with a linear model significantly improves 288-domain data mixing and is on par or better with 24 domains on a 1B target benchmark.For CPile validation loss, the linear fit under vertex sampling is on par with or better than nonlinear regression with Dirichlet sampling.Source: §4.3 · p. 6

4.4 Scaling to a 4B Qwen3.5 model benchmarked on code and math capabilities

  • In the 4B Qwen3.5 experiments, LinMix delivers strong performance across language, coding, and math benchmarks and transfers capabilities to held-out HumanEval and MATH benchmarks.The setup uses a 416B-token corpus spanning 40 domains and target models trained on 80B tokens.Source: §4.4 · p. 6

5 Theoretical Analysis

The theoretical analysis frames regression-based mixing as domain attribution and uses experimental-design theory to explain why vertex sampling is informative. It characterizes the error of Dirichlet designs and the optimality scope of vertex sampling.

5 Theoretical Analysis

  • LINMIX shifts domain weights by target: it downweights math and code for HellaSwag, and upweights programming, software, math, or science for MBPP and GSM8K.These are the domain-weight patterns described for the respective benchmarks.Source: Figure 4 · p. 6
  • Domain-attribution coefficients average sample-level attribution scores, linking regression-based data mixing to linear datamodels and cooperative-game values.The paper relates this framework to Shapley and Banzhaf data values.Source: §5 · p. 6
  • Dirichlet sampling's attribution error grows as concentration increases and mixes approach the prior, while capped sparse mixes are rejected at many domains.At 288 domains, no surrogate fits the Dirichlet design well; vertex sampling remains best in expectation under equal caps.Source: §5 · p. 7

6 Conclusion

The conclusion reports that an interpretable linear model suffices under vertex sampling, while noting proxy-training cost and limited target-model runs. Related work situates LINMIX among experimental design, mixing, attribution, and data-selection methods.

6 Conclusion

  • The conclusion says an interpretable linear model suffices for vertex-sampled data mixing and identifies proxy-swarm cost and one target-model run as key limitations.The authors suggest nonlinear optimization of the vertex-sampled matrix as a possible route to higher target performance.Source: §6 · p. 7

A Related Work

  • The paper grounds vertex sampling in D-optimal experimental design, which maximizes det(Π⊤Π) for first-degree mixture models.For equal caps, the paper proves its random-permutation vertex scheme optimal in expectation; unequal caps remain future work.Source: A Related Work · p. 15
  • LINMIX fixes proxy mixtures before observing outcomes, unlike adaptive-search methods that allocate runs iteratively.The fixed design supports a useful linear fit across hundreds of domains; combining it with adaptive search is left for future work.Source: A Related Work · p. 15
  • LINMIX's allocation rule connects to prior ranking-based methods and to capped water-filling for positively attributed domains under isoelastic utility.For r = 0, the framework recovers cap saturation in rank order; for r > 0, the utility maximizer assigns positive weight to every positively attributed domain.Source: A Related Work · p. 15
  • LINMIX extends datamodel-style attribution from individual examples to domain proportions estimated from proxy runs, with larger-model use depending on cross-scale transfer.The related work notes that transfer depends on proxy-training hyperparameters.Source: A Related Work · p. 16
  • The paper distinguishes domain mixing from example-level selection and filtering, which can be combined by curating the corpus before optimizing its domain mix.Related work covers selection methods that combine quality and diversity criteria or account for interactions between examples.Source: A Related Work · p. 16

B Proofs

For r in (0, 1], the optimal mix is unique and assigns weight only to domains with positive estimated attribution. Its closed form allocates the budget across that support subject to domain caps.

B Proofs

  • The unique optimal mix assigns zero weight to every domain with nonpositive estimated attribution and positive weight to every positively attributed domain.These conclusions assume the positively attributed domains’ caps can absorb the full budget; the positive-weight result applies for r in (0, 1].Source: B Proofs · pp. 16–17
  • The optimum exists even at r = 1, where the objective is log-linear on the positive-attribution support and assigns −∞ utility to zero weight there.A feasible mix proportional to the caps provides positive weight to every domain in that support.Source: B Proofs · p. 16
  • At r = 0, the linear objective is maximized at a truncated-simplex vertex, and greedy allocation can leave lower-attribution domains in the positive set at zero weight.With estimated attributions (2, 1, −3) and unit caps, the optimum is (1, 0, 0).Source: B Proofs · p. 17
  • For r in (0, 1], the optimal weight on each positively attributed domain is the minimum of its cap and an attribution-scaled allocation set by a unique budget-normalizing scalar.The common scalar equalizes marginal utilities except where a domain cap binds.Source: B Proofs · p. 17

C Experimental Setup

The experiments use topic-based and finer topic–format domains, alongside a 40-domain setup that adds code and math data. Proxy and target models are evaluated on benchmark and held-out data.

C Experimental Setup

  • WebOrganizer yields 195B tokens; its prior data mixes include 24 topic domains and 288 topic–format domains, while the larger half of 576 pairs covers 94% of the corpus.The 40-domain corpus adds 168B Stack-Edu tokens and a 53B-token highest-quality Nemotron-CC-Math subset, bringing total data size to 416B tokens.Source: C Experimental Setup · p. 17, Figure 5 · p. 18
  • The experiments train language models on tokenized, single-domain sequences of 2048 tokens and use a 32B-token target budget to set domain caps.The caps are 5.98 and 5.63 at 24 and 288 domains, and 3.40 at 40 domains.Source: C Experimental Setup · p. 17
  • Proxy models are scored by BPB on benchmark training splits, while target models are scored by accuracy and BPB on held-out sets.For the Common Pile, separate 150k-sequence and 300k-sequence held-out splits are used for fitting proxy surrogates and scoring target models, respectively.Source: C Experimental Setup · p. 18, Table 4 · p. 18

D Additional Results

Additional results compare vertex and Dirichlet sampling, examine utility transformations, and report transfer across proxy scales. The theoretical analysis establishes vertex-sampling optimality and conditions under which Dirichlet sampling increases attribution error.

D Additional Results

  • At 24 domains, the π utility concentrates weight on four domains at their caps, while √π or ln π spreads weight for the 32B-token target budget.A larger budget would spread the π mix through lower caps; coefficient standard-error scaling adds at most half a point of HellaSwag accuracy.Source: D Additional Results · p. 19, Table 5 · p. 21, Table 6 · p. 21
  • At 40 domains, vertex sampling reaches ρ ≥ 0.94 with every surrogate at m = 512, while tuned Dirichlet sampling stays at ρ ≤ 0.68 on three benchmarks and 0.80 on CPile.The 4B Qwen3.5 proxy experiments evaluate HellaSwag, MBPP, GSM8K, and held-out Common Pile loss.Source: Figure 7 · p. 19, D Additional Results · p. 19
  • At 288 domains, the linear fit from vertex-sampled mixes transfers its ranking across proxy scales, whereas the Dirichlet fit’s ranking does not.The comparison fits HellaSwag BPB on 30M proxy models and predicts outcomes for 60M and 100M proxy swarms.Source: Figure 8 · p. 20

E Theoretical Analysis

  • Vertex sampling is optimal on the simplex and, in expectation, on the simplex with equal caps; Dirichlet attribution error grows as concentration rises and can grow with domain count under a more-skewed prior.The stated domain-count condition is when the prior is more skewed than the Dirichlet temperature can flatten.Source: E Theoretical Analysis · p. 22

E.1 From datamodels to domain attribution

The section derives domain-level attribution from linear datamodels and explains how proxy-mix design controls its estimation error. Domain coefficients aggregate sequence-level attribution, while the design matrix determines how reliably they can be estimated.

E.1 From datamodels to domain attribution

  • The expected proxy-run outcome is linear in its domain mix, with each domain coefficient equal to the proxy budget times its mean sequence-level attribution.Evaluating the linear relation at a single-domain mix identifies that domain coefficient as the expected outcome of training on that domain alone.Source: §E.1 · p. 22
  • Domain attribution interpolates between sequence-level data attribution at the finest partition and one undifferentiated corpus at a single domain.The domain partition affects sequence-selection noise; finer partitions can require more proxy runs because estimation error grows with the number of domains.Source: §E.1 · p. 23

E.2 The design matrix governs the estimation error

  • The least-squares attribution estimate is unbiased, and its error is bounded by noise level and the inverse design information, especially the smallest eigenvalue of the normalized Gram matrix.A small eigenvalue signals a domain direction that varies little across proxy runs; the condition number controls the bound up to a factor of d.Source: §E.2 · p. 24

E.3 Pitfalls of Dirichlet sampling

Dirichlet designs can estimate domain attributions poorly when mixes cluster near the prior or undersample rare domains. Vertex designs address the design problem, with optimality guarantees on the simplex and, under equal caps, the capped simplex.

E.3 Pitfalls of Dirichlet sampling

  • For Zipfian priors, Dirichlet attribution error grows by order 1/c from under-dispersion and, when p > 1, by a further factor dp−1 from tail starvation.The supplied figure summarizes the concentration and tail-exponent regimes; the theorem’s bounds match up to constants.Source: Figure 11 · p. 25, §E.3 · pp. 26–27
  • At the experimental prior fits with τ = 2, tail starvation is negligible, but cap constraints force near-prior Dirichlet ranges that remain about 4 and 30 times worse than the best capped design at 288 and 40 domains.The 24-domain setup is in the p < 1 regime, and both samplers fit well there; at 288 domains, no surrogate fits the Dirichlet swarm.Source: §E.3 · p. 27

E.4 On optimal sampling

  • On the uncapped simplex, the balanced vertex design is uniquely optimal across trace error, determinant, minimum eigenvalue, and worst-case prediction variance.Its prediction variance for every mix is bounded by vmax d/m, which no other design improves in the worst case.Source: §E.4 · p. 28
  • With equal caps and uniform prior, vertex sampling is optimal on the capped simplex; its expected information matrix reaches the stated bounds for determinant, minimum eigenvalue, and trace of the inverse.The cap makes each run cover multiple domains, increasing the reference trace-of-inverse error by about d/(κ−1) relative to the uncapped value.Source: §E.4 · pp. 28–29
  • With uniform prior, equal caps κ/d, and independently drawn runs (q=d/κ integer), the matrix has eigenvalue 1/d along 1_d and (κ−1)/[d(d−1)] on its orthogonal complement, with multiplicity d−1.Source: §E.4 · p. 28
  • The near-optimality guarantee for a sampled capped design requires m to exceed a multiple of d ln d, and the authors note this sufficient condition exceeds their swarm sizes.In practice, they select the best-conditioned design from many samples; the theorem does not cover that selection procedure.Source: §E.4 · pp. 29–30

F Limitations

The paper identifies proxy-swarm training cost and gaps between its attribution-error theory and broader practical decisions as limitations.

F Limitations

  • Training a proxy swarm costs compute for all regression-based methods, even though its cost-to-performance trade-off can improve with target scale and budget.The authors estimate that a 512 × 30M proxy swarm would use 0.1k GPU-hours, or 4% of a 4B model trained on 120B tokens.Source: F Limitations · p. 30
  • The theory bounds attribution-estimation error, not prediction error, and its noise bound depends on how the corpus is partitioned into domains.Consequently, the d^2/m rate alone does not determine how many domains to use.Source: F Limitations · p. 30
  • Whether the optimal mix is dense in positively attributed domains depends on the benchmark, corpus size, and pretraining token budget.The authors treat curvature as a utility trade-off; LinMix’s hyperparameter-free property concerns proxy-mix sampling and regression fitting.Source: F Limitations · p. 30

G Compute Resources

The paper reports proxy and target training costs across model sizes and estimates total project compute.

G Compute Resources

  • The proxy swarms cost 0.7k GPU-hours at 30M parameters, 1.4k at 60M, 1.8k for 100M OLMo2, and 2.7k for 100M Qwen3.5.The Qwen3.5-like 100M proxy takes 2.6 GPU-hours per run, 1.5 times the OLMo2 proxy’s cost.Source: G Compute Resources · p. 30
  • The reported target-model runs account for about 7.0k GPU-hours at 1B OLMo2, 1.1k at 1B Qwen3.5-like, and 17k at 4B.These totals cover the stated runs behind Tables 1, 5, 6, and 2.Source: G Compute Resources · p. 30
  • The reported experiments cost about 34k GPU-hours, while the authors estimate the whole project used about 51k GPU-hours including preliminary work and failed jobs.Models were trained on NVIDIA GH200 96GB GPUs without special optimizations beyond default code implementations.Source: G Compute Resources · p. 30

Related papers

Loading 2610.37390v1…