Explaining Data Mixing Scaling Laws

Summary (Overview)

• Unified theoretical framework: The authors propose a theoretical framework to explain data mixing scaling laws, extending two single-domain theoretical perspectives (Quantization Model and Linear Regression Model) to multi-domain settings.
• Two key mechanisms identified: The framework identifies Capacity Competition (finite model capacity couples domain losses globally) and Noise Reduction (optimal weights shift toward harder-to-learn domains) as the two governing factors of domain losses.
• Superior empirical performance: The proposed models achieve lower Mean Relative Error (MRE) than existing empirical baselines (1.53% vs. 2.21% for the best baseline) while using significantly fewer free parameters.
• Cross-scale extrapolation: The framework successfully extrapolates across scales—predicting highly effective mixtures for large, unseen scales using parameters fitted exclusively on smaller ones.
• "Shared Head, Disjoint Tail" assumption: The paper formalizes the natural structure that domains overlap on fundamental skills while diverging on specialized skills, enabling tractable theoretical analysis.

Introduction and Theoretical Foundation

Background and Motivation

Large foundation models are trained on multi-domain data mixtures, where the proportion of each domain (the "data mixture") critically impacts performance. Discovering optimal mixtures is costly and lacks principled methodologies. Existing approaches fall into two paradigms:

• Online adaptation: Dynamically adjust domain weights during training (e.g., ODM, Skill-it, Aioli, ADO, PiKE)
• Offline prediction: Fit empirical scaling laws to predict loss landscapes a priori (e.g., Shukor et al., 2026; Ye et al., 2025; Ge et al., 2025)

The paper identifies a critical gap: while empirical laws exist, they act as "black boxes" with unclear generalization properties and lack rigorous theoretical understanding of domain interaction mechanics.

Theoretical Foundations

The paper builds on two primary theoretical frameworks for standard neural scaling laws:

1. Quantization Model (Michaud et al., 2023) : Frames learning as sequential acquisition of discrete "quanta" of skills following a Zipfian distribution:

$$p(q_k) \propto k^{-\alpha}, \quad (\alpha > 1)$$

The loss scales as:

$$L(N) = c \sum_{k=N+1}^{\infty} p(q_k) \approx c \cdot N^{-(\alpha-1)}$$

2. Linear Regression Model (Lin et al., 2024a; Bordelon et al., 2024) : Derives scaling laws from spectral analysis of data covariance matrices under one-pass SGD. The eigenvalues follow a power law $\lambda_k \propto k^{-\alpha}$, and the loss decomposes into model scaling and data scaling terms:

$$L(N, D) \approx \underbrace{O\left(\frac{1}{N^{a_1}}\right)}_{\text{Model Scaling}} + \underbrace{O\left(\frac{1}{D^{a_2}}\right)}_{\text{Data Scaling}} + E$$

Problem Setup

Given $K$ data domains, a mixture weight vector $h \in \Delta^{K-1}$ (probability simplex), model size $N$, and training tokens $D$, the goal is to predict domain loss $L_i(h \mid N, D)$ and find the optimal mixture:

$$h^{*} = \arg\min_{h \in \Delta^{K-1}} \sum_{i=1}^{K} w_i f_i(h \mid N, D)$$

Existing empirical laws include:

• Additive Law (Shukor et al., 2026): $L_i \approx E_i + \left(\sum_{j=1}^{K} C_{ij} h_j^{\gamma_{ij}}\right)^{-1}$
• Exponential Law (Ye et al., 2025): $L_i \approx c_i + k_i \exp\left(\sum_{j=1}^{K} t_{ij} h_j\right)$

Methodology

The "Shared Head, Disjoint Tail" Structure

The paper's central assumption: domains share a common foundation of basic knowledge (head) while diverging in specialized topics (tail):

• Power-Law Distribution: Skills within each domain follow a power-law frequency distribution
• Shared Head: High-probability fundamental skills overlap across domains (e.g., basic grammar, logic, arithmetic)
• Disjoint Tail: Rare, specialized knowledge becomes increasingly domain-specific

Extended Quantization Model

Each domain $i$ has skills indexed by $k_i \in [1, \infty)$ with power-law density:

$$p_i(k_i) = (\alpha_i - 1) k_i^{-\alpha_i}, \quad \text{for } \alpha_i > 1$$

The model allocates capacity by choosing coverage thresholds $x_i \geq 1$, solving:

$$\begin{array}{ll} \underset{x}{\min} & \sum_{i=1}^{K} h_i c_i x_i^{-b_i} \\ \text{s.t.} & \sum_{i=1}^{K} (x_i - H) \leq N - H, \\ & x_i \geq H, \quad \forall i. \end{array} \tag{1}$$

where $b_i = \alpha_i - 1$. The test loss on domain $i$ is:

$$L_i(h) = c_i (x_i^*(h))^{-b_i} + E_i$$

Key insight—Capacity Competition: The constraint $\sum(x_i - H) \leq N - H$ couples domains, creating non-trivial interaction. For similar exponents, an approximate closed form exists:

$$x_i^*(h) \approx \frac{(b_i c_i h_i N^{\bar{b}+1})^{\frac{1}{b_i+1}}}{\left(\sum_{k=1}^{K} (b_k c_k h_k)^{\frac{1}{b_k+1}}\right)^{\frac{\bar{b}+1}{b_i+1}}}$$

Limitation: This model alone predicts $h^* \equiv w$ (train on target distribution), contradicting empirical observations.

Extended Linear Regression Model

To address this limitation, the model incorporates data-dependent noise. Under the spectral assumption of shared head/disjoint tails, each domain's tail eigenvalues in the mixture are:

$$\lambda(\mathbf{H}(h), u_k^{(i)}) = h_i k^{-\alpha_i}$$

Theorem 4.1 (Informal): For a projected linear model trained via one-pass SGD on $D$ samples from mixture $h$, the expected test loss on domain $i$ satisfies:

$$L_i(h \mid N, D) \approx c_i x_i^*(h, N)^{-b_i} + A_i (D h_i)^{-a_i} + E_i \tag{3}$$

where $a_i, A_i, E_i$ are constants depending on $\alpha_i$.

Key insight—Noise Reduction: The noise term $A_i(Dh_i)^{-a_i}$ breaks the symmetry of the capacity allocation problem, shifting optimal weights toward harder-to-learn domains (larger $A_i$, smaller $\alpha_i$).

Optimization

The optimal mixture is found via bi-level optimization:

$$h^{*} = \arg\min_{h} \sum_{i=1}^{K} w_i \left(c_i x_i^*(h, N)^{-b_i} + A_i (D h_i)^{-a_i} + E_i\right) \tag{4}$$

Proposition 4.2 (Gradient Characterization): The gradient of the outer objective is:

$$\nabla_k \mathcal{J}(h) = -w_k a_k A_k D^{-a_k} h_k^{-a_k-1} + \frac{\lambda x_k^*}{h_k (b_k + 1)} \left(\bar{R} - \frac{w_k}{h_k}\right)$$

where $\bar{R} = \frac{\sum_{j=1}^{K} \frac{x_j^*}{b_j+1} \left(\frac{w_j}{h_j}\right)}{\sum_{j=1}^{K} \frac{x_j^*}{b_j+1}}$

This is solved using Online Mirror Descent (OMD).

Empirical Validation / Results

Experimental Setup

Fitting Accuracy Results

Table 2: Comparison of fitting accuracy on 64 1B-parameter models (K=17 domains from Pile)

The Extended Linear Regression Model achieves the lowest MRE (1.53%) and MAE (0.034) while using $5K$ parameters versus $K(2K+1)$ for the Additive baseline—a significant reduction.

Optimal Mixture Prediction

The mixture predicted by the Extended Linear Regression Model consistently yields the best held-out performance across all configurations:

• 200M model on 4 domains: Lowest test perplexity
• 122M model on 7 domains: Lowest test perplexity
• 310M model on 7 domains: Lowest test perplexity
• 17-domain Pile setting: Lowest test loss on Pile-CC domain

Cross-Scale Extrapolation Results

4-domain setting (200M/8B → 700M/16B):

• The extrapolated mixture for the target 700M/16B scale strictly outperforms both the static mixture (derived for 200M/8B) and the Additive Law baseline
• Confirms the framework correctly predicts how optimal mixtures shift with scale

7-domain setting (122M/10B → 1B/30B):

• The extrapolated mixture achieves the same test loss as the state-of-the-art empirical baseline (Additive Law fitted with a massive range of scales including the target 1B/30B itself)
• This is highly significant: the baseline had access to target-scale data, while our model relied solely on 122M/10B proxy losses

Robustness to Tail Overlap

A synthetic stress test (Figure 1) shows the model's MRE remains consistently low and stable even as tail overlap increases up to 40%, demonstrating robustness to violations of the disjoint-tail assumption.

Theoretical and Practical Implications

Theoretical Contributions

1. Unified framework: Bridges the gap between empirical data mixing laws and theoretical understanding by extending established single-domain scaling law theories to multi-domain settings.

2. Capacity Competition as mechanism: Formalizes how finite model capacity creates global coupling between domain losses—a non-trivial interaction that purely empirical laws capture only through flexible curve fitting.

3. Noise Reduction as mechanism: Explains why optimal mixtures deviate from target distributions: harder-to-learn domains require more training data to reduce noise, shifting optimal weights away from the target distribution.

4. Parameter efficiency: The theoretically grounded model achieves superior accuracy with $5K$ parameters versus $K(2K+1)$ for the best baseline—a quadratic-to-linear reduction in parameter count.

Practical Implications

1. Compute-efficient mixture selection: Practitioners can determine optimal data mixtures using small-scale proxy models, then extrapolate to larger scales without expensive trial-and-error.

2. Interpretable predictions: Unlike black-box empirical laws, the framework provides mechanistic understanding of why certain mixtures perform better, enabling more informed decisions.

3. Scale-aware optimization: The framework explicitly captures how optimal mixtures shift with model size and data budget, enabling scale-specific mixture design.

Conclusion

Main Takeaways

• The "Shared Head, Disjoint Tail" structure provides a tractable and empirically validated assumption for modeling multi-domain data.
• Two mechanisms—Capacity Competition and Noise Reduction—govern domain losses in data mixing, explaining both domain interaction and optimal mixture deviation from target distributions.
• The Extended Linear Regression Model achieves state-of-the-art fitting accuracy (MRE = 1.53%) and optimal mixture prediction with significantly fewer parameters than empirical baselines.
• The framework successfully extrapolates across scales, matching the performance of empirical laws fitted with target-scale data while using only small-scale proxy data.

Future Directions

1. Unseen domains and downstream tasks: Extending the framework to predict effects on unseen domains or downstream task performance.

2. Explicit domain overlap: Incorporating explicit modeling of information overlap to refine predictions for highly correlated domains.

3. Reliable fitting algorithms: Developing more robust estimation techniques (e.g., convex relaxations or analytical approximations) to replace the current non-convex optimization, which is computationally intensive and initialization-sensitive.