LinMix: Data Mixing as Linear Domain Attribution

Models and Data Code

Introduction

The composition of a pretraining dataset, drawn from an ever-growing data corpus, is one of the most decisive factors governing the downstream capabilities of large foundation models ([21] (Team Olmo 2025)). Given the impressive gains in model training efficiency and benchmark performance obtained by curating an optimal dataset, data mixing has emerged as a well-founded optimization problem. Naturally, a source dataset can be split into the so-called domain subsets based on their web origin (e.g., GitHub, arXiv, Fan et al., [17]), semantic type (e.g., code, math, Li et al., [42]), or statistical profile (e.g., by clustering data representations, Diao et al., [15]). Data mixing is then commonly formulated as a domain-weighting problem, where each domain receives a fraction of the target model’s token budget and the model is pretrained on the resulting mixture ([75, 42, 66, 76, 62]).

In this work, we study regression-based data mixing ([46, 73, 77, 9]), which consists of two stages: (i) training a swarm of small proxy models on randomly sampled mixes, and (ii) selecting the target mix by optimizing a regression model fit to predict performance from the mix. This procedure improves both the benchmark performance and the training efficiency of language models, yet two of its design choices remain poorly understood and are addressed ad hoc in practice. First, it is unclear from which distribution the proxy mixes should be drawn, a question we frame as the classic problem of regression design Kiefer and Wolfowitz [38], where the sampled mixes govern how much a fixed budget of proxy runs reveals about each domain’s effect. Second, we show that popular methods do not scale to hundreds of fine-grained domains, where their Dirichlet sampling and optimization hyperparameters lack working defaults and require per-setup heuristic tuning, which becomes increasingly difficult as the number of domains increases.

Domain attribution. In parallel, data attribution quantifies how much a model’s behavior depends on its training data [40, 24]. We formalize regression-based data mixing as a form of data attribution, in which domains replace individual training samples. A linear fit to the proxy runs thus yields a domain-level datamodel Ilyas et al. [31], whose coefficients measure each domain’s marginal contribution to the target outcome. Consequently, a reliable regression design identifies these interpretable domain weights, an objective this work addresses from first principles.

Contributions. We introduce LINMIX as a simple and effective data mixing method whose vertex proxy sampling scheme enables selection of the optimal mix for language model pretraining via an interpretable linear model (§ 3). Figure 1 illustrates its design. Pretraining experiments with 1B-parameter OLMo2 and Qwen3.5 language models demonstrate that LINMIX outperforms data mixing methods relying on Dirichlet sampling across reasoning, code, and math benchmarks (§ 4). Furthermore, our theoretical analysis reveals fundamental limitations of Dirichlet-based mixing (§ 5). Beyond requiring heuristic hyperparameter tuning, it severely inflates attribution error whenever proxy mixes concentrate around the prior or when the prior becomes skewed (Theorem 1). Finally, we formally establish the theoretical optimality of vertex sampling in LINMIX (Theorem 2).

Background

Language models are pretrained on a corpus of $n$ sequences partitioned into $d$ domains $\mathcal{D}_1,\ldots,\mathcal{D}_d$ with $n_k := |\mathcal{D}_k|$ sequences each. The natural proportions $(\pi_{\mathrm{prior}})_k := n_k/n$ define the prior data mix, a point on the probability simplex $\Delta^{d-1}$. Whereas training on the corpus as-is draws data according to $\pi_{\mathrm{prior}}$, data mixing seeks to reweight this prior into domain weights $\pi_{\mathrm{target}} \in\Delta^{d-1}$ that maximize the downstream performance of the target model [75, 17, 1]. However, evaluating even a single candidate mix at the target scale would require a full pretraining run over $b_{\mathrm{target}}$ sequences. Instead, data mixing techniques train a swarm of $m$ small proxy models, each on $b_{\mathrm{proxy}} \ll b_{\mathrm{target}}$ sequences drawn according to a candidate mix $\pi_i$, and use their performance to predict $\pi_{\mathrm{target}}$; the candidate mixes form the rows of the design matrix $\Pi\in\mathbb{R}^{m \times d}$.

Sampling $\Pi$. State-of-the-art approaches such as RegMix Liu et al. [46] and OLMix Chen et al. [9] populate the rows of $\Pi$ via Dirichlet sampling, drawing mixes from a tempered version of the prior for each $i = 1,\ldots,m$:

The scale $\gamma_i$ is the total concentration of the Dirichlet distribution, where small $\gamma_i$ produces sparse mixes dominated by a few domains, and large $\gamma_i$ tight perturbations of the tempered prior. The temperature $\tau$ flattens the prior so that dominant domains (e.g., CommonCrawl) do not absorb most of the probability mass, with $\tau= 2$ by default [46, 73].

Up-weighting a domain far beyond its prior share can only be realized at the target scale by repeating that domain’s data, which brings diminishing returns beyond a few passes [52]. Recent methods therefore restrict mixes to the truncated simplex $\mathcal{P}_{\mathrm{cap}} := \{ \boldsymbol{\pi} \in\Delta^{d-1} \mid\pi_k \le u_k \text{ for all } k\}$ with per-domain caps $u_k := \min(1, \kappa(\pi_{\mathrm{prior}})_k)$ for an up-sampling factor $\kappa\ge1$. A mix at the cap makes at most $\kappa\cdot b_{\mathrm{target}}/n$ passes over any domain’s data; in practice, Liu et al. [46] hardcode $\kappa= 15$, Wettig et al. [73] derive $\kappa\approx n/b_{\mathrm{target}}$ so that no data repeats at the target model scale, and Chen et al. [9] set $\kappa= 4n/b_{\mathrm{target}}$, i.e., at most four passes.

Optimizing $\boldsymbol{\pi}_{\mathrm{target}}$. Each proxy model $i$ is trained on its mix $\boldsymbol{\pi}_i$ and evaluated to obtain outcome $y_i$, e.g., validation loss or bits-per-byte, yielding the proxy dataset $(\Pi, \boldsymbol{y})$ with $\boldsymbol{y} \in\mathbb{R}^m$. Regression-based data mixing methods fit a parameterized model $f_\theta: \Delta^{d-1} \to\mathbb{R}$ to predict the outcome from the mix, and then predict the target mix by minimizing the fitted model over the truncated simplex:

In practice, REGMIX solves this optimization approximately, evaluating a gradient-boosted $f_{\hat{\theta}}$ on a large pool of Dirichlet-sampled candidate mixes and averaging the Top-$K$ best-predicted ones for robustness [46]. OLMIX [9] adopts the exponential mixing law of Ye et al. [77] as $f_\theta(\boldsymbol{\pi}) = \exp\left(\sum_{j=1}^{d}\theta_j\pi_j\right) + \theta_0$, which is convex in $\boldsymbol{\pi}$ and can hence be minimized exactly over $\mathcal{P}_{\mathrm{cap}}$. Finally, $\mathcal{D}_{\mathrm{target}}$ is constructed by drawing $\hat{\pi}_{\mathrm{target},k} \cdot b_{\mathrm{target}}$ sequences from each domain $D_k$ and the target model is trained on this dataset. Note that every step of this approach inherits its quality from the design $\Pi$, which determines how well $\theta$ can be identified from only $m$ runs; yet how to sample $\Pi$ has received little attention, which we address in this paper.

Data Mixing as Linear Domain Attribution

We formalize data mixing as a domain attribution problem seeking an interpretable linear model whose coefficients measure each domain’s contribution to the target outcome. We empirically find that such a linear fit cannot be supported using Dirichlet-sampled proxy runs, and prior work compensates by using nonlinear surrogates and per-setup hyperparameter tuning, which we report in § 4. Instead of changing the surrogate model, we change the design by sampling $\Pi$ to maximize the information the proxy runs carry about the domain effects, after which nonlinear models offer no better fit than the linear one. Reliable attribution in turn enables the constrained optimization of $\boldsymbol{\pi}_{\mathrm{target}}$, where Dirichlet-based methods degrade once $d$ reaches hundreds of fine-grained domains.

Sampling $\Pi$. Unlike prior work, we sample the $m$ mixes at the vertices of $\mathcal{P}_{\mathrm{cap}}$, aiming to maximize their spread measured by $\det(\Pi^\mathsf{T}\Pi)$ without sampling hyperparameters. This choice follows the D-optimality principle of experimental design [38], which we elaborate on in Appendix § 8. For each run $i$, we draw a uniform random permutation $\sigma_i$ of domains $\{1,\ldots,d\}$ and fill domains greedily in that order, saturating each at its cap until the unit budget is exhausted:

Each resulting $\boldsymbol{\pi}_i$ has at most one strictly interior coordinate, with the rest pinned to $u_k$ or $0$. Such a sparse design is near-orthogonal in expectation, rather than centered at $\boldsymbol{\pi}_{\mathrm{prior}}$, as in Dirichlet sampling. To further improve $\det(\Pi^\mathsf{T}\Pi)$, we pair the proxy runs, letting run $2i$ reuse the permutation of run $2i-1$ in reverse, i.e. $\sigma_{2i-1}(k) := \sigma_{2i-1}(d-k+1)$, so that domains left empty by one fill are saturated first by its counterpart. The paired fills are anti-correlated by construction, with largely disjoint supports, and thus contribute near-orthogonal directions to $\Pi$, keeping $\Pi^\mathsf{T}\Pi$ well conditioned.

Optimizing $\boldsymbol{\pi}_{\mathrm{target}}$. Fitting the linear model $f_\beta(\boldsymbol{\pi}) = \beta^\mathsf{T}\boldsymbol{\pi}$ in Eq. (2) by least squares yields the attribution estimate $\hat{\beta} = (\Pi^\mathsf{T}\Pi)^{-1}\Pi^\mathsf{T}\boldsymbol{y}$, whose covariance scales with $(\Pi^\mathsf{T}\Pi)^{-1}$. We always center $\hat{\beta}$ to zero mean so that $\hat{\beta}_k$ measures the marginal value of domain $k$ relative to the average. We also negate $y$ for target metrics where lower is better, so that a positive $\hat{\beta}_k$ indicates a domain of above-average value. The mix selection then becomes a linear program over the truncated simplex,

solved by filling domains to their caps in decreasing order of $\hat{\beta}_k$ until the budget is exhausted, so the solution is itself a vertex of $\mathcal{P}_{\mathrm{cap}}$. This allocation inherently treats each additional domain sequence as equally valuable, which is appropriate for most cases, especially quality-filtered data.

Transforming the utility of $\pi_k$. Beyond sampling $\Pi$ and optimizing $\pi_{\mathrm{target}}$, a third design choice worth studying is the utility assigned to each domain weight, which can capture the diminishing returns of repeating data from a given domain. For example, if we assume that data value decays with domain repetition, the linear utility $\pi_k$ can be transformed to the logarithmic one $\ln\pi_k$, so that each additional sequence of a domain is worth less than the previous one. Theoretically, these two extremes can be framed as special cases of the general isoelastic objective

which recovers the linear objective in Eq. (4) as $r \to0$ and the log-linear form $\sum_{k:\hat{\beta}_k>0}\hat{\beta}_k \ln\pi_k$ at $r=1$. Proposition 1 establishes that positive domain attribution determines the support of the optimal mix, while the curvature $r$ only allocates the budget within it (see an example in App. Figure 9). We discuss the classical roots of this theory and its closest data selection counterpart in Appendix § 8.

As $r \to0$, the solution in Corollary 1 collapses to the greedy fill of Eq. (3) with domains ordered by decreasing $\hat{\beta}_k$ (for distinct $\hat{\beta}_k$). Optimal mixes thus saturate the caps of their top domains, such as our sampled vertices, and predicting with $\hat{\pi}_{\mathrm{target}}$ interpolates rather than extrapolates. Unlike in Dirichlet sampling, optima here do not fall outside the convex hull of mixes clustered around the prior.

Experiments

Data. We use the 200B-token WebOrganizer corpus (Wettig et al., [73]) built from CommonCrawl web data and cleaned for language model pretraining. Each document is assigned one of 24 topics (e.g., science, politics) and one of 24 formats (e.g., academic writing, news), and we split the corpus into 2048-token sequences, each drawn from a single domain. We study two mixing granularities: 24 coarse domains given by the topics alone, and 288 fine-grained domains formed by the larger half of the 24 × 24 topic–format pairs, covering 94% of the corpus (cf. App. Figure 5). Removing the practically empty tail from the prior data mix reduces noise when estimating domain weights.

Language models. Following Chen et al. [9], we build on the open source OLMo 2 implementation (Team OLMo et al., [64]) to train 3M and 60M proxy models and 1B target models. All models are trained past the 20 tokens-per-parameter Chinchilla ratio (Hoffmann et al., [29]) to avoid undertraining (Chen et al., [11]), on 1B, 2B, and 32B tokens, respectively. This setup is comparable to prior work that pairs 27M (Liu et al., [46]) and 50M (Wettig et al., [73]) proxies with 1B targets. Following Wettig et al. [73], we set the up-sampling cap to $\kappa\approx n/b_{\mathrm{target}}$ so that no sequence repeats at the target scale, a non-arbitrary choice shared by all sampling methods we compare. Hyperparameters and further training details are given in Appendix § 10.

Target benchmarks. We optimize data mixes for three targets: ARC (easy and challenge, Clark et al., [12]) for scientific question answering; HellaSwag (Zellers et al., [81]) for commonsense reasoning; and HellaSwag combined with PIQA (Bisk et al., [5]) as a two-benchmark target. In each case, we optimize the mix against proxy-model bits-per-byte (BPB, Gao et al., [22]) metric on the benchmark’s training split and report the target model’s accuracy on a held-out set: test for ARC, and validation for HellaSwag and PIQA, whose test labels are withheld. We also optimize for BPB validation loss on the Common Pile (CPile, Kandpal et al., [32]) data corpus to diversify the setup.

Dirichlet sampling of proxy data mixes requires tuning hyperparameters

Before training proxy models, we need to sample the design matrix $\Pi$. Figure 2 shows the fraction of Dirichlet-sampled domain mixes that violate the cap constraint of $\mathcal{P}_{\mathrm{cap}}$ and must be rejected. For 288 domains, the baseline algorithm does not converge with the default hyperparameters $\tau= 2$, $\gamma_{\mathrm{lo}} = 0.1$, $\gamma_{\mathrm{hi}} = 5$, or even $\gamma_{\mathrm{hi}} = 10$ (cf. the hatched area in Figure 2). Since related work gives no guidance on how to tune them, we set $\tau= 2$ and propose to select $\gamma$ by minimizing the condition number of $\Pi$, which yields mixes much closer to the prior at $\gamma_{\mathrm{lo}} = 40$ and $\gamma_{\mathrm{hi}} = 4000$.

Vertex sampling outperforms Dirichlet, enabling a linear model fit

We train 512 × 30M and 1024 × 60M proxy models per sampler for 24 and 288 domains, respectively. Figure 3 shows the cross-validated fit between predicted and observed benchmark BPB (Spearman $\rho$) under each design as the number of proxy models grows. At 24 domains, every surrogate fits well under both samplers, but at 288 domains, no surrogate helps Dirichlet sampling even when $\gamma$ is tuned ($\rho\le0.39$ at $m = 1024$). Conversely, our hyperparameter-free vertex sampling puts the interpretable linear model on par with the best nonlinear alternative (see also $R^2$ in App. Figure 6). 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. As an ablation in App. Figure 8, we train 30M and 100M models on 288 domains, fit a linear model from the 30M swarm, and use it to rank the outcomes of 60M and 100M models at $\rho\ge0.95$, while the Dirichlet fit does not transfer ($\rho\le0.42$). We visualize the optimized domain weights in App. Figure 9, where relaxing curvature $\tau$ retains the full positively attributed support (Prop. 1). Based on it, we transform the utility on 24 domains with $\ln\pi$ for HellaSwag, which adds Social Life, Fashion & Beauty, Games, and Industrial to the four domains at their caps (Health, Home & Hobbies, Food & Dining, Adult); and with $\sqrt{\pi}$ for CPile, which adds a sixth Hardware domain to the data mix.

LinMix outperforms data mixing baselines on a 1B target model benchmark

We compare LinMix to a nonlinear RegMix, which outperforms several data selection methods and was applied in WebOrganizer, as well as Olmix, which benchmarks several data mixing approaches and finds the exponential model to be state-of-the-art. Table 1 shows target benchmark performance where LinMix with a linear model significantly improves data mixing with 288 domains, and is on par or better with 24 domains. This result is consistent with the proxy metric values reported in Section 4.2, which here translate into downstream performance. An interpretable linear regression suffices for data mixing that targets the CPile validation loss under vertex sampling, i.e. it is on par with or better than complex Dirichlet sampling with nonlinear regression. We ablate on two optimization factors in Appendix D: Table 5 indicates that scaling the linear model coefficients by standard error could improve performance, and Table 6 demonstrates how transforming the utility of $\pi$ towards $\ln\pi$ relaxes the 24-domain support for HellaSwag (Proposition 1), yielding better results.

Method	24 domains (512 × 30M models) ARC Acc↑	BPB↓	HellaSwag Acc↑	BPB↓	CPile BPB↓	288 domains (1024 × 60M proxy models) ARC Acc↑	BPB↓	HellaSwag Acc↑	BPB↓	HS + PIQA Acc↑	BPB↓	CPile BPB↓
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
	+2.8%	−0.6%	+0.9%	−0.5%	−0.6%	+13.8%	+10.9%	+11.8%	+3.2%	+7.6%	+4.1%	+1.9%

Sampling	Scaling	Acc↑	BPB↓
Dirichlet	$\hat{\beta}$	.597	.764
	$\hat{\beta}/\mathrm{SE}$	.602	.761
Vertex	$\hat{\beta}$	.615	.758
	$\hat{\beta}/\mathrm{SE}$	.617	.759

Utility	24 domains		288 domains
	Acc↑	BPB↓	Acc↑	BPB↓
$\pi$ ($r = 0$)	.543	.795	.615	.758
$\sqrt{\pi}$ ($r = .5$)	.565	.782	.611	.758
$\ln\pi$ ($r = 1$)	.568	.785	.610	.759

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

We demonstrate the broader applicability of LinMix by doubling the data corpus and training a target model with ten times more compute. The 24-topic WebOrganizer dataset is extended with Nemotron-CC-Math (Karimi Mahabadi et al., [35]) and Stack-Edu (Allal et al., [2]) split into 15 programming languages, yielding a 416B-token corpus spanning 40 domains. We pretrain 4B Qwen3.5 (Qwen Team, [57]) target models on 80B tokens at a 20$\times$ ratio to allow for more diverse experiments. We train 512$\times$ 100M proxy models per sampler to prevent under-parametrization, given the more challenging coding and math capabilities. Mixes are optimized towards training splits of HellaSwag (Austin et al., [3]), and GSM8K (Cobbe et al., [13]). Figure 4 shows an example of optimized domain weights targeting three diverse model capabilities. It illustrates that RegMix obtains similar mixes across the three target capabilities, while the exponential law of Olmix differentiates between targets, but to a lesser degree than LinMix. Table 2 shows that LinMix delivers strong performance across the language, coding, and math benchmarks, transferring these capabilities to held-out benchmarks such as HumanEval (Chen et al., [8]) and MATH (Hendrycks et al., [28]; Lewkowycz et al., [41]). We additionally train a 1B-parameter Qwen3.5-like model on 32B tokens targeting Cile from the same proxy model swarm and report its results in the last column.

Method	Target: HellaSwag HellaSwag Acc↑	BPB↓	PIQA Acc↑	BPB↓	Target: MBPP (code) MBPP P@1↑	BPB↓	HumanEval P@1↑	BPB↓	Target: GSM8K (math) GSM8K EM↑	BPB↓	MATH EM↑	BPB↓	CPile BPB↓
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
	+4.4%	+2.0%	−0.1%	+1.3%	+46.1%	+9.3%	+15.0%	+7.0%	+14.4%	+2.1%	+31.6%	+4.3%	+0.5%

Theoretical Analysis

We formalize the domain attribution view of data mixing from first principles in Appendix E and summarize these results here. Starting from linear datamodels (Ilyas et al., [31]), we show that the expected outcome of a proxy run is linear in its domain mix, with coefficients that average sample-level attribution scores (Proposition 2, Corollary 2). Domain attribution thus interpolates between data attribution and a single-domain corpus, preserving the linear structure in between. We further relate it to the Shapley and Banzhaf data values of cooperative game theory in Remark 1.

With a finite budget of proxy runs, estimating the domain attribution is a linear regression whose error is governed by the design $\Pi$ (Lemma 1). Dirichlet sampling inflates this error as the concentration $\gamma$ grows and keeps its mixes close to the prior (Theorem 1). With many domains, Dirichlet mixes cannot stay far from the prior because sparse mixes at a small $\gamma$ violate the per-domain caps and are rejected. We therefore tune $\gamma$ empirically to a near-prior range (§ 4.1), for which the expected error of Dirichlet sampling is theoretically several times the lowest error possible with equal caps (Appendix § 12.§ 3). This matches Figure 3, where both samplers lead to reliable fits at 24 domains, but no surrogate fits the Dirichlet design well at 288 domains. In theory, vertex sampling without caps trains each proxy model on a single domain and thus estimates every domain attribution independently and equally well, which no other design can improve on (Proposition 3). In practice, the caps force every mix to spread over several domains and reveal less about each of them, but vertex sampling remains the best possible design in expectation for equal caps (Theorem 2).

Under these idealizations, Theorems 1 and 2 help explain why the linear fit of LINMIX can be more reliable than one based on Dirichlet sampling. Dirichlet mixes concentrate around the prior, while vertices carry as much evenly spread information as the caps allow.

Conclusion

We demonstrate that an interpretable linear model suffices for data mixing under vertex sampling and validate it across language model pretraining benchmarks. We further develop the theoretical foundations of LINMIX and relate data mixing to the attribution literature. While this paper contributes to the interpretability of data mixing, our results suggest that optimizing a vertex-sampled matrix with a nonlinear model could yield even better target model performance, which we leave for future work. We discuss limitations of our work in Appendix § 13, most importantly the cost of training a proxy swarm and the single training run per target model.

Acknowledgments and Disclosure of Funding

Hubert Baniecki was supported by the Foundation for Polish Science (FNP), and the state budget within the Polish Ministry of Education and Science program “Pearls of Science” project number PN/01/0087/2022. Work on this project is financially supported by the Foundation for Polish Science (FNP) grant ‘Centre for Credible AI’ No. FENG.02.01-IP.05-0058/24, and we gratefully acknowledge the Polish high-performance computing infrastructure PLGrid (HPC Center: ACK Cyfronet AGH) for providing computer facilities and support within the computational grant no. PLG/2026/019857.

Related Work

Design of experiments. Sampling $\Pi$ is a classic problem in statistics. Because each row is a vector of domain proportions that sum to one, the proxy swarm is a mixture experiment [59]. Choosing inputs that minimize the variance of a regression estimate is the subject of optimal experimental design. A design is $D$-optimal when it maximizes $\det(\Pi^\top\Pi)$, which minimizes the volume of the confidence ellipsoid of $\hat{\beta}$ [38]. For approximate designs, Kiefer and Wolfowitz [39] show that maximizing this determinant also minimizes the worst-case prediction variance. The D-optimal design for the first-degree mixture model places equal weight on the vertices of the simplex [37]. When upper bounds such as our caps $u_k$ truncate the simplex, the extreme vertices designs of McLean and Anderson [49] likewise use the vertices of the constrained region. LINMIX follows this principle at a scale where computing exact optimal designs on $\mathcal{P}_{\mathrm{cap}}$ is impractical, sampling vertices generated by uniformly random domain permutations. We prove that this sampling scheme is optimal in expectation on $\mathcal{P}_{\mathrm{cap}}$ with equal caps Theorem 2 and leave the case of unequal caps to future work. In § 4.1, we tune Dirichlet sampling by minimizing the condition number of $\Pi$, a criterion closely related to D-optimality and $\det(\Pi^\top\Pi)$. Both degrade together because nearly-collinear mixes reduce the smallest singular value of $\Pi$, thereby collapsing the determinant and inflating the condition number. LINMIX also samples proxy mixes in pairs by filling domains in a random order and then in reverse order, following the idea of antithetic variates [27]. This relates to pairing sampled coalitions with their complements that has improved Shapley value approximation [20], though our capped mixes need not be exact complements. More broadly, regression design is of interest for improving attribution, e.g. Musco and Witter [53] improve regression-based Shapley value approximation via leverage-score sampling of coalitions.

Data mixing. Data mixing methods differ in how they estimate domain value and spend compute to choose a mix [47]. Online methods adapt the mixture during training. DoReMi [75] learns weights during proxy training through group distributionally robust optimization, DoGE [17] and Tan et al. [62] use gradient-based updates guided by generalization or validation performance, and PiKE [44] re-weights data sources from their gradients. MixMin [66] optimizes a weighted ensemble of proxy predictions, while Chameleon [76] derives weights from relationships between domain embeddings. Data Mixing Laws [77], AutoScale [33], and Shukor et al. (2025) use scaling relationships to predict mixture performance or select mixtures at larger training scales. Mixture optimization also extends to tokenizers [74] and multimodal midtraining [72]. Adaptive search allocates new proxy runs to promising mixtures through iterative predictor fitting [15] or Bayesian optimization [10, 78]. LINMIX instead fixes its proxy mixtures before observing their outcomes, using vertex sampling to estimate domain effects across the feasible region. This design enables a useful linear fit across hundreds of domains; combining it with adaptive search is a direction for future work.

Theory of optimal regression-based data mixing. Regression-based data mixing methods select the target mix by minimizing a fitted surrogate model, offering little theoretical insight into the optimal mix itself [46, 73, 9]. Recently, Chen et al. [9] bound the performance loss when a previously optimized mix is reused after the domain set changes, taking the mixing law and its exact minimization as given. Thrush et al. [65] rank domains by correlating the per-domain losses of public models with their benchmark scores instead, and prove that such rankings are optimally used by saturating the caps in rank order. In our framework, this rule is the vertex solution of Eq. (5) at $r = 0$, with the ranking estimated from a designed experiment on training mixtures rather than a correlation across uncontrolled models. For $r > 0$, the objective is a weighted isoelastic utility, also known as the alpha-fair family of resource allocation [50], whose marginal utility grows unbounded at zero weight, so the maximizer places positive weight on every positively attributed domain Proposition 1. Corollary 1 solves the KKT conditions in closed form, yielding a capped water-filling allocation [cf. 7].

Data attribution. Our work relates to the parallel literature on data attribution, which typically studies the influence of individual training examples rather than entire domains (Ghorbani and Zou, 2019; Feldman and Zhang, [19]; Wang and Jia, [68]; Li and Yu, [43]). We identify a connection between regression-based data mixing and datamodels (Ilyas et al., [31]; Saunshi et al., [58]), which predict model outputs from training-set membership and show that a linear model can capture these dependencies. DsDm (Engstrom et al., [16]) turns this predictive view into data selection, using gradient-based TRAK estimates (Park et al., [54]) to rank individual examples and select a fixed-size subset. LINMIX shares this principle but estimates domain proportions from proxy runs and allocates tokens under domain caps. Using these estimates for a larger target model requires that data effects transfer across scales, as studied by Khaddaj et al. [36]. Transfer depends on proxy training hyperparameters, as lower learning rates improve agreement with the dataset rankings of larger models (Wang et al., [70]). With a limited number of runs (Wang et al., [69]), the question becomes which mixtures reveal enough about each domain to support the allocation, which vertex sampling answers by design.

Data selection and filtering. Data quality and filtering methods improve training corpora by selecting or weighting individual examples. DataComp ([21, 42]) benchmarks how data curation affects large model performance. Methods that combine selection criteria include Meta-rater (Zhuang et al., [82]), which learns combinations of quality scores from proxy runs, and Fan et al. [18], who jointly optimize quality and diversity through learned selection masks, while Group-MATES (Yu et al., [80]) further accounts for interactions between examples. During training, selection can adapt by coordinating multiple criteria (Bai et al., [4]), updating example weights (Yu et al., [79]), or scoring examples by their optimizer updates (Wang et al., [71]). For fine-tuning, Ling et al. [45] select data using a diversity reward, while LimaCost (Moon et al., [51]) measures instruction-data value against a reference dataset using gradients. Filtering and mixing can be combined by curating the corpus before optimizing its domain mix, as in our experiments.

Proofs

Proof. Throughout, let $\tilde{\pi}$ denote the feasible mix that spreads the budget over $\mathcal{A}$ in proportion to the caps, $\tilde{\pi}_k := u_k/u_{\mathcal{A}}$ for $k \in\mathcal{A}$ and $\tilde{\pi}_k := 0$ otherwise. Its entries sum to one and $u_{\mathcal{A}} > 1$ gives $\tilde{\pi}_k < u_k$, so it is feasible and puts positive weight on every domain in $\mathcal{A}$.

Existence. For $r \in(0,1)$, the objective $v_r(\pi) := \sum_{k=1}^{d} \hat{\beta}_k \frac{\pi_k^{1-r}-1}{1-r}$ is continuous on the compact set $\mathcal{P}_{\mathrm{cap}}$, so a maximizer exists. For $r = 1$, the specialized objective defines the objective as $v_1(\pi) := \sum_{k \in\mathcal{A}} \hat{\beta}_k \ln\pi_k$ with the convention $\ln0 = -\infty$. This convention continues the case $r < 1$, where all maximizers lie on the face $\{\pi_k = 0 \text{ for } k \notin\mathcal{A}\}$. On this face the omitted terms are constants that do not affect the maximizer, and the remaining sum converges pointwise to $v_1$ as $r \to1$. $v_1$ is upper semicontinuous and bounded above by zero, so it attains a finite maximum, as $v_1(\tilde{\pi}) > -\infty$.

Positive weight on every domain in $\mathcal{A}$. The utility of a domain in $\mathcal{A}$ has an unbounded derivative at zero weight, so leaving such a domain empty cannot be optimal. To make this precise, let $\pi^\star$ be any maximizer and suppose $\pi^\star_j = 0$ for some $j \in\mathcal{A}$. For $r = 1$ this gives $v_1(\pi^\star) = -\infty< v_1(\tilde{\pi})$, a contradiction. For $r \in(0,1)$, pick any $i$ with $\pi^\star_i > 0$ and move a small weight $\varepsilon\in(0,\min(\pi^\star_i/2,u_j))$ from $i$ to $j$. Only two terms change. The $j$-th gains $\hat{\beta}_j \frac{\varepsilon^{1-r}}{1-r} > 0$. The $i$-th loses at most $\hat{\beta}_i^+(\pi^\star_i/2)^{-r}\varepsilon$ with $\hat{\beta}_i^+ := \max(\hat{\beta}_i,0)$, by the mean value theorem on $[\pi^\star_i/2,\pi^\star_i]$. The gain is roughly $\varepsilon^{1-r}$ while the loss is only proportional to $\varepsilon$, so the move strictly improves the objective for small $\varepsilon$, again a contradiction. Hence $\pi^\star_k > 0$ for all $k \in\mathcal{A}$.

Zero weight outside $\mathcal{A}$. Conversely, weight outside $\mathcal{A}$ adds no value and is better spent inside. Suppose $\pi^\star_j > 0$ for some $j \notin\mathcal{A}$. Because $\sum_{k \in\mathcal{A}} \pi^\star_k \le1 < u_{\mathcal{A}}$, some domain $i \in\mathcal{A}$ sits strictly below its cap. Moving the weight $\varepsilon:= \min(\pi^\star_j, u_i - \pi^\star_i) > 0$ from $j$ to $i$ is feasible. The $j$-th term does not decrease, since $\hat{\beta}_j \le0$ makes it non-increasing in $\pi_j$, and for $r = 1$ it is absent from $\nu_1$ altogether. The $i$-th term strictly increases, a contradiction. Hence $\pi_k^\star= 0$ for all $k \notin\mathcal{A}$.

Uniqueness. By the two steps above, every maximizer lies in the convex set $\mathcal{F} := \{\boldsymbol{\pi} \in\mathcal{P}_{\mathrm{cap}} : \pi_k = 0 \text{ for } k \notin\mathcal{A}\}$ and is strictly positive on $\mathcal{A}$. On $\mathcal{F}$, the objective is a sum of strictly concave functions of the coordinates in $\mathcal{A}$, since $\hat{\beta}_k > 0$ and $r \in(0, 1]$. Two distinct maximizers would therefore be beaten by their midpoint, so the maximizer is unique.

Linear limit. At $r = 0$, the objective $\nu_0(\boldsymbol{\pi}) = \hat{\boldsymbol{\beta}}^{\mathsf{T}}\boldsymbol{\pi} - \hat{\boldsymbol{\beta}}^{\mathsf{T}}\mathbf{1}$ is linear, so it is maximized at a vertex of the truncated simplex. The greedy fill solving Eq. (4) can then exhaust the budget before reaching the lowest-value domains in $\mathcal{A}$ and leave them at zero weight. For example, $\hat{\boldsymbol{\beta}} = (2,1,-3)$ with $\boldsymbol{u} = (1,1,1)$ yields the optimum $(1,0,0)$.

Proof. Fix $r \in(0,1]$. The claimed strict concavity on the support was established in the uniqueness step of Proposition 1, as each summand over $\mathcal{A}$ is strictly concave. By Proposition 1, the maximizer $\boldsymbol{\pi}^\star$ is unique, lies in $\mathcal{F}$, and satisfies $\pi_k^\star> 0$ for all $k \in\mathcal{A}$. On $\mathcal{F}$, the problem maximizes this concave objective under the affine constraints $\sum_{k \in\mathcal{A}} \pi_k = 1$ and $\pi_k \le u_k$. Because every $\pi_k^\star$ is positive, the objective is differentiable at $\boldsymbol{\pi}^\star$ and the constraints $\pi_k \ge0$ are inactive, so $\boldsymbol{\pi}^\star$ satisfies the KKT conditions (cf. Boyd and Vandenberghe [7]) with a scalar $\omega$ for the budget and multipliers $\mu_k \ge0$ for the caps. The marginal utility of domain $k$ is $\hat{\beta}_k(\pi_k^\star)^{-r}$, so stationarity reads $\hat{\beta}_k(\pi_k^\star)^{-r} = \omega+ \mu_k$ for every $k \in\mathcal{A}$. In words, the optimum equalizes marginal utilities at the common level $\omega$, except where a cap binds. Because $u_{\mathcal{A}} > 1 = \sum_{k \in\mathcal{A}} \pi_k^\star$, some domain lies strictly below its cap. For this domain, complementary slackness gives $\mu_k = 0$, so $\omega= \hat{\beta}_k(\pi_k^\star)^{-r} > 0$. Now the closed form can be read off. If $\pi_k^\star< u_k$, then $\mu_k = 0$ and stationarity gives $\pi_k^\star= (\hat{\beta}_k/\omega)^{1/r}$. If $\pi_k^\star= u_k$, then $\mu_k \ge0$ gives $(\hat{\beta}_k/\omega)^{1/r} \ge u_k$. In both cases, $\pi_k^\star= \min\left(u_k,(\hat{\beta}_k/\omega)^{1/r}\right)$. Since $\boldsymbol{\pi}^\star$ is feasible, summing this form over $k \in\mathcal{A}$ shows that $\omega$ enforces the normalization $\sum_{k \in\mathcal{A}} \min\left(u_k,(\hat{\beta}_k/\omega)^{1/r}\right) = 1$. Finally, $\omega$ is unique. The sum $\sum_{k \in\mathcal{A}} \min\left(u_k,(\hat{\beta}_k/\omega)^{1/r}\right)$ is continuous and non-increasing in $\omega$. At any $\omega$ where it equals $1 < u_{\mathcal{A}}$, some term lies strictly below its cap and there it decreases strictly in $\omega$. Therefore the value $1$ is attained at exactly one $\omega$.

Experimental Setup

Pretraining data. We tokenize data into single-domain sequences of 2048 tokens with the GPT-NeoX-20B tokenizer (Black et al., [6]), which has a vocabulary of 50,277 tokens. This way, the WebOrganizer corpus (Wettig et al., [73]) yields 195B tokens, and the 40-domain corpus adds 168B tokens of Stack-Edu (Allal et al., [2]) across 15 programming languages, plus the 53B-token highest-quality subset of Nemotron-CC-Math (Karimi Mahabadi et al., [35]), bringing the total data size to 416B. Figure 5 shows the resulting prior data mixes.

Sampling data mixes. The caps $\kappa$ become 5.98 and 5.63 for the 32B target token budget at 24 and 288 domains, and $\kappa= 3.40$ at 40 domains. We initially derived the latter for a 120B budget, allowing us to further train the models beyond the 80B tokens used. Dirichlet sampling uses $\tau= 2$ and $[\gamma_{\mathrm{lo}}, \gamma_{\mathrm{hi}}] = [0.1,10]$, $[40,4000]$, and $[100,10^4]$ at 24, 288, and 40 domains, respectively. The 40-domain range follows the criterion of § 4.1: the default parameters reject over 99.9% of draws, and the condition number of $\boldsymbol{\Pi}$ is minimized near $\gamma= 10^2$. We keep the 100% span of the tuned range, over which the rejection rate falls from 0.96 to 0.01. For both samplers, we draw $10^4$ designs and keep the best-conditioned $\boldsymbol{\Pi}$, which gives the most stable regression coefficients.

Language models. Table 3 lists model architectures and hyperparameters. Both model families follow their reference implementations (Team OLMo et al., [64]; Qwen Team, [57]); Qwen3.5 interleaves three gated DeltaNet layers per attention layer. All models use AdamW with $\beta_1 = 0.9$ and weight decay 0.1. The peak learning rate scales with the inverse cube root of the model parameter count (Porian et al., [55]). The schedule warms up linearly over the first 10% of steps, decays by cosine to 10% of the peak, and then linearly to zero over the final 10%.

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

Evaluation. We use the LM Evaluation Harness (Gao et al., [23]) for a standardized evaluation. Proxy models are scored by BPB on the training split of each benchmark, and target models by accuracy and BPB on their held-out set (Table 4). Shots are drawn from the training split, except MBPP and MATH, which use their standard fixed prompts, and HumanEval BPB, which relies on three hand-written examples. For the two-benchmark targets, ARC (easy and challenge) and HS+PIQA, we fit a surrogate per benchmark, minimize the mean of their predictions, and report the mean of the two target model scores. For the 4B models, math and code benchmarks are scored using greedy decoding, with exact match for math and pass@1 for code. We draw two disjoint held-out sets from the Common Pile dataset (Kandpal et al., [32]): a 150k-sequence split on which we score the proxy models to fit the surrogates, and a 300k-sequence split reserved for scoring the target models.

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
HumanEval	—	test (164)	0	pass@1
MATH	—	test (5,000)	4	exact match

Additional Results

Figure 6 extends Figure 3 with $R^2$ measurements, and Figure 7 repeats it for the four targets of the 4B experiments at 40 domains, where vertex sampling reaches $\rho\ge0.94$ with every surrogate at $m = 512$, while the tuned Dirichlet swarm stays at $\rho\le0.68$ on the three benchmarks and 0.80 on CPile. Figure 8 shows that the linear fit of the vertex swarm transfers across proxy scales. Figure 9 and Figure 10 visualize the interpretable domain weights of the optimized mixes at 24 and 288 domains and for the Common Pile loss at 40 domains.

Tables 5 and 6 ablate the two optimization choices of LINMIX on the HellaSwag target. Table 5 tests our hypothesis that dividing the coefficients $\hat{\beta}$ by their standard errors, which discounts uncertain domains, improves the mix. The gain is at most half a point of accuracy for both the vertex swarm and the tuned Dirichlet swarm with $\gamma\in[40, 4000]$. Table 6 shows that the interpretable transformation of $\pi$ matters at 24 domains, where $\pi$ ($r = 0$) concentrates the mix on four domains at their caps $u_k$. This support is too narrow for the target budget of 32B tokens, which a larger budget would spread through lower caps, and $\sqrt{\pi}$ or $\ln\pi$ ($r > 0$) spreads at the current one.

Theoretical Analysis

We formalize the domain attribution view of data mixing from first principles. Starting from linear datamodels, we show that the expected outcome of a proxy run is linear in its domain mix, with coefficients that average sample-level attribution scores (Proposition 2, Corollary 2), and that estimating them is a linear regression whose error is governed by the design $\Pi$ (Lemma 1). Dirichlet sampling inflates this error as the concentration $\gamma$ grows and keeps its mixes close to the prior, and also as the number of domains grows when the prior is more skewed than the temperature $\tau$ can flatten (Theorem 1). Vertex sampling is instead optimal on the simplex (Proposition 3) and, in expectation, on the simplex with equal caps (Theorem 2).

From datamodels to domain attribution

Datamodeling (Ilyas et al. [31]) attributes model outcomes to training data counterfactually, in a regime where many training runs on random subsets of the data are feasible. Fix a learning algorithm and a target input $z$, e.g. a benchmark question, and let $y(S; z)$ denote the loss on $z$ of a model trained on a subset $S$ of the corpus $\mathcal{D} := \mathcal{D}_1 \cup\cdots\cup\mathcal{D}_d$ of $n$ sequences. A datamodel is fit to the pairs $(\mathbf{1}_{S_i}, y(S_i; z))$ collected over runs $i = 1, \ldots, m$, where $\mathbf{1}_{S_i} \in\{0,1\}^n$ indicates the sequences included in run $i$. Empirically, a linear fit predicts the outcomes of held-out runs well and motivates the model

where $\beta(z) \in\mathbb{R}^n$ collects the attribution scores, i.e., the marginal effect of including each sequence on the loss at $z$, and $\eta(z)$ is a noise term with zero mean given $S$ that captures the randomness of training. Follow-up work relies on this local linearity to relate attribution scores to model outcomes (Karchmer et al. [34]), and we adopt it as the assumption underlying this analysis. As all subsets considered below have the same size, an intercept can be absorbed into $\beta(z)$ and is omitted.

Consider a proxy run on the mix $\boldsymbol{\pi}$ that trains on $b := b_{\mathrm{proxy}}$ sequences, of which $b\pi_k$ are drawn uniformly without replacement from domain $\mathcal{D}_k$, independently across domains, where we assume that $b\pi_k \le n_k$ is an integer. Let its outcome $y$ be the benchmark loss, i.e., the average of $y(S; z)$ over the set of benchmark inputs $\mathcal{Z}$. Let $\beta^{\mathcal{Z}} := \frac{1}{|\mathcal{Z}|}\sum_{z\in\mathcal{Z}}\beta(z)$ denote the benchmark-averaged attribution scores, and let $\mu_k$ and $\sigma_k^2 := \frac{1}{n_k}\sum_{s\in\mathcal{D}_k}(\beta_s^{\mathcal{Z}}-\mu_k)^2$ denote their mean and variance over the sequences of domain $k$. Proposition 2 formalizes the data attribution view of data mixing.

Proof. Linearity. For a fixed subset $S$, averaging Eq. (6) over $z \in\mathcal{Z}$ gives $y = \langle\beta^{\mathcal{Z}}, \mathbf{1}_S\rangle+ \bar{\eta}$ with $\bar{\eta} := \frac{1}{|\mathcal{Z}|}\sum_{z\in\mathcal{Z}}\eta(z)$ and $\mathbb{E}[\bar{\eta}\mid S] = 0$, because the same subset produces the loss at every $z$ and the average commutes with the inner product. Given $\boldsymbol{\pi}$, the subset $S$ contains $b\pi_k$ sequences drawn uniformly without replacement from $\mathcal{D}_k$, so every $s \in\mathcal{D}_k$ is included with probability $\Pr[s \in S] = b\pi_k/n_k$. By linearity of expectation,

Evaluating this identity at $\boldsymbol{\pi} = e_k$ identifies $\bar{\beta}_k$ as the expected outcome of a run on domain $k$ alone, which is feasible whenever $n_k \ge b$.

Variance. By the law of total variance and $\mathbb{E}[y\mid S] = \langle\beta^{\mathcal{Z}},\mathbf{1}_S\rangle$, $\operatorname{Var}(y\mid\boldsymbol{\pi}) = \operatorname{Var}(\langle\beta^{\mathcal{Z}},\mathbf{1}_S\rangle\mid\boldsymbol{\pi}) + \mathbb{E}[\operatorname{Var}(\bar{\eta}\mid S)\mid\boldsymbol{\pi}]$, where the second term is at most $\sigma_\eta^2$. The domains are sampled independently, and within domain $k$ the sum of $\beta_s^{\mathcal{Z}}$ over a simple random sample of $b\pi_k$ out of $n_k$ sequences has variance $b\pi_k\sigma_k^2\frac{n_k-b\pi_k}{n_k-1}\le b\pi_k\sigma_k^2$ by the finite population variance formula ([14]), where the fraction is read as zero for $n_k = 1$. Summing over $k$ bounds the first term by $b\sum_k\pi_k\sigma_k^2\le b\max_k\sigma_k^2$. $\square$

The linear model $f_{\beta}(\pi) = \beta^{\mathsf{T}}\pi$ fit in Eq. (2) is a datamodel whose units are domains instead of sequences, with coefficients that are budget-scaled averages of sample-level attribution scores. The partition into domains is arbitrary, since the only assumption is uniform sampling within each domain, and it affects only the noise, which stems from the randomness of training and from which sequences realize the mix, which we call the noise from sequence selection. Domains that are homogeneous in attribution reduce the latter, but their coefficients take more proxy runs to estimate, as the estimation error grows with $d$ (Lemma 1); the finest partition recovers data attribution itself.

Proof. With $n_k = 1$ for every $k$, Proposition 2 applies with $\mu_k = \beta_k^{Z}$ and $\sigma_k^{2} = 0$, and $b\pi_k \in\{0,1\}$ forces $\pi= \mathbf{1}_{S}/b$. With caps $u_k = 1/b$, the greedy fill of Eq. (3) assigns $1/b$ to the first $b$ sequences of a uniformly random permutation, which form a uniformly random subset of size $b$.

Domain attribution thus interpolates between data attribution at $d = n$ and a single undifferentiated corpus at $d = 1$, preserving the linear structure in between. Remark 1 relates domain attribution to the data values of cooperative game theory, such as the Shapley and Banzhaf values.

Remark 1 (Relation of domain attribution to Shapley and Banzhaf data values). Data valuation scores a sequence $s$ by a weighted average of its marginal contributions $y(S \cup\{s\}; z) - y(S; z)$ over subsets $S$, with weights that define the Shapley value (Ghorbani and Zou, [24]), the Banzhaf value (Wang and Jia, [68]), or the weighted Banzhaf value (Li and Yu, [43]), which includes every other sequence in $S$ independently with probability $\alpha$. If the linear Datamodel of Eq. (6) holds with an intercept for subsets of every size, these values coincide with the attribution score $\beta_s(z)$, and by Proposition 2 the domain attribution $\bar{\beta}_k$ is $b$ times the average value of the sequences in domain $k$, up to an additive constant shared by all domains. Without the linear assumption, the slope coefficients of the affine least squares approximation of $y(S; z)$ over all subsets $S$ are the Banzhaf values (Hammer and Holzman, [26]), and under independent inclusion with probability $\alpha$ they are the weighted Banzhaf values with weight $\alpha$ (Marichal and Mathonet, [48]). Vertex sampling with equal caps resembles independent inclusion with probability about $1/\kappa$, except that every run contains the same number of domains. For nonlinear outcomes, its least squares coefficients are therefore a fixed-size analogue of the weighted Banzhaf value. While a Banzhaf value captures the effect of adding a sequence to a random training set, differences between domain coefficients capture the effect of moving budget between domains at a fixed training size, which is the question that data mixing asks.

The design matrix governs the estimation error

By Proposition 2, the outcomes of $m$ proxy runs satisfy $\mathbf{y} = \mathbf{\Pi}\bar{\beta} + \boldsymbol{\xi}$ for a noise vector $\boldsymbol{\xi} \in\mathbb{R}^{m}$, and the attribution estimate $\hat{\beta} = (\mathbf{\Pi}^{\mathsf{T}}\mathbf{\Pi})^{-1}\mathbf{\Pi}^{\mathsf{T}}\mathbf{y}$ is the least squares estimator of $\bar{\beta}$. We make the standard assumptions of linear regression, writing $\mathbf{H}_{\Pi} := \frac{1}{m}\mathbf{\Pi}^{\mathsf{T}}\mathbf{\Pi}$ for the normalized information matrix, i.e., the Gram matrix of the mixes, whose inverse scales the covariance of $\hat{\beta}$.

Proposition 2 provides the first two conditions for independent proxy runs, where $v_{\max} = \sigma_{\eta}^{2} + b\max_k \sigma_k^{2}$ is admissible and, for homoscedastic noise $\Sigma_{\Pi} = \sigma^{2}I_m$, $v_{\min} = v_{\max} = \sigma^{2}$, while the third condition is a property of the design that we discuss below. The simplex is an affine hyperplane that avoids the origin, so $d$ mixes can span $\mathbb{R}^{d}$, with one direction for the overall outcome level and $d - 1$ for contrasts between domains.

Proof. This is the standard variance calculation for least squares. Substituting $\boldsymbol{y}=\boldsymbol{\Pi}\bar{\boldsymbol{\beta}}+\boldsymbol{\xi}$ gives $\hat{\boldsymbol{\beta}}-\bar{\boldsymbol{\beta}}=(\boldsymbol{\Pi}^{\mathsf{T}}\boldsymbol{\Pi})^{-1}\boldsymbol{\Pi}^{\mathsf{T}}\boldsymbol{\xi}$, which has zero conditional mean and expected squared norm $\operatorname{tr}((\boldsymbol{\Pi}^{\mathsf{T}}\boldsymbol{\Pi})^{-1}\boldsymbol{\Pi}^{\mathsf{T}}\boldsymbol{\Sigma}_{\boldsymbol{\Pi}}\boldsymbol{\Pi}(\boldsymbol{\Pi}^{\mathsf{T}}\boldsymbol{\Pi})^{-1})$. The semidefinite bounds on $\boldsymbol{\Sigma}_{\boldsymbol{\Pi}}$ sandwich this trace between $v_{\min}$ and $v_{\max}$ times $\operatorname{tr}((\boldsymbol{\Pi}^{\mathsf{T}}\boldsymbol{\Pi})^{-1})$, with equality for $\boldsymbol{\Sigma}_{\boldsymbol{\Pi}}=\sigma^2\boldsymbol{I}_m$. The last inequality holds as the $d$ eigenvalues of $(\boldsymbol{\Pi}^{\mathsf{T}}\boldsymbol{\Pi})^{-1}=\frac{1}{m}\boldsymbol{H}_{\boldsymbol{\Pi}}^{-1}$ are at most $1/(m\lambda_{\min}(\boldsymbol{H}_{\boldsymbol{\Pi}}))$. $\square$

The last inequality bounds the error by the smallest eigenvalue of $\boldsymbol{H}_{\boldsymbol{\Pi}}$, which is small when some combination of domains barely varies across the proxy runs, i.e., when the runs leave a direction of the simplex unexplored. In optimal experimental design ([38, 56]), minimizing the trace in Eq. (8), maximizing this eigenvalue, and maximizing $\det(\boldsymbol{H}_{\boldsymbol{\Pi}})$ are known as the A-, E-, and D-criteria. This eigenvalue is what the condition number $\operatorname{cond}(\boldsymbol{\Pi})=\sqrt{\lambda_{\max}(\boldsymbol{H}_{\boldsymbol{\Pi}})/\lambda_{\min}(\boldsymbol{H}_{\boldsymbol{\Pi}})}$ controls: every mix sums to one, so $\boldsymbol{\Pi}\mathbf{1}_d=\mathbf{1}_m$ and $\mathbf{1}_d^{\mathsf{T}}\boldsymbol{H}_{\boldsymbol{\Pi}}\mathbf{1}_d=1$, hence $\lambda_{\max}(\boldsymbol{H}_{\boldsymbol{\Pi}})\ge1/d$ and $1/\lambda_{\min}(\boldsymbol{H}_{\boldsymbol{\Pi}})\le d\operatorname{cond}(\boldsymbol{\Pi})^2$. Minimizing the condition number of $\boldsymbol{\Pi}$, as in § 4.1, therefore controls the last bound of Eq. (8) up to the factor $d$. The noise level and the number of domains are fixed by the data and its partition, so the number of runs $m$ and the design $\boldsymbol{\Pi}$ are the only two levers of a regression-based data mixing method. We next compare the designs of Dirichlet and vertex sampling on these quantities.

Pitfalls of Dirichlet sampling

Dirichlet sampling in Eq. (1) draws each mix as $\boldsymbol{\pi}_i\sim\operatorname{Dirichlet}(\gamma_i\bar{\boldsymbol{a}})$ around the tempered prior $\bar{\boldsymbol{a}}:=\boldsymbol{\pi}^{\mathrm{prior}}/\lVert\boldsymbol{\pi}^{\mathrm{prior}}\rVert_1$, with a log-uniform concentration $\gamma_i\in[\gamma_{\mathrm{lo}},\gamma_{\mathrm{hi}}]$. We analyze it without the caps and the rounding of $\hat{\pi}_k$ to integers, which keeps the moments in closed form, and start with the expected information matrix $\bar{\boldsymbol{H}}=\mathbb{E}[\boldsymbol{H}_{\boldsymbol{\Pi}}]$.

Proof. Second moment. For $\boldsymbol{\pi}\sim\operatorname{Dirichlet}(\boldsymbol{\alpha})$ with $\alpha_0:=\sum_k\alpha_k$, the moments $\mathbb{E}[\pi_j\pi_k]=\frac{\alpha_j\alpha_k}{\alpha_0(\alpha_0+1)}$ for $j\ne k$ and $\mathbb{E}[\pi_k^2]=\frac{\alpha_k(\alpha_k+1)}{\alpha_0(\alpha_0+1)}$ read in matrix form as $\mathbb{E}[\boldsymbol{\pi}\boldsymbol{\pi}^{\mathsf{T}}]=\frac{\boldsymbol{\alpha}\boldsymbol{\alpha}^{\mathsf{T}}+\operatorname{diag}(\boldsymbol{\alpha})}{\alpha_0(\alpha_0+1)}$. With $\boldsymbol{\alpha}=\gamma\bar{\boldsymbol{a}}$ and $\alpha_0=\gamma$, since $\bar{\boldsymbol{a}}\in\Delta^{d-1}$,

The two weights sum to one, so averaging over $\gamma$ gives Eq. (9) with $c=\mathbb{E}_{\gamma}[1/(1+\gamma)]\in(0,1)$. For $\log\gamma$ uniform on $[\log\gamma_{\mathrm{lo}},\log\gamma_{\mathrm{hi}}]$, substituting $\gamma=e^t$ and using $\frac{\mathrm{d}}{\mathrm{d}t}[t-\ln(1+e^t)]=\frac{1}{1+e^t}$ yields the closed form

Inverse. Let $\boldsymbol{D}:=\operatorname{diag}(\bar{\boldsymbol{a}})$, which is positive definite as $\bar{a}_k>0$, so $\bar{\boldsymbol{H}}\succeq c\boldsymbol{D}\succ0$. The Sherman–Morrison formula (Sherman and Morrison, [60]) applied to the rank-one update $\bar{\boldsymbol{H}}=$

$cD+(1-c)\bar{a}\bar{a}^{\mathsf{T}}$, together with $D^{-1}\bar{a}=\mathbf{1}_d$ and $\bar{a}^{\mathsf{T}}\mathbf{1}_d=1$, gives

and taking the trace yields $\operatorname{tr}(\bar{H}^{-1})=\frac{1}{c}\left(\sum_k\bar{a}_k^{-1}-(1-c)d\right)$.

Spectrum. The lower bound follows from $\bar{H}\succeq cD$. For the upper bound, let $k_1,k_2$ index the two smallest entries of $\bar{a}$ and pick a unit vector $w\in\operatorname{span}(e_{k_1},e_{k_2})$ orthogonal to $\bar{a}$, which exists as orthogonality is a single linear constraint on a two-dimensional space. Then $w^{\mathsf{T}}\bar{H}w=c\left(\bar{a}_{k_1}w_{k_1}^{2}+\bar{a}_{k_2}w_{k_2}^{2}\right)\le c\bar{a}_{(2)}$, and the Rayleigh principle gives $\lambda_{\min}(\bar{H})\le c\bar{a}_{(2)}$.

By Lemma 2, $\bar{H}$ averages two designs with weights $c$ and $1-c$: single-domain runs at $e_k$ drawn with probability $\bar{a}_k$, and the tempered prior repeated in every run. These two terms lead to two pitfalls of Dirichlet sampling (Figure 11). (i) Under-dispersion: a large concentration keeps the mixes close to the tempered prior, so $c$ is small and the information in every direction orthogonal to $\bar{a}$ shrinks by the factor $c$. (ii) Tail starvation: even at $c\to1$, rare domains with a small share $\bar{a}_k$ are seldom explored, so their attributions are estimated poorly and dominate the error, as $\operatorname{tr}(\bar{H}^{-1})$ in Lemma 2 grows with $\sum_k\bar{a}_k^{-1}$.

To quantify the second pitfall, we model the prior mix as Zipfian ([83]), $(\pi_{\mathrm{prior}})_k\propto k^{-s}$ for domains ranked by size (Figure 12). Tempering preserves this form, as $(k^{-s})^{1/\tau}=k^{-p}$, so the tempered prior is $\bar{a}_k=k^{-p}/\zeta_d(p)$ with the flatter exponent $p:=s/\tau$ and the truncated zeta function $\zeta_d(p):=\sum_{k=1}^{d}k^{-p}$. Its rarest domain has the share $\bar{a}_{\min}=d^{-p}/\zeta_d(p)$, which shrinks as $d$ grows, so by Lemma 2, $\lambda_{\min}(\bar{H})\ge c\,d^{-p}/\zeta_d(p)$. To bound the error of a sampled design, which depends on $H_{\Pi}$ rather than on $\bar{H}$, we use the following auxiliary results.

Proof. Expanding $(X^{-1}w-\bar{X}^{-1}w)^{\mathsf{T}}X(X^{-1}w-\bar{X}^{-1}w)\ge0$ gives $w^{\mathsf{T}}X^{-1}w\ge2w^{\mathsf{T}}\bar{X}^{-1}w-w^{\mathsf{T}}\bar{X}^{-1}X\bar{X}^{-1}w$, whose expectation is $w^{\mathsf{T}}\bar{X}^{-1}w$; summing over $w=e_1,\ldots,e_d$ gives the trace inequality.

$\lambda_{\min}\left(\sum_{i=1}^{m}\mathbb{E}[X_i]\right)$. Then, for every $t \in[0,1)$,

Proof. The first inequality is Theorem 1.1 of Tropp [67]. The second one follows from $(1-t)\ln(1-t) \ge t^2/2-t$ for $t \in[0,1)$, which holds as both sides vanish at $t=0$ and the derivative of their difference, $-\ln(1-t)-t$, is nonnegative.

Proof. The summands $X_i:=\pi_i\pi_i^{\mathsf{T}}$ of $mH_\Pi$ are independent and positive semidefinite with $\mathbb{E}[X_i]=\bar{H}$ and $\lambda_{\max}(X_i)=\lVert\pi_i\rVert_2^2\le R$, so Lemma 4 with $\mu_{\min}=m\lambda_{\min}(\bar{H})$ and $t=1/2$ bounds the probability of $\lambda_{\min}(H_\Pi)\le\frac{1}{2}\lambda_{\min}(\bar{H})$ by $d e^{-m\lambda_{\min}(\bar{H})/(8R)}\le\delta$, using the assumed lower bound on $m$. On the complementary event, $\lambda_{\min}(H_\Pi)>0$, so $\Pi$ has full column rank and the last bound of Eq. (8) gives the error bound.

Proof. Recall that $\bar{a}_{\min}=d^{-p}/\zeta_d(p)$ and $\lambda_{\min}(\bar{H})\geq c d^{-p}/\zeta_d(p)$.

Part (i). Every mix satisfies $\|\pi_i\|_2^2\leq\|\pi_i\|_1^2=1$, so Corollary 3 applies with $R=1$. Its condition on $m$ is implied by the assumed one, as $8\ln(d/\delta)/\lambda_{\min}(\bar{H})\leq8d^p\zeta_d(p)\ln(d/\delta)/c$, and its error bound reads

Part (ii). For $m\geq d$, the design has full column rank almost surely, because the Dirichlet distribution has a density on the simplex, which does not contain the origin, so any $d$ of its draws are almost surely linearly independent. Hence Eq. (8) applies almost surely, and Lemma 3 with $X=\Pi^\top\Pi$ and $\bar{X}=m\bar{H}$ gives

using the trace formula of Lemma 2 with $\bar{a}_k^{-1}=k^p\zeta_d(p)$. The claim follows from $\sum_{k=1}^{d}k^p\geq\int_0^d x^p\,dx=\frac{d^{1+p}}{1+p}$.

Rate. Since $\zeta_d(p)\geq d\cdot d^{-p}$, we have $d^{1+p}\zeta_d(p)\geq d^2$, so for $d\geq2(1+p)$ the subtracted term satisfies

and the right-hand side of Eq. (12) is at least $\frac{v_{\min}}{mc}\frac{d^{1+p}\zeta_d(p)}{2(1+p)}$. Together with Eq. (11), both bounds are within the factors $\frac{1}{2(1+p)}$ and $2$ of $\frac{\sigma^2}{mc}d^{1+p}\zeta_d(p)$ for $v_{\min}=\sigma^2=v_{\max}$.

Finally, comparing the sum $\zeta_d(p)$ with the integral of $x^{-p}$ gives $\zeta_d(p)=\frac{d^{1-p}}{1-p}(1+o(1))$ for $p<1$, $\zeta_d(p)=\ln d+O(1)$ for $p=1$, and $\zeta_d(p)=\zeta(p)(1+o(1))$ for $p>1$ with the Riemann zeta function $\zeta$, which yields Eq. (13). □

Theorem 1 shows that, compared to the optimal error $\sigma^2d^2/m$ on the simplex (Proposition 3 below), under-dispersion increases the error of Dirichlet sampling by a factor of order $1/c$, and tail starvation by a further factor $d^{p-1}$ when $p>1$, i.e., when the prior is too skewed for the temperature $\tau$ to flatten (Figure 11). Both factors also appear in the lower bound of part (ii), so the two bounds match up to constants. For the uniform prior of Theorem 2, $p=0$ and only the factor $1/c$ remains. In our experiments, tail starvation is negligible: the Zipf fits in Figure 12 give $s\approx0.6$ at 24 domains and $s\approx0.9$ at 288 and 40 domains, so with $\tau=2$ all three setups lie in the regime $p<1$. Under-dispersion, however, is forced by the caps. Without caps, a smaller concentration always helps, as it increases $c$, and the range $\gamma\in[0.1,10]$ used at 24 domains gives $c=0.5$. At 288 and 40 domains, small concentrations produce sparse mixes that violate the caps and are almost all rejected, so we tune the ranges with the criterion of § 4.1 to $\gamma\in[40,4000]$ and $[100,10^4]$, for which $c\approx0.005$ and $0.002$. By Lemma 2 and the proof of part (ii), the expected error of the uncapped Dirichlet design is then at least $\operatorname{tr}(\bar{H}^{-1})/d^2\approx2$, $230$, and $540$ times the optimum $\sigma^2d^2/m$ for the empirical priors and $\gamma$ ranges of our three setups. Part of this gap is unavoidable, as under equal caps every design on $\mathcal{P}_{\mathrm{cap}}$ pays a factor of about $d/(k-1)$ (Theorem 2). Under the equal-cap approximation, Dirichlet sampling thus remains about 4 and 30 times worse than the best capped design at 288 and 40 domains. This matches Figure 3, where both samplers fit well at 24 domains but no surrogate fits the Dirichlet swarm at 288, and § 4.2, where near-prior Dirichlet mixes vary HellaSwag BPB five times less than vertex mixes.

On optimal sampling

Without caps, the vertices of the simplex are the pure-domain mixes $e_1,\ldots,e_d$, and the greedy fill of Eq. (3) places each run at a uniformly random vertex. Allocating exactly $m/d$ runs to each vertex gives the balanced vertex design, the classical optimal design for first-degree mixture models [59, 37]. We restate it because its error $\sigma^2d^2/m$ is the reference optimum in Appendix E.3 and Theorem 2 builds on its proof.

runs on domain $k$, and

with equality for homoscedastic noise. Among all designs with $m$ rows in $\Delta^{d-1}$, the balanced vertex design is, up to the order of the runs, the unique design that minimizes $\operatorname{tr}((\boldsymbol{\Pi}^{\mathsf{T}}\boldsymbol{\Pi})^{-1})$, maximizes $\det(\boldsymbol{\Pi}^{\mathsf{T}}\boldsymbol{\Pi})$, maximizes $\lambda_{\min}(\boldsymbol{\Pi}^{\mathsf{T}}\boldsymbol{\Pi})$, and minimizes the worst-case prediction variance $\max_{\boldsymbol{\pi}\in\Delta^{d-1}}\boldsymbol{\pi}^{\mathsf{T}}(\boldsymbol{\Pi}^{\mathsf{T}}\boldsymbol{\Pi})^{-1}\boldsymbol{\pi}$.

Proof. Balanced vertex design. Each row of $\boldsymbol{\Pi}$ is a standard basis vector, so $(\boldsymbol{\Pi}^{\mathsf{T}}\boldsymbol{\Pi})_{kl}$ counts the runs placed at $e_k$ if $k=l$ and vanishes otherwise, giving $\boldsymbol{\Pi}^{\mathsf{T}}\boldsymbol{\Pi}=\frac{m}{d}I_d$. Then $\hat{\boldsymbol{\beta}}=\frac{d}{m}\boldsymbol{\Pi}^{\mathsf{T}}\boldsymbol{y}$ averages the outcomes of the runs on each domain, and $\operatorname{tr}((\boldsymbol{\Pi}^{\mathsf{T}}\boldsymbol{\Pi})^{-1})=d^2/m$, so Eq. (14) and its equality case follow from Lemma 1.

Optimality. Let $\boldsymbol{\Pi}$ be any design with rows $\boldsymbol{\pi}_i\in\Delta^{d-1}$ and $M:=\boldsymbol{\Pi}^{\mathsf{T}}\boldsymbol{\Pi}$ positive definite with eigenvalues $\lambda_1,\ldots,\lambda_d>0$; designs with a singular $M$ are never optimal, as all four quantities are then degenerate. The diagonal of $M$ sums to

with equality if and only if every row has a single nonzero entry, i.e., is a vertex. For the trace of the inverse, the Cauchy–Schwarz inequality $(\sum_j\lambda_j)(\sum_j\lambda_j^{-1})\ge d^2$ and Eq. (15) give $\operatorname{tr}(M^{-1})\ge d^2/\operatorname{tr}(M)\ge d^2/m$, with equality if and only if all eigenvalues are equal and $\operatorname{tr}(M)=m$, i.e., $M=\frac{m}{d}I_d$. For the determinant, the inequality of arithmetic and geometric means gives $\det(M)=\prod_j\lambda_j\le(\operatorname{tr}(M)/d)^d\le(m/d)^d$, with the same equality condition. For the smallest eigenvalue, $\lambda_{\min}(M)\le\operatorname{tr}(M)/d\le m/d$, again with the same equality condition. For the worst-case prediction variance, $\boldsymbol{\pi}\mapsto\boldsymbol{\pi}^{\mathsf{T}}M^{-1}\boldsymbol{\pi}$ is convex, so its maximum over the simplex is attained at a vertex, and $\max_k(M^{-1})_{kk}\ge\frac{1}{d}\operatorname{tr}(M^{-1})\ge d/m$, with equality if and only if $M^{-1}$ has a constant diagonal and $\operatorname{tr}(M^{-1})=d^2/m$, which again forces $M=\frac{m}{d}I_d$. Finally, $M=\frac{m}{d}I_d$ requires every row to be a vertex by the equality case of Eq. (15), and then $M=\operatorname{diag}(m_1,\ldots,m_d)$ with $m_k$ runs at $e_k$ forces $m_k=m/d$, i.e., the balanced vertex design up to the order of its rows, which attains all four bounds.

Proposition 3 formalizes why vertex sampling works: measuring one domain at a time makes the estimates $\hat{\beta}_k$ uncorrelated for independent runs, and measuring every domain equally often gives each the same variance, at most $v_{\max}d/m$. The balanced vertex design is optimal under all classical criteria at once (A, D, E, and the worst-case prediction variance; [56]), which in general favor different designs, so the choice of criterion does not matter here. Since $\boldsymbol{\pi}_{\mathrm{target}}$ is chosen by optimizing the prediction $\langle\hat{\boldsymbol{\beta}},\boldsymbol{\pi}\rangle$ over mixes, the last criterion matters most for data mixing: the prediction of every mix $\boldsymbol{\pi}$ has variance at most $v_{\max}d/m$, which no other design can improve in the worst case. In practice, the caps are smaller than one, so a run cannot train on a single domain, and LinMix samples the vertices of $\mathcal{P}_{\mathrm{cap}}$, each of which fills several domains up to their caps. Theorem 2 shows that the expected information matrix of vertex sampling is optimal on the capped simplex for a uniform prior.

random such design, satisfies

with equality for $H=\overline{V}$, so that the expected design of vertex sampling minimizes $\operatorname{tr}(H^{-1})$ and maximizes $\lambda_{\min}(H)$ and $\det(H)$ over $P_{\mathrm{cap}}$.

(iii) Under the noise conditions of Assumption 1, for any $\delta\in(0,1)$, if $m\ge\frac{8\kappa(d-1)}{\kappa-1}\ln(d/\delta)$, then with probability at least $1-\delta$ over the draw of $\Pi$, $\lambda_{\min}(H_{\Pi})\ge\frac{\kappa-1}{2d(d-1)}$ and

Proof. Part (i). With equal caps $u:=\kappa/d=1/q$, the greedy fill assigns the cap to the first $q$ domains of the permutation $\sigma_i$, which exhausts the budget as $qu=1$, and zero to the remaining domains. The first $q$ elements of a uniformly random permutation form a uniformly random $q$-subset $S_i$, so $\pi_i=u\,\mathbf{1}_{S_i}$ with $\Pr[k\in S_i]=q/d$ and $\Pr[k,l\in S_i]=q(q-1)/(d(d-1))$ for $k\ne l$. Hence

whose eigenvalue along $\mathbf{1}_d$ is $u^2\left(\frac{q}{d}+\frac{q(q-1)}{d}\right)=u^2q^2/d=1/d$, and whose eigenvalue on the orthogonal complement is $u^2\frac{q}{d}\frac{d-q}{d-1}=\frac{q}{d}\frac{d-q}{d(d-1)}=\frac{\kappa-1}{d(d-1)}$, using $u^2q=1/q$ and $q=d/\kappa$. The trace of $\overline{V}^{-1}$ follows by summing the reciprocals with multiplicities $1$ and $d-1$.

Part (ii). Let $H=\sum_jw_j\pi_j\pi_j^\top$ be positive definite with weights $w_j\ge0$ summing to one and $\pi_j\in P_{\mathrm{cap}}$. Two constraints hold for every such $H$. First, every mix sums to one, so $\mathbf{1}_d^\top H\mathbf{1}_d=\sum_jw_j(\pi_j^\top\mathbf{1}_d)^2=1$, i.e., the unit vector $v:=\mathbf{1}_d/\sqrt{d}$ satisfies $v^\top Hv=1/d$. Second, for $\pi\in P_{\mathrm{cap}}$, $\lVert\pi\rVert_2^2=\sum_k\pi_k^2\le\sum_k u\pi_k=u$, with equality if and only if every $\pi_k\in\{0,u\}$, i.e., $\pi$ is a vertex, so $\operatorname{tr}(H)\le u=\kappa/d$. Let $p_1,\ldots,p_{d-1}$ be an orthonormal basis of the orthogonal complement of $v$, collected in $P\in\mathbb{R}^{d\times(d-1)}$, so that $Q:=[v\ P]$ is orthogonal, and let $C:=P^\top HP$ be the lower-right block of $Q^\top HQ$. The two constraints give $\operatorname{tr}(C)=\operatorname{tr}(H)-v^\top Hv\le(\kappa-1)/d$. For the smallest eigenvalue, restricting the Rayleigh quotient of $H$ to the range of $P$ gives $\lambda_{\min}(H)\le\lambda_{\min}(C)\le\operatorname{tr}(C)/(d-1)\le(\kappa-1)/(d(d-1))$.

For the determinant, Fischer’s inequality for the positive definite block matrix $Q^\top HQ$ (Horn and Johnson, [30], Section 7.8) and the inequality of arithmetic and geometric means give $\det(H)=\det(Q^\top HQ)\le(v^\top Hv)\det(C)\le\frac{1}{d}\left(\frac{\operatorname{tr}(C)}{d-1}\right)^{d-1}\le\frac{1}{d}\left(\frac{\kappa-1}{d(d-1)}\right)^{d-1}$. For the trace of the inverse, $\operatorname{tr}(H^{-1})=\operatorname{tr}(Q^\top H^{-1}Q)=v^\top H^{-1}v+\sum_{j=1}^{d-1}p_j^\top H^{-1}p_j$. For any unit vector $x$, the Cauchy–Schwarz inequality $1=(x^\top H^{1/2}H^{-1/2}x)^2\le(x^\top Hx)(x^\top H^{-1}x)$ gives $x^\top H^{-1}x\ge1/(x^\top Hx)$. Applied to $v$, this gives $v^\top H^{-1}v\ge d$; applied to each $p_j$ and combined with $(\sum_jx_j)(\sum_j1/x_j)\ge(d-1)^2$ for positive $x_j$, it gives $\sum_jp_j^\top H^{-1}p_j\ge\sum_j1/(p_j^\top Hp_j)\ge(d-1)^2/\operatorname{tr}(C)\ge d(d-1)^2/(\kappa-1)$. Finally, $\overline{V}$ lies in the convex hull as an expectation, and by part (i) it has $v$ as an eigenvector with eigenvalue $1/d$ and $C=\frac{\kappa-1}{d(d-1)}I_{d-1}$, so all three bounds hold with equality.

Part (iii). Every vertex has $\lVert\pi_i\rVert_2^2=qu^2=u$, so Corollary 3 applies with $R=u=\kappa/d$ and $\lambda_{\min}(\overline{V})=\frac{\kappa-1}{d(d-1)}$ from part (i), for which its condition on $m$ and its error bound are exactly those claimed. $\square$

Theorem 2 extends Proposition 3 to the capped simplex: without caps, $q=1$ and part (ii) reduces to the bounds of the proposition. The caps force each run to spread over $q=d/\kappa$ domains, so a run reveals only about $1/q$ as much about any single domain. This is why $\operatorname{tr}(\overline{V}^{-1})$ in part (i) exceeds the uncapped value $d^2$ by a factor of about $d/(\kappa-1)$, which is 4.5, 62, and 16 in our three setups if we treat their caps as equal. Vertex sampling is nevertheless optimal for two reasons. First, every mix sums to one, so all designs carry the same information about the overall outcome level ($1/d$ along $\mathbf{1}_d$) and differ only in the information about contrasts between domains. Vertices attain the largest norm $\lVert\boldsymbol{\pi}\rVert_2^2 = \kappa/d$ on $\mathcal{P}_{\mathrm{cap}}$, so a design of vertices has the largest total information $\operatorname{tr}(H)$, which leaves vertex sampling the most information for these contrasts; mixes that are not vertices have a strictly smaller norm and are never optimal. Second, it spreads this information evenly across these contrasts, which is what the three quantities in Eq. (16) reward. § 3 calls the design near-orthogonal in expectation because the expected off-diagonal entries of $\Pi^\top\Pi$ are only $(q-1)/(d-1) \approx1/\kappa$ times the diagonal ones. Part (iii) shows that a sampled design is nearly optimal once $m$ exceeds a multiple of $d\ln d$: its estimation error is at most $2d/(d-1) \approx2$ times the smallest value any design on $\mathcal{P}_{\mathrm{cap}}$ can achieve under homoscedastic noise. This guarantee is conservative, as its constants come from the matrix Chernoff bound and its condition on $m$ exceeds the size of our swarms; in practice, we select the best-conditioned of many sampled designs (Appendix § 10), which controls the realized design directly but is not covered by the theorem. Pairing runs by reversed permutations, as in § 3, does not affect the theorem. The reverse of a uniform permutation is uniform, so $\bar{V}$ stays the same, and for $\kappa\ge2$ the two runs of a pair have disjoint supports, so part (iii) still holds with the $m/2$ independent pairs as summands, each with expectation $2\bar{V}$ and $R=u$. We leave non-uniform priors with unequal caps to future work, as they break the symmetry that gives $\bar{V}$ its closed form.

Limitations

A limitation shared by all regression-based methods is the cost of training a proxy swarm, which we quantify here by scaling proxy model size and budget. In practice, the cost-to-performance trade-off of data mixing becomes only more favorable as the target model size and pretraining budget increase. For example, pretraining a 4B model on 120B tokens for 2.5k GPU hours could be designed with a swarm of $512 \times30$M models computed for 0.1k GPU hours (4% of the target cost; cf. Appendix § 14). Naturally, we train the target models using a single random repetition to enable more experimental setups. Our theory bounds the error of the estimated attribution $\hat{\beta}$, not the error of the predictions made with it. Moreover, the noise level $v_{\max}$ in our bounds depends on how the corpus is split into domains, so the $d^2/m$ rate alone cannot tell how many domains to use. Finally, whether the optimal mix is dense in positively attributed domains depends on the target benchmark, corpus size, and pretraining token budget. We thus view the curvature $r$ that controls this as a useful utility trade-off, while the hyperparameter-free property of LINMIX applies to the design of proxy-mix sampling and regression model fitting.

Compute Resources

All models are trained on NVIDIA GH200 96GB GPUs without special optimizations beyond the default code implementations. Training one OLMo2 proxy model takes 0.2 GPU-hours for 30M parameters (1B tokens), 0.7 for 60M (2B), and 1.7 for 100M (3.2B). The Qwen3.5-like 100M proxy takes 2.6 GPU-hours on the same 3.2B tokens, 1.5 times as much as OLMo2, which we attribute to its hybrid attention with gated DeltaNet layers that add constant overhead independent of model size. Model evaluation adds further cost per proxy, depending on the benchmark and its sample size. The proxy swarms cost 0.7k GPU-hours at 30M (1024 models at 24 domains and 2048 at 288 domains), 1.4k at 60M (2048 models), 1.8k at 100M for OLMo2 (1024 models), and 2.7k at 100M for Qwen3.5 (1024 models). Pretraining a single 1B target model for 32B tokens takes 233 GPU-hours using 16 GPUs, so the 23 runs behind Table 1 and the 7 runs behind Table 5 and 6 amount to ca. 7.0k. The 1B Qwen3.5-like target model takes 272 GPU-hours on the same 32B tokens, so the four runs behind the CPile column of Table 2 amount to ca. 1.1k. One 4B target run (80B tokens) takes 1.7k GPU-hours using 32 GPUs, so the 10 runs behind Table 2 amount to ca. 17k. Together with model evaluations, the reported experiments cost ca. 34k GPU-hours; we estimate that the whole project consumed ca. 51k GPU-hours, including preliminary work and failed jobs.