Summary (Overview)

• This paper introduces ByDoRe (Bayesian Domain Reweighting), a novel framework for optimizing data mixtures in Large Language Model (LLM) pre-training by formulating domain weights as probabilistic distributions rather than deterministic vectors.
• The method uses a Gamma-Dirichlet hierarchical Bayesian model with a prior prediction network to stabilize and accelerate domain weight learning, addressing the instability and computational prohibitive issues of direct optimization methods like DoReMi.
• ByDoRe achieves state-of-the-art performance on general-purpose language modeling (48.50% average accuracy) while consuming only 0.8% of the training cost of the SOTA RegMix method, and improves specialized target performance by 5.2 percentage points over RegMix with only 1.2% of its training cost.
• The framework provides theoretical convergence guarantees for both the meta-objective and training objective at $\mathcal{O}(1/\sqrt{T})$ rates.
• Experimental results demonstrate that ByDoRe is the only method that correctly identifies Wikipedia as a primary knowledge source, balancing it against Pile-CC, whereas existing methods overfit to validation distributions.

Introduction and Theoretical Foundation

Background and Motivation

The performance of Large Language Models is fundamentally influenced by the distributional composition of multi-domain pre-training data. Early approaches (GPT-3, Llama 2) relied on manual heuristics for data mixing, but these static configurations fail to capture intricate synergies between domains as data complexity grows.

Limitations of Existing Approaches

1. Function-Fitting Methods (RegMix, AutoScale, MDE): These fit a proxy function mapping domain weights to validation losses, relying on structural assumptions such as:

• Rank invariance: Domain importance rankings remain consistent across scales
• Scaling laws: Power-law extrapolation of performance

However, as shown in Figure 1, these assumptions are frequently violated. The paper demonstrates empirically that:

• Domain rankings observed at small scales are reversed at 10x larger data scales
• Power-law fits derived from small-scale data diverge significantly from actual observed losses when extrapolated

2. Direct Optimization Methods (DoReMi): These directly optimize domain weights through distributionally robust optimization:

$$\min_{\theta} \max_{\mathbf{w} \in \Delta^{K-1}} \sum_{i=1}^{K} w_i \left[ \frac{\sum_{x \in D_i} \bigl(\ell_{\theta}(x) - \ell_{\mathrm{ref}}(x)\bigr)}{\sum_{x \in D_i} |x|} \right] \tag{1}$$

But suffer from:

• Unstable optimization trajectories: High-variance oscillations due to one-pass training
• Prohibitive computational overhead: Often requires traversing nearly the entire trillion-token corpus

Theoretical Foundation: Bayesian Formulation

The paper reformulates domain weighting from a probabilistic perspective. Instead of optimizing a deterministic weight vector, ByDoRe infers a probability distribution over domain mixtures. The key insight is that modeling uncertainty reduces sensitivity to stochastic batch-level loss fluctuations.

The problem setup defines:

• Training corpus: $\mathcal{D} = \cup_{k=1}^{K} D_k$ partitioned into K domains
• Domain-weight vector: $\mathbf{w} \in \Delta^{K-1}$ where $w_k \geq 0$ and $\sum_{k=1}^{K} w_k = 1$
• Additive mixture distribution: $p_{\mathbf{w}}(x) = \sum_{k=1}^{K} w_k p(x \mid D_k)$

Methodology

The Gamma-Dirichlet Hierarchical Bayesian Model

The variational distribution is parameterized as:

$$q_{\Phi}(\mathbf{w}) = \int \operatorname{Dir}(\mathbf{w} \mid \boldsymbol{\beta}) \operatorname{Gamma}\left(\boldsymbol{\beta} \mid \mathbf{a}_{\Phi}(\mathbf{s}), \mathbf{b}_{\Phi}(\mathbf{s})\right) d\boldsymbol{\beta} \tag{6}$$

The framework consists of three steps:

(i) Prior-parameter prediction: A prior prediction network $\mathcal{G}_{\Phi}$ generates Gamma hyperparameters from empirical domain signals:

$$\left(\mathbf{a}_{\Phi}(\mathbf{s}), \mathbf{b}_{\Phi}(\mathbf{s})\right) = \mathcal{G}_{\Phi}(\mathbf{s})$$

The network consists of a linear projection, two-layer Transformer encoder, and output layer with Softplus activation. Input signals are domain-specific losses at iteration t:

$$\mathbf{s}_t = \left[ \ell_{\theta_t}(b_{t,1}), \dots, \ell_{\theta_t}(b_{t,K}) \right]$$

(ii) Concentration-variable sampling: $\boldsymbol{\beta} \sim \operatorname{Gamma}(\mathbf{a}_{\Phi}(\mathbf{s}), \mathbf{b}_{\Phi}(\mathbf{s}))$

(iii) Domain-weight sampling: $\mathbf{w} \sim \operatorname{Dir}(\boldsymbol{\beta})$

Learning Algorithm

The training alternates between two updates:

Main Model Update (Eq. 7):

$$B_t \sim p_{\mathbf{w}_t}^{\otimes |B_t|}, \qquad \theta_{t+1} = \theta_t - \eta \nabla_{\theta} \mathcal{L}_{\mathrm{tr}}(\theta_t; B_t)$$

Domain-Weight Distribution Update (Eq. 8-9): Uses a central finite-difference approximation of the meta-gradient:

$$\Phi_{t+1} = \Phi_t - \gamma \Bigg(\frac{\eta}{2\epsilon} \Big[ \nabla_{\Phi} \mathcal{L}_{\mathrm{tr}}(\theta_t + \epsilon\Delta, \bar{\mathbf{w}}_t) - \nabla_{\Phi} \mathcal{L}_{\mathrm{tr}}(\theta_t - \epsilon\Delta, \bar{\mathbf{w}}_t) \Big] + \lambda \nabla_{\Phi} \mathcal{R}_{\mathrm{KL}}(\Phi_t) \Bigg)$$

where $\Delta = \nabla_{\theta} \log p(\mathcal{D}_{\mathrm{val}} \mid \hat{\theta}(\mathbf{w}))$ and the expected weights are computed as:

$$\bar{w}_{t,k} = \frac{a_{t,k}/b_{t,k}}{\sum_{j=1}^{K} a_{t,j}/b_{t,j}} \tag{10}$$

Theoretical Guarantees

Theorem 2.1 (Meta-Objective Convergence): The ELBO satisfies:

$$\min_{0 \leq t \leq T-1} \mathbb{E}\left[ \|\nabla_{\Phi}\mathcal{L}(q_{\Phi_t})\|^2 \right] \leq \mathcal{O}\left(\frac{1}{\sqrt{T}}\right) + \mathcal{O}(\epsilon^4)$$

Theorem 2.2 (Training Objective Convergence): Model parameters satisfy:

$$\min_{0 \leq t \leq T-1} \mathbb{E}\left[ \|\nabla_{\theta}\mathcal{L}_{\mathrm{tr}}(\theta_t, \bar{\mathbf{w}}_t)\|^2 \right] \leq \mathcal{O}\left(\frac{1}{\sqrt{T}}\right) + \mathcal{O}(\epsilon^2)$$

Empirical Validation / Results

Experimental Setup

• Dataset: The Pile (17 public domains)
• Two scenarios:
  • Case 1 (General-Purpose): Full mixture of 17 domains
  • Case 2 (Specialized Targeting): Restricted to GitHub and Enron Emails
• Two-stage evaluation: Search phase (1M-parameter proxy model, 0.4B tokens for ByDoRe) → Verification phase (1B-parameter target model, 150B tokens)
• Evaluation: 13 downstream benchmarks (HellaSwag, PIQA, WinoGrande, RACE, SciQ, etc.)

Case 1: General-Purpose Language Modeling

Table 2: Zero-shot performance comparison (Case 1)

Key findings:

• ByDoRe achieves 48.50% average accuracy, outperforming all baselines
• 127x reduction in search FLOPs compared to RegMix ($2.40\times10^{16}$ vs $3.07\times10^{18}$)
• ByDoRe is the only method that correctly identifies Wikipedia (en) as a primary knowledge source (weight 0.31), balanced against Pile-CC (0.59)
• RegMix and DoReMi overfit by assigning disproportionately high weight to Pile-CC (0.87 and 0.61) while neglecting Wikipedia (0.02 and 0.07)

Case 2: Specialized Targeting and Robustness

Table 3: Zero-shot performance comparison (Case 2)

Key findings:

• Function-fitting methods collapse: RegMix (33.15%), AutoScale (35.45%), MDE (35.72%) are all outperformed by human heuristics (37.72%) because Enron Emails violate rank invariance and scaling law assumptions
• ByDoRe achieves 38.35%, exceeding both human baseline and DoReMi (38.18%)
• Lowest computational overhead: $1.44\times10^{16}$ FLOPs

Ablation Studies

1. Probabilistic Modeling Suppresses Jitter: Without the probabilistic layer, weights exhibit violent oscillations and overfit to Pile-CC (near-unity weight). The Bayesian formulation smooths the optimization landscape.

2. Physical Sampling Intervention: Intra-batch weighted sampling (strictly maintaining domain proportions within each batch) achieves the fastest and lowest validation loss compared to:

• Inter-batch weighted sampling
• Intra-batch loss re-weighting
• Inter-batch loss re-weighting

Theoretical and Practical Implications

Theoretical Contributions

1. Novel Bayesian formulation of domain weighting that treats weights as latent random variables, providing a principled framework for uncertainty quantification in data mixing
2. Theoretical convergence guarantees at $\mathcal{O}(1/\sqrt{T})$ rates for both meta-objective and training objective
3. Bridges the gap between direct optimization methods (unstable but assumption-free) and function-fitting methods (efficient but assumption-dependent)

Practical Implications

1. Computational efficiency: ByDoRe achieves SOTA performance with only 0.8% of RegMix's training cost and is 4x faster than the fastest baseline (MDE)
2. Robustness: Does not rely on fragile assumptions like rank invariance or scaling laws, making it applicable to specialized domains
3. Stability: The Gamma-Dirichlet hierarchy effectively suppresses variance-induced jitters, enabling rapid convergence
4. Scalability: The dynamic feedback loop between weight updates and data sampling enables efficient adaptation without accessing the full-scale corpus

Conclusion

ByDoRe addresses the critical challenge of data mixture optimization for LLM pre-training by:

• Formulating domain weights through a Gamma-Dirichlet hierarchical Bayesian model
• Adaptively inferring posterior distributions from observations via a prior prediction network
• Effectively suppressing optimization jitters and marginalizing batch-level stochasticity

The method achieves superior performance across diverse data regimes with substantially lower computational overhead than existing approaches. The authors suggest this work opens promising directions for data optimization in the LLM era, including data mixing and selection across all stages of LLM training.

Future directions include:

• Extending the framework to other stages of LLM training
• Exploring applications beyond data mixing to data selection
• Scaling to even larger model and data regimes