# Harnessing the Potential of Optimizing Data Mixtures via [ByDoRe]

> ByDoRe's Bayesian domain reweighting achieves state-of-the-art LLM pretraining performance with only 0.8% of RegMix's compute, without relying on fragile scaling-law assumptions.

- **Source:** [arXiv](https://arxiv.org/abs/2607.27928)
- **Published:** 2026-09-19
- **Permalink:** https://picx.dev/p/1C0Zbp
- **Whiteboard:** https://picx.dev/p/1C0Zbp/image

## Summary

## 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)**

| Task | Human | Pile-CC | DoReMi | RegMix | AutoScale | MDE | **Ours** |
|------|-------|---------|--------|--------|-----------|-----|---------|
| ARC-Easy | 49.14 | 49.98 | 51.05 | 50.28 | 46.94 | 47.60 | **51.32** |
| COPA | 66.17 | 67.17 | 67.17 | 69.00 | 68.00 | 63.33 | **70.17** |
| HellaSwag | 37.93 | 42.26 | 41.88 | **43.65** | 37.99 | 37.63 | 41.58 |
| LAMBADA | 28.06 | 35.99 | 34.12 | **34.35** | 27.56 | 26.30 | 31.49 |
| LogiQA | 24.37 | 26.47 | 27.27 | 27.21 | 26.22 | 26.70 | **28.11** |
| MultiRC | 55.08 | 52.53 | 52.76 | 52.94 | 53.00 | 50.30 | **53.43** |
| OpenBookQA | 28.30 | 30.37 | **30.13** | 29.33 | 27.67 | 29.00 | 29.07 |
| PIQA | 64.91 | 68.91 | 68.06 | **69.12** | 65.29 | 66.23 | 68.03 |
| QQP | 38.99 | 43.77 | 46.92 | **50.19** | 49.41 | 46.24 | 49.08 |
| RACE | 30.88 | 31.24 | 30.69 | 31.18 | 29.44 | 30.40 | **31.33** |
| SciQ | 80.92 | 80.57 | **80.83** | 78.35 | 77.17 | 77.40 | 80.93 |
| Social IQA | 38.79 | 40.11 | **40.56** | 40.02 | 38.72 | 38.05 | 39.63 |
| WinoGrande | 51.16 | 52.38 | 51.80 | 51.24 | **52.58** | 51.87 | 51.68 |
| **Overall Avg** | 45.75 | 47.83 | 47.94 | 48.22 | 46.15 | 45.47 | **48.50** |
| **FLOPs cost** | - | - | $2.10\times10^{18}$ | $3.07\times10^{18}$ | $2.55\times10^{17}$ | $1.02\times10^{17}$ | **$2.40\times10^{16}$** |

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)**

| Task | Human | DoReMi | RegMix | AutoScale | MDE | **Ours** |
|------|-------|--------|--------|-----------|-----|---------|
| ARC-Easy | 30.74 | 31.33 | 27.38 | 29.69 | 28.79 | **32.46** |
| COPA | **56.50** | 53.67 | 53.00 | 50.17 | 57.17 | 53.50 |
| HellaSwag | 27.46 | 28.68 | 25.77 | 26.95 | 26.91 | **28.73** |
| LAMBADA | 5.50 | **9.07** | 0.07 | 3.12 | 2.56 | 8.30 |
| LogiQA | 26.09 | 26.75 | 24.01 | 24.32 | 23.40 | **28.14** |
| MultiRC | 56.20 | 54.37 | 56.51 | 51.10 | 55.48 | **56.87** |
| OpenBookQA | 23.53 | **25.07** | 24.87 | 24.67 | 24.50 | 24.40 |
| PIQA | 54.23 | **55.62** | 51.91 | 53.92 | 53.19 | 54.61 |
| QQP | **46.56** | 38.80 | 37.30 | 37.33 | 37.76 | 39.00 |
| RACE | 22.98 | 24.91 | 21.48 | 23.98 | 23.41 | **24.58** |
| SciQ | 56.27 | **64.58** | 25.07 | 51.75 | 48.72 | 62.77 |
| Social IQA | 34.37 | 34.66 | 33.30 | 33.19 | 33.53 | **35.39** |
| WinoGrande | 49.93 | 48.75 | **50.26** | 48.76 | 48.88 | 49.76 |
| **Overall Avg** | 37.72 | 38.18 | 33.15 | 35.45 | 35.72 | **38.35** |
| **FLOPs cost** | - | $1.68\times10^{17}$ | $1.22\times10^{18}$ | $1.02\times10^{17}$ | $4.08\times10^{16}$ | **$1.44\times10^{16}$** |

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

---

_Markdown view of https://picx.dev/p/1C0Zbp, served by PicX — AI-generated visual whiteboard summaries of research papers._