When do data mixtures improve scaling laws? Insights from high-dimensional regression

1 Introduction

• Prior work has examined empirical mixture selection and scaling laws, but the conditions for faster scaling than either dataset alone remain unresolved.Existing theoretical results include cases where mixture composition changes only the prefactor, as well as models predicting improved exponents.Source: §1 · p. 1
• The study models multiple domains with a shared regression function, heterogeneous covariances and noise levels, and sample sizes that may grow at different rates.It characterizes minimax risk for general covariances and mixed-ridge performance for commuting covariances.Source: §1 · pp. 1–2
• Mixing datasets can strictly improve the scaling exponent over either dataset alone, under conditions on power-law spectra and relative dataset growth.The analysis also identifies regimes where ridge regression attains the minimax-optimal exponent.Source: §1 · p. 2
• Decoder-only transformer simulations show faster decreases in target-domain test loss when target and auxiliary data are mixed than when either dataset is used alone.The experiments are described as qualitatively similar to the theory.Source: §1 · p. 2

2 Related work

• Empirical work selects or adapts domain weights, predicts mixture performance, and models loss as a function of domain proportions and training scale.Examples include DoReMi, DoGE, Chameleon, Skill-It, Aioli, ADO, RegMix, MixMin, and AutoScale.Source: §2 · p. 2
• Theoretical work studies distribution-weighted predictors, adaptive sampling, and mixture gains in nonparametric regression, while faster scaling exponents remain unestablished for high-dimensional regression.This paper accommodates decaying covariance spectra and dataset sizes growing at different rates.Source: §2 · p. 2

3 Problem setup

• Each of K independent datasets contains n_i samples from a domain distribution P_i on R^d, with mean-zero features and covariance C_i.The dimension may be infinite, with C_i trace-class on ℓ2.Source: §3 · p. 2
• All domains share the same target parameter, and the test distribution can be restricted to a single target domain using π*=e1.The setup evaluates excess test risk for an estimator.Source: §3 · pp. 2–3
• The paper characterizes minimax risk and the risk of a mixed ridge estimator, whose conditional noise-averaged risk has a bias-variance decomposition.Importance-weighted objectives also describe domain reweighting.Source: §3 · p. 3
• The design assumption requires concentration for every domain and positive semidefinite operator, and covers independent sub-Gaussian coordinates and convex Lipschitz concentration.The shared-parameter model also corresponds to kernel ridge regression when inputs are mapped through a reproducing-kernel feature map.Source: §3 · p. 3

4 General results

• The minimax risk is bounded above and below by constant multiples of L(M), and is equivalent to L(M) when E0 is bounded.The stated bounds are L(M)/π² and (1 + CE0)L(M); E0 is O(1) under the conditions in Remark 4.1.Source: §4 · pp. 3–4
• At the minimax-rate level, heterogeneous datasets are equivalent to a single Gaussian sequence model with information operator M.The proof compares the regression model with a Gaussian sequence model and controls additional variance from random designs.Source: §4 · p. 4
• Each coordinate’s estimation cost depends on its test weight relative to the information from all domains; unestimated coordinates incur a radius-weighted cost.Domain information is greater along directions where its covariance is large relative to the ellipsoid constraint.Source: §4 · p. 4
• For commuting covariances, mixed-ridge error admits deterministic equivalents, derived by controlling coupled covariance-weighted resolvent traces.The deterministic proxy replaces each dataset with its population covariance scaled by an effective sample size and deflated by the effective dimension it must fit.Source: §4 · pp. 4–5
• The ridge error bound is additive O(eK), and Theorem 2 assumes commuting covariances; multiplicative and noncommuting guarantees remain conjectural.Under power-law spectra and optimal ridge regularization, the paper states that the bound is o(R).Source: §4 · p. 5

• The model specializes to a target test distribution and an auxiliary domain with aligned power-law spectra, where δ = α1 − α2 > 0 means the auxiliary spectrum has heavier tails.The target regularity is s in the target geometry; the model assumes α1, α2 > 1 and 2s > δ.Source: §5 · p. 5
• The minimax scaling laws compare target-only, auxiliary-only, and mixed risks, with individual-dataset rates expressed through exponents Γtar and Γaux.For δ = 1, the auxiliary-only rate includes a log n factor.Source: §5 · p. 5
• A linear estimator attains the minimax rate but requires knowledge of the population covariances; the analysis then asks when ridge matches it without that knowledge.Source: §5 · p. 5
• The mixed minimax rate beats both single-dataset rates if and only if δ > 1 and γ2 lies in (γc, 1).Here γc = 1 − δ/(1 + 2s).Source: §5 · pp. 5–6
• For optimal ridge, the approximation error is o(R), so the deterministic-equivalent rate also holds for the mixed estimator’s excess test risk.Source: §5 · p. 6
• Ridge achieves the improved mixed minimax exponent whenever δ > 1 and γc < γ2 < 1, for every regularity s, including s > α1.In this regime, n2 = o(n); γ2 > γc supplies enough auxiliary samples to improve on target-only risk, while γ2 < 1 leaves target data useful relative to auxiliary-only training.Source: §5 · pp. 6–7
• Ridge may miss the minimax rate when source regularity exceeds the capacity exponent of the spectrum it fits; the truncations involve min{s, α1} and min{s, α2}.The authors expect iterated Tikhonov or early stopping to remove these truncations.Source: §5 · p. 7

Linear models

• In linear-model panel (a), mixed ridge matches the mixed minimax rate and strictly improves on either dataset’s minimax rate, while panels (b) and (c) show no improvement.In those panels, the mixed rate coincides with the auxiliary-only and target-only rates, respectively.Source: §6 · p. 7
• On CIFAR-10 and ImageNet-100 features, the deterministic equivalent remains valid beyond the technical assumptions made for Theorem 2.The result is reported in Figure 2.Source: §6 · p. 7

Language modeling

• The language-model experiment trains an 81.5M-parameter GPT-2-style model from scratch on SlimPajama, using StackExchange as the target domain and n tokens from that domain.Figure 3 plots target-domain test NLL against reference target budget n in panels (a)–(c) and compares normalized target and auxiliary token frequencies in panel (d).Source: §6 · p. 8, Figure 3 · p. 8
• For γ2 = 1.25, the language-model mixture improves the scaling law over either dataset alone, but this improvement is not noticeable for γ2 = 0.75 or 1.75.The auxiliary domain also assigns relatively more probability to tokens infrequent in the target; mixtures with closer token-frequency distributions do not improve the scaling law.Source: §6 · p. 8

7 Conclusion

• Mixing improves the scaling exponent when the auxiliary spectrum is heavy tailed (δ > 1) and its sample size grows at an appropriate rate, γc < γ2 < 1.Whenever the mixed minimax rate improves on either dataset alone, ridge regression is minimax optimal; transformer experiments show qualitatively agreeing behavior.Source: §7 · pp. 8–9

A.1 Proof of Theorem 5

• Theorem 5 gives an exact Gaussian linear minimax error characterization for any positive semidefinite N.The proof introduces a variational characterization of that error.Source: A Proof of Theorem 1 · p. 13

A.2 Proof of Theorem 1

• The constructed regression estimator has worst-case risk bounded by vL(M), where v = 1 + κE0 and κ depends on the concentration constants in Assumption 1.The proof controls the additional variance from random designs uniformly over the parameter ball.Source: §A.2 · pp. 16–17
• A van Trees-type lower bound and the constructed upper bound show that regression minimax risk is within constant factors of the Gaussian sequence minimax risk.The conclusion states that their unrestricted risks have the same order when the concentration constants and E0 are uniformly bounded.Source: §A.2 · p. 20

A.3 Extension to infinite dimension

• In infinite dimension, finite-rank approximations and continuity of the variational objective extend the lower-bound argument to arbitrary feasible A.The proof concludes with R* ≥ π^-2L(M/π^2) ≥ π^-2L(M).Source: §A.3 · p. 20

B.1 Invertibility of L

• Nonnegativity and symmetry of K imply that L is positive definite and therefore invertible.The proof uses the fixed-point equations and a weighted Cauchy–Schwarz bound.Source: §B.1 · p. 21

B.2 Proof of Theorem 2

• Theorem 2’s variance calculation applies Theorem 6’s third bound with Cπ* and the target dataset i.The proof averages over independent, centered training noises in its preceding calculation.Source: §B.2 · pp. 21–22

C Proof of Theorem 3

• Restricting the supremum to diagonal operators preserves its value because diagonalization keeps feasibility and does not decrease the objective.The proof then expresses the optimization through the operators’ diagonal entries.Source: C Proof of Theorem 3 · p. 23
• A cutoff where the variance and bias terms cross yields matching upper and lower bounds for the minimax risk.The proof chooses the first index m* with V_m* ≥ T_m* and bounds the risk by the corresponding bias term.Source: C Proof of Theorem 3 · p. 24
• The proof establishes distinct rates for target-only and auxiliary-only data and derives mixed-data rates across the remaining parameter regimes.It states that the cases cover all parameter regimes and establish the claimed rates.Source: C Proof of Theorem 3 · pp. 24–25
• At the boundary γ₂ = γc, the mixed-data risk is n^-Γtar, without a logarithmic factor.The preceding target-data cutoff applies at this boundary.Source: C Proof of Theorem 3 · p. 25

D Proof of Theorem 4

• For each regime, an explicit feasible λ_n with B0(λ_n) ≍ V0(λ_n) and ε_n = o(1) yields the tight scaling law of optimal ridge risk.The proof verifies feasibility for the balanced cutoff sequences used in its constructions.Source: D Proof of Theorem 4 · pp. 26, 34
• The proof establishes (20) by showing μ_i(λ_n)/n_i = Θ(1) and d_j(λ_n) ≍ h_j + λ_n, reducing risk control to the two spectral sums.These comparisons yield V1 ≍ V0(λ_n), and the upper and lower bounds establish (20).Source: D Proof of Theorem 4 · pp. 27–28
• When δ > 0, J separates target-dominated low frequencies from auxiliary-dominated high frequencies if n2 ≤ n; if n2 > n, the auxiliary dataset dominates throughout.For m ≥ J, the bias satisfies B0(h_m) ≍ m^-2s.Source: D Proof of Theorem 4 · p. 30
• When δ = 0, the bias contribution has order m^-2aℓ1(m), while the variance sum has order m/(n + n2).This case has α1 = α2 and a = b.Source: D Proof of Theorem 4 · p. 32
• When δ < 0, both bias contributions are bounded by a constant times m^-2s, and the tail term gives B0(h_m) ≍ m^-2s.For n2 ≥ n, the auxiliary dataset dominates below J and the target dataset above J; for n2 < n, J = 1 and the target dominates throughout.Source: D Proof of Theorem 4 · p. 32
• For the stated regime, the polynomial exponent is Γ = 2a/(1+2a), with J ≍ n^(1/(1+2a)) and J_n ≍ J^(-2a).These relations identify the exponent and cutoff scale in the proof.Source: D Proof of Theorem 4 · p. 34
• In the fixed-signal critical case s = α1, δ = 1, and γ2 = γ+c, the scaling has an additional factor of order log log n.The proof then turns to the regime γ2 > γ+c.Source: D Proof of Theorem 4 · p. 35
• When δ = 0, the two datasets have the same spectral decay, and the proof identifies this as the optimal-scaling case.The following exponent expression is incomplete in the supplied passage, so no exponent value is stated here.Source: D Proof of Theorem 4 · p. 36
• Every constructed cutoff has m → ∞ and B0(hm) ≍ V0(hm), yielding feasibility ϵn = o(1); the displayed logarithmic factors are bounded between a positive constant and O(1 + log n).Applying (20) with λn = hm proves the stated polynomial exponents.Source: D Proof of Theorem 4 · p. 37

D.1 A technical lemma

• The technical lemma gives power-law sum comparisons: prefix sums grow as m^u for u > 0, as 1 + log m for u = 0, and remain bounded for u < 0; tail sums diverge for u ≥ 0.The interval-sum estimate also extends to real endpoints by replacing them with their integer floors.Source: §D.1 · p. 37

E Proof of Lemma 1

• High-probability bounds on the normalized covariance errors support the risk approximation, while the fixed-point comparison is deterministic.The proof reports probability bounds for the covariance estimates and states that the second comparison is deterministic.Source: E Proof of Lemma 1 · pp. 38, 42, 43
• The risk comparison reduces to controlling normalized empirical covariance errors and comparing fixed-point quantities with their deterministic counterparts.The resulting bound compares R₁(λₙ) with B₀ + Vσ, with relative error controlled by ε and λₙε.Source: E Proof of Lemma 1 · pp. 39, 44

F.1 Heuristic derivation of deterministic equivalences

• The heuristic derivation applies first- and second-order deterministic equivalents to obtain the ridge test-error bias and variance formulas.The derivation identifies these formulas as the main results of Theorem 2.Source: §F.1 · p. 45

F.2 Main theorem and proof sketch

• Theorem 6 establishes deterministic equivalents for three functionals from Theorem 2 under commuting covariances and the stated assumptions.The theorem assumes each covariance satisfies Assumption 1 and the sample-size regularization parameters satisfy Assumption 2.Source: §F.2 · p. 46
• With fixed K and comparable dataset sizes n_i ≍ n, all three approximation errors ε_1, ε_2, and ε_3 scale as O(n^-1/2).Source: §F.2 · p. 46

F.3.1 Deterministic part of Tr(AG)

• For Tr(AG), the proof splits error into deterministic and martingale parts, using leave-one-out resolvents and conditional estimates for the deterministic part.Source: §F.2 · p. 46, §F.3.1 · p. 47

F.3.2 Martingale part of Tr(AG)

• The martingale bound for Tr(AG) uses a filtration over observations, conditional concentration, and deleted-resolvent bounds.Source: §F.3.2 · pp. 48–49

F.4 Deterministic equivalence of Tr(AGCkG)

• The proof of the Tr(AGC_kG) deterministic equivalent likewise separates deterministic and martingale contributions.Source: §F.4 · p. 49

F.4.1 Deterministic part of Tr(AGCkG)

• For the deterministic part of Tr(AGC_kG), covariance-weighted row-increment bounds control the errors before the estimates are combined.Source: §F.4.1 · pp. 52–53

F.4.2 Martingale part of Tr(AGCkG)

• For the martingale part of Tr(AGC_kG), the proof separates cross and product terms and combines their bounds with the deterministic estimate.Source: §F.4.2 · pp. 53–55

F.5.1 Deterministic part of Tr(AGX⊤

• For the deterministic part of Tr(AGX⊤), the proof replaces a denominator by (1 + κ_k)^2 and approximates Tr(AHC_kH) by Tr(AGC_kG).Source: §F.5.1 · pp. 55–56

F.5.2 Martingale part of Tr(AGX⊤

• Combining moment and deterministic bounds yields probabilistic bounds for the martingale terms via Markov’s inequality.The proof first establishes bounds for the martingale components and then sums them; the stated probability guarantee holds for sufficiently large M and n_min.Source: §F.5.2 · p. 57

F.6 Technical lemmas

• For independent features satisfying Assumption 1, Lemma 4 gives a high-probability bound for PSD matrices independent of those features.The passage states probability at least 1 − n^-D over the randomness of the two features.Source: §F.6 · p. 58
• The technical lemmas also establish resolvent and trace bounds when up to two observations are removed.The stated bounds apply to deterministic PSD matrices; the probabilistic trace bound follows using resolvent estimates and a union bound over datasets.Source: §F.6 · p. 60

G.1 Details of numerical experiments

• The target data consists of StackExchange documents; the OOD auxiliary mixture uses six other SlimPajama domains with weights 0.637, 0.099, 0.095, 0.076, 0.048, and 0.045, respectively.The domains are GitHub, C4, arXiv, Common Crawl, Books, and Wikipedia.Source: §G.1 · p. 62
• For the near-target auxiliary construction, the mixture combines 90% independent StackExchange data with 10% of the six-domain OOD mixture.Source: §G.1 · p. 62
• Documents are normalized and split by normalized-content hashes at 98%/1%/1%, then tokenized, packed into 512-token blocks, and filtered for incomplete or low-diversity blocks.Reported dataset sizes count prediction tokens remaining after preprocessing.Source: §G.1 · p. 62
• All reported runs use an 81.5M-parameter decoder-only transformer and match the cosine learning-rate schedule to each run’s total token budget.The budgets are H = n for Target-only, H = n2 for Aux-only, and H = n + n2 for Mixture.Source: §G.1 · p. 62
• Figure 4 plots target-test NLL against reference target budget and compares normalized token frequencies; Table 1 reports target-test OLS scaling slopes and mixture improvements against both single-domain baselines.More-negative slopes indicate faster target-NLL decrease.Source: Table 1 · p. 63, Figure 4 · p. 63

G.2 Experiments on datasets with close distributions

• For γ2 = 1.25, the mixture does not improve the scaling law when target and auxiliary token frequencies are close.Table 1 shows this phenomenon across different γ2 settings; the result supports the finding that improvement requires sufficiently different spectral decay rates.Source: §G.2 · p. 63