LinMix: Data Mixing as Linear Domain Attribution

Hubert Baniecki; Przemyslaw Biecek; Martin Pawelczyk

cs.LGOct 27, 2026

LinMix: Data Mixing as Linear Domain Attribution

Authors: Hubert Baniecki, Przemyslaw Biecek, Martin Pawelczyk

Organizations: University of Warsaw · Centre for Credible AI, Warsaw University of Technology · Faculty of Computer Science, University of Vienna

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.

Figures & tables

Figure 1: Data mixing for language model pretraining with linear domain attribution. We show that an interpretable linear regression model suffices for data mixing under vertex sampling.

Figure 2: Dirichlet sampling with default $\gamma,\tau$ hyperparameters does not converge for 288 domains. Fraction of Dirichlet-sampled domain mixes that violate the target model cap constraint. We choose new Dirichlet hyperparameters by minimizing the condition number of matrix $\boldsymbol{\Pi}$ (blue star).

Figure 3: Optimizer performance under the vertex and Dirichlet samplers. Spearman correlation is averaged across 5 random repetitions of 10-fold cross-validation under a fixed proxy model budget. Nonlinear is implemented using LightGBM, as in RegMix , and exponential fits the exponential law to the target, as in Olmix . The corresponding

R^{2}

values are reported in App. Figure 6 .

Figure 4: Domain weights relative to the prior mix across data mixing methods and target benchmarks. LinMix downweights math and code domains for language modeling (blue in HellaSwag), upweights programming languages and software domains for coding (red in MBPP), as well as math and science domains for math capabilities (red in GSM8K).

	24 domains ( $512\times$ 30M models)					288 domains ( $1024\times$ 60M proxy models)
	ARC		HellaSwag		CPile	ARC		HellaSwag		HS + PIQA		CPile
Method	Acc $\uparrow$	BPB $\downarrow$	Acc $\uparrow$	BPB $\downarrow$	BPB $\downarrow$	Acc $\uparrow$	BPB $\downarrow$	Acc $\uparrow$	BPB $\downarrow$	Acc $\uparrow$	BPB $\downarrow$	BPB $\downarrow$
Prior	.437	.864	.543	.788	.888	.436	.870	.541	.788	.634	.902	.890
RegMix	.496	.780	.563	.781	.842	.449	.850	.547	.784	.638	.896	.878
Olmix	.486	.785	.550	.782	.860	.435	.872	.550	.783	.641	.892	.850
LinMix	.510	.785	.568	.785	.847	.511	.757	.615	.758	.690	.855	.834

Table 1: Performance comparison of data mixing methods across

\mathbf{23\times}

1B par. OLMo2 target models pretrained on 32B tokens. Normalized accuracy and bits-per-byte (BPB) are measured on a held-out set per target benchmark. Color denotes relative difference to the best baseline.

	Target: HellaSwag				Target: MBPP (code)				Target: GSM8K (math)				CPile
	HellaSwag		PIQA		MBPP		HumanEval		GSM8K		MATH		CPile
Method	Acc $\uparrow$	BPB $\downarrow$	Acc $\uparrow$	BPB $\downarrow$	P@1 $\uparrow$	BPB $\downarrow$	P@1 $\uparrow$	BPB $\downarrow$	EM $\uparrow$	BPB $\downarrow$	EM $\uparrow$	BPB $\downarrow$	BPB $\downarrow$
Prior	.632	.743	.764	.952	.304	.375	.213	.285	.214	.403	.162	.463	.814
RegMix	.659	.735	.775	.933	.304	.379	.213	.290	.135	.418	.124	.487	.808
Olmix	.667	.731	.783	.932	.302	.368	.244	.277	.294	.391	.202	.446	.801
LinMix	.696	.716	.782	.920	.444	.334	.280	.258	.337	.383	.266	.427	.797

Table 2: Performance of data mixing methods across

\mathbf{10\times}

4B Qwen3.5 target models pretrained on 80B tokens. Mixes are optimized with 40 domains for each target benchmark and evaluated on its held-out set and a held-out benchmark. The additional CPile column reports

4\times

1B Qwen3.5 models optimized for held-out Common Pile validation loss. Color denotes relative diff. to the best baseline.

Appendix figures & tables11 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 5: Prior data mix distributions. Domain shares for the three setups: 24 domains (topics), 288 domains formed from topic and format pairs, and 40 domains extending topics with code and math data. The dashed line marks the removal of the distribution’s tail, where the kept larger half of all 576 pairs covers 94% of the corpus.

Model	Params	Layers	Heads	Hidden size	Batch	Tokens	Peak LR	$\beta_{2}$
OLMo2 30M	29.2M	4	8	256	64	1B	$3.6\times 10^{-3}$	0.99
OLMo2 60M	57.6M	8	8	384	128	2B	$2.9\times 10^{-3}$	0.98
OLMo2 100M	102.0M	12	8	512	128	3.2B	$2.4\times 10^{-3}$	0.97
OLMo2 1B	1.41B	18	16	2048	512	32B	$1.1\times 10^{-3}$	0.95
Qwen3.5 100M	104.5M	9+3	4	576	128	3.2B	$2.4\times 10^{-3}$	0.97
Qwen3.5 1B	1.33B	12+4	16	2048	512	32B	$1.1\times 10^{-3}$	0.95

Appendix

Table 3: Model architectures and training hyperparameters. Parameters include embeddings. Qwen3.5 layer counts list gated DeltaNet + full-attention layers. Batch gives sequences per step.

Benchmark	Proxy split (size)	Held-out set (size)	Shots	Target metric
ARC-Easy	train (2,251)	test (2,376)	5	normalized acc.
ARC-Challenge	train (1,119)	test (1,172)	5	normalized acc.
HellaSwag	train (39,905)	validation (10,042)	5	normalized acc.
PIQA	train (16,113)	validation (1,838)	5	normalized acc.
MBPP	train (374)	test (500)	3	pass@1
GSM8K	train (7,473)	test (1,319)	5	exact match

Appendix

Table 4: Benchmark splits, sample sizes, and few-shot counts. Dashes mark benchmarks used only as held-out checks.

Figure 6: Performance of optimizers under Vertex and Dirichlet samplers. Extended Figure 3 . Average Spearman correlation and

R^{2}

across 5 random repetitions of 10-fold cross-validation under a fixed proxy model budget. Nonlinear is implemented using LightGBM, as in RegMix , and exponential fits the exponential law to the target, as in Olmix .

Figure 7: Optimizer performance on the targets of the 4B experiments. Figure 6 for the 40-domain setup, where

512\times

100M Qwen3.5 proxy models per sampler are fit to HellaSwag, MBPP, GSM8K, and the held-out Common Pile loss (CPile). Spearman correlation and

R^{2}

are averaged across 5 random repetitions of 10-fold cross-validation.

Figure 8: The linear attribution fit transfers across proxy scales only under vertex sampling. A linear model is fit to predict HellaSwag BPB from the mix over the

1024\times

30M proxy models at 288 domains, and predicts the outcomes of the 60M and 100M swarms trained on the same mixes. Absolute BPB drops with proxy scale, so we offset predictions by one constant per swarm without changing rankings. Dirichlet-sampled mixes yield a weak in-sample fit whose ranking does not transfer, whereas vertex sampling yields a ranking that transfers across model scales.

Figure 9: Domain weights optimized for the 1B target benchmarks. Domain weights relative to the prior mix across data mixing methods and targets on the 200B-token corpus split into 24 and 288 domains, including the held-out Common Pile loss (CPile). For 288 domains, rows retain the 24-topic order and sort each topic’s formats by increasing share of the prior data mix. Topic names are at the group centers, and since a topic group with few formats is shorter than its label, it can overlap vertically. Beyond $\pi\;(r=0)$ , the support of the target mix is the set of all positively attributed domains; the mixes at $\sqrt{\pi}\;(r=0.5)$ and $\ln\pi\;(r=1)$ thus share this support and differ only in the domain weight magnitudes (Proposition 1 ). Based on this, Table 1 trains the LinMix $\ln\pi$ mix for HellaSwag at 24 domains, the $\sqrt{\pi}$ mix for CPile at 24 domains, and the $\pi$ mix otherwise.

Figure 10: Domain weights optimized for validation loss at 40 domains. Domain weights relative to the prior mix across data mixing methods for BPB on the Common Pile. The three LinMix columns vary the utility transformation.

Table 15

Figure 11: Two pitfalls of Dirichlet sampling. Order of the estimation error in Theorem 1 as a function of the Dirichlet concentration

\gamma

and the tempered tail exponent

p=s/\tau

of the Zipfian prior. The concentration is fixed, so that

c=1/(1+\gamma)

, and the boundaries at

\gamma=1

and

p=1

are schematic. For

\gamma\gg 1

, the mixes cluster around the prior, and the error grows with

1/c\approx\gamma

. For

p>1

, the rarest domains receive too little mass, and the error grows super-quadratically in

d

Figure 12: Prior data mixes are approximately Zipfian. Domain share against rank on log-log axes for the three setups, with least-squares fits of Zipf’s law

(\boldsymbol{\pi}_{\text{prior}})_{k}\propto k^{-s}

in log-log coordinates. The 288-domain setup retains the larger half of the

576

topic–format pairs, which truncates its tail.

Data Mixture Optimization

Data mixture optimization studies how to allocate training examples across domains and tasks to improve model performance under data and compute constraints.

See all 40 papers

Data Mixing for Large Language Models Pretraining: A Survey and Outlook
Mar 25, 2026Zhuo Chen, Yuxuan Miao, Supryadi +1Data Mixture Optimization Language Model Pretraining
The International Joint Institute of Tianjin University, Tianjin University, Tianjin 300072, China · TJUNLP Lab, School of Computer Science and Technology, Tianjin University, Tianjin 300350, China
Data Mixing as Mixture Experiment: Response Surface Methodology and Optimal Design for Large Language Model Pretraining
Aug 24, 2026Yicheng Mao, Hongru DuData Mixture Optimization Language Model Pretraining
Department of Mathematics and Statistics, University of Calgary, University Drive NW, Calgary, T2N 1N4, Canada · Department of Systems and Information Engineering, University of Virginia, Charlottesville, VA, USA
CausalMix: Data Mixture as Causal Inference for Language Model Training
Jul 1, 2026Zinan Tang, Yukun Zhang, Shaomian Zheng +7Data Mixture Optimization LLM Training
Tsinghua University · Ant Group · Renmin University of China

Mar 25, 2026J/K move · Enter open · S save

Data Mixing for Large Language Models Pretraining: A Survey and Outlook

Zhuo Chen, Yuxuan Miao, Supryadi, Deyi Xiong

The International Joint Institute of Tianjin University, Tianjin University, Tianjin 300072, China · TJUNLP Lab, School of Computer Science and Technology, Tianjin University, Tianjin 300350, China

Data Mixture Optimization Language Model Pretraining

Open paper AskPDF arXiv

Large language models (LLMs) rely on pretraining on massive and heterogeneous corpora, where training data composition has a decisive impact on training efficiency and downstream generalization under realistic compute and data budget constraints. Unlike sample-level data selection, data mixing optimizes domain-level sampling weights to allocate limited budgets more effectively. In recent years, a growing body of work has proposed principled data mixing methods for LLM pretraining; however, the literature remains fragmented and lacks a dedicated, systematic survey. This paper provides a comprehensive review of data mixing for LLM pretraining. We first formalize data mixture optimization as a bilevel problem on the probability simplex and clarify the role of data mixing in the pretraining pipeline, and briefly explain how existing methods make this formulation tractable in practice. We then introduce a fine-grained taxonomy that organizes existing methods along two dimensions: static versus dynamic mixing. Static mixing is further categorized into rule-based and learning-based methods, while dynamic mixing is grouped into adaptive and externally guided families. For each class, we summarize representative approaches and analyze their strengths and limitations from a performance-cost trade-off perspective. Building on this analysis, we highlight challenges that cut across methods, including limited transferability across data domains, optimization objectives, models, and validation sets, as well as unstandardized evaluation protocols and benchmarks, and the inherent tension between performance gains and cost control in learning-based methods. Finally, we outline several exploratory directions, including finer-grained domain partitioning and inverse data mixing, as well as pipeline-aware designs, aiming to provide conceptual and methodological insights for future research.

Abstract

Data Mixture Optimization

Explore similar work

Data Mixing for Large Language Models Pretraining: A Survey and Outlook

Data Mixing as Mixture Experiment: Response Surface Methodology and Optimal Design for Large Language Model Pretraining

CausalMix: Data Mixture as Causal Inference for Language Model Training