LinMix: Data Mixing as Linear Domain Attribution
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.
Related
- Linear and Nonlinear Phase Mixing for the Gravitational Vlasov–Poisson System Under an External Kepler Potential
- Linear time mixing and pre-cutoff for oriented kinetically constrained models
- Near linear mixing of random-cluster Glauber dynamics on random regular graphs
- Null-controllability for some linear parabolic systems with controls acting on different parts of the domain and its boundary
- Policy Iteration for Domain Randomized Linear Quadratic Systems
BibTeX
@misc{baniecki2026linmix,
title = {{LinMix: Data Mixing as Linear Domain Attribution}},
author = {Hubert Baniecki and Przemyslaw Biecek and Martin Pawelczyk},
howpublished = {arXiv},
year = {2026},
month = oct,
doi = {10.48550/arxiv.2610.37390},
eprint = {2610.37390},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2610.37390},
note = {CC BY 4.0},
}
RIS
TY - UNPB
TI - LinMix: Data Mixing as Linear Domain Attribution
AU - Hubert Baniecki
AU - Przemyslaw Biecek
AU - Martin Pawelczyk
PY - 2026
DA - 2026/10/27
PB - arXiv
DO - 10.48550/arxiv.2610.37390
UR - https://arxiv.org/abs/2610.37390
AB - 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.
ER -