🤖 machine learning

Harnessing the Potential of Optimizing Data Mixtures via Bayesian Domain Reweighting

This paper introduces a Bayesian domain reweighting method that leverages Gamma priors and Dirichlet distributions to infer optimal pre-training data mixtures, achieving stable and efficient optimization with significantly lower data consumption compared to existing function-fitting or direct search approaches.

Original authors: Xiang Yuan, Kaiqing Lei, Zhenyu Jin, Jun Shu, Deyu Meng, Zongben Xu

Published 2026-07-31

📖 7 min read🧠 Deep dive

CC BY 4.0

Original authors: Xiang Yuan, Kaiqing Lei, Zhenyu Jin, Jun Shu, Deyu Meng, Zongben Xu

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a super-smart robot how to speak like a human. You have a massive library of books, websites, code, and emails to feed it. But here's the tricky part: not all books are created equal. If you feed the robot 99% cooking recipes and 1% math textbooks, it will become a chef who can't do algebra. If you feed it mostly old emails, it might learn to sound like a corporate robot. This is the world of Large Language Models (LLMs), the AI brains behind tools like chatbots and writing assistants. The secret sauce to making them brilliant isn't just how much data they eat, but what mix of data they consume.

For a long time, humans tried to guess the perfect recipe by hand, mixing in a little bit of code, a dash of news, and a sprinkle of Wikipedia. But as the libraries of data grew into the billions, human guessing stopped working. The relationships between different types of data became too complex. Some researchers tried to solve this by building a "crystal ball" that predicts how much of each data type to use based on how well the robot performs on small tests. Others tried to just let the robot taste everything and adjust the mix on the fly. Unfortunately, the crystal balls often broke when the data got too big, and the "taste-and-adjust" method was so chaotic and slow that it was practically impossible to finish the job.

This paper introduces a new way to find the perfect data recipe, called ByDoRe. Instead of guessing or trying to predict the future with rigid rules, the authors treat the data mix like a game of probability. They use a clever mathematical trick called Bayesian inference, which is like having a smart assistant that doesn't just pick one answer, but keeps a whole range of "maybe" answers in mind, updating its confidence as it learns. By using this method, the team found a way to stabilize the learning process, making it faster and more reliable than previous attempts. Their experiments suggest that this approach can find a better data mix while using a tiny fraction of the computer power required by other methods, even when the data behaves in unpredictable ways.

The Problem: The Robot's Picky Eater

Imagine you are a chef trying to create the ultimate smoothie. You have a blender, a mountain of fruit, and a recipe book. If you just throw everything in at random, the smoothie might taste terrible. If you follow a recipe written by a human chef from 10 years ago, it might be okay, but it won't be perfect for today's taste buds.

In the world of AI, the "smoothie" is the training data, and the "chef" is the algorithm that decides how much of each ingredient (like code, news, or stories) to put in. Early AI models relied on human chefs who guessed the ratios. But as the amount of data exploded, these guesses became useless. The AI needed a smarter way to figure out the mix.

Some scientists tried to build a "predictor" that looks at a small taste of the smoothie and guesses the perfect recipe for the whole batch. They assumed that if a certain fruit (say, Wikipedia) tastes good in a small cup, it will be the most important ingredient in the giant pitcher. But the paper shows this is often wrong. Just because something tastes good in a small sample doesn't mean it's the best choice for the whole thing. In fact, the paper found that as the data gets bigger, the "importance" of different ingredients can flip-flop. What was the best ingredient at a small scale might become the worst at a large scale. This is like thinking a tiny bit of salt makes a soup perfect, only to realize that in a giant pot, it makes it inedible.

Other scientists tried a different approach: they let the AI taste the data and adjust the mix in real-time. This is like the chef tasting the smoothie every second and adding more sugar or fruit. The problem? The chef gets dizzy. The adjustments are so wild and jittery that the smoothie never settles into a good flavor. It takes forever to get it right, and the computer gets so tired (expensive to run) that it's not worth it.

The Solution: The Bayesian "Smart Guess"

The authors of this paper, Xiang Yuan and his team, decided to stop guessing and stop trying to force the data into a rigid box. Instead, they built a system called ByDoRe (Bayesian Domain Reweighting).

Think of ByDoRe not as a chef who picks one specific recipe, but as a chef who keeps a mental map of all possible good recipes. Instead of saying, "I am 100% sure we need 30% Wikipedia," the system says, "I think we need about 30% Wikipedia, but it could be anywhere between 25% and 35% depending on how the robot is feeling."

This "mental map" is based on a mathematical concept called a Dirichlet distribution. Imagine a bag of marbles where the colors represent different data types. Instead of pulling out one marble and saying, "This is the color we need," the system looks at the whole bag and understands the probability of each color being the right one. This helps smooth out the wild swings and jittery adjustments that made the previous methods fail.

To make this even faster, they added a "smart predictor" (a neural network) that learns from the robot's current performance to guess what the "mental map" should look like next. It's like having a sous-chef who watches the robot taste the smoothie and immediately suggests, "Hey, the robot liked the code part, let's nudge the probability of code up a little bit."

What They Found: A Smoother, Faster Ride

The team tested their new method on a massive dataset called The Pile, which contains 17 different types of data, from scientific papers to chat logs. They compared ByDoRe against the best existing methods, including the jittery "taste-and-adjust" method (DoReMi) and the "predictor" methods (RegMix, AutoScale).

Here is what they discovered:

It's Stable: While the old "taste-and-adjust" method was like a rollercoaster, with the data mix jumping up and down wildly, ByDoRe was a smooth train ride. The weights (the recipe ratios) settled down quickly and stayed there.
It's Cheaper: This is the big one. To find the best mix, the old methods required a massive amount of computer power. ByDoRe managed to find a better mix using only 0.8% of the computer power (measured in FLOPs) that the top competitor, RegMix, needed. In Case 2, where the data was tricky, ByDoRe used only 1.2% of the cost of the best competitor.
It Works When Others Fail: The paper tested a scenario where the data behaved badly (violating the "rules" that other methods rely on). The old methods crashed and gave terrible results. ByDoRe, however, adapted and found a great mix, even beating human experts in some tests.

In one specific test, ByDoRe achieved an average score of 48.50% on various language tasks, beating the previous best automated method (RegMix) which scored 48.22%. Even more impressively, in a specialized test where the data was very different from the norm, ByDoRe scored 38.35%, while the other automated methods struggled to reach 35% or lower.

Why This Matters

The paper suggests that we don't need to rely on rigid rules or expensive, chaotic guessing to teach AI how to eat the right data. By using a probabilistic approach that embraces uncertainty rather than fighting it, we can find better data mixes much faster and with less computing power.

The authors conclude that this method offers a "stable and agile framework" for the future. It doesn't just work for one type of data; it seems robust enough to handle the messy, unpredictable reality of real-world information. While the paper doesn't claim to have solved every problem in AI, it suggests a promising new direction: treating data selection not as a math problem with one right answer, but as a dynamic, probabilistic journey where the AI learns to trust its own evolving intuition.

Technical Summary: Harnessing the Potential of Optimizing Data Mixtures via Bayesian Domain Reweighting

Problem Statement

The performance of Large Language Models (LLMs) is fundamentally determined by the distributional composition of multi-domain pre-training data. While early models relied on manual heuristics to determine data mixture ratios, these static configurations fail to capture complex synergies between domains as data scales increase.

Existing automated approaches attempt to optimize domain weights but face two primary limitations:

Function-Fitting Methods (e.g., RegMix, AutoScale): These methods fit a proxy function (assuming rank invariance or scaling laws) between domain weights and validation losses on small data subsets to extrapolate to full-scale training. However, empirical evidence shows that domain importance shifts non-monotonically and violates rank invariance as data scales grow, leading to significant estimation bias and performance degradation in specialized or complex scenarios.
Direct Optimization Methods (e.g., DoReMi): These methods optimize weights directly from data using distributionally robust optimization. While assumption-free, they suffer from unstable optimization trajectories due to high-variance stochastic loss signals in one-pass LLM training. This results in violent oscillations, slow convergence, and prohibitive computational costs, often requiring exhaustive traversal of the corpus to achieve stability.

Methodology: ByDoRe Framework

The authors propose ByDoRe (Bayesian Domain Reweighting), a framework that reformulates domain weighting as a probabilistic inference problem rather than a deterministic optimization task.

Core Components

Hierarchical Bayesian Model:
- Instead of learning a deterministic weight vector $w$ , ByDoRe infers a probability distribution over domain mixtures.
- The domain weights $w$ are modeled as latent random variables following a Dirichlet distribution, governed by concentration parameters $\beta$ .
- The parameters $\beta$ themselves follow Gamma distributions, whose shape and rate hyperparameters are determined by observed domain-specific signals (e.g., batch-level losses).
Prior Prediction Network:
- Inspired by empirical Bayes, the authors introduce a neural network $G_\Phi$ (a prior prediction network) that adaptively determines the hyperparameters of the Gamma distribution based on empirical domain signals $s_t$ (average losses from small batches).
- This allows the model to learn "prior-assignment laws" that capture long-term data patterns while filtering out transient batch-level variance.
Variational Inference and Optimization:
- The goal is to maximize the Evidence Lower Bound (ELBO) of the posterior predictive distribution on a validation set.
- The optimization involves alternating updates:
  - Model Update: The proxy model parameters $\theta$ are updated using a training batch constructed via the sampled weights $w_t \sim q_\Phi(w)$ .
  - Prior Network Update: The parameters $\Phi$ of the prior prediction network are updated to maximize the ELBO.
- Gradient Approximation: To avoid the prohibitive cost of computing Hessian-vector products required for exact meta-gradients, the authors employ a central finite-difference approximation. This allows the meta-gradient to be computed using only first-order gradients of the training loss.
Physical Sampling Intervention:
- The framework intervenes directly in the physical data sampling process. For every training step, the domain proportions within the batch are strictly maintained according to the sampled weights $w_t$ , ensuring the model receives high-quality, representative gradients.

Key Contributions

Stabilized Optimization: ByDoRe addresses the instability and slow convergence of direct optimization methods (like DoReMi) by introducing a probabilistic layer. The hierarchical Bayesian formulation smooths the optimization landscape, eliminating the negative effects of deterministic weight learning and reducing sensitivity to stochastic batch-level fluctuations.
Assumption-Free Efficiency: Unlike function-fitting methods that rely on violated structural assumptions (rank invariance, scaling laws), ByDoRe learns directly from data observations without restrictive priors. It achieves this with significantly lower computational overhead than both function-fitting and direct optimization baselines.
Superior Performance with Low Cost: The method achieves state-of-the-art (SOTA) results in general-purpose and specialized targeting scenarios while consuming a fraction of the training cost required by existing methods.

Experimental Results

The authors evaluated ByDoRe on The Pile dataset (17 domains) across two scenarios:

Case 1: General-Purpose Language Modeling:
- ByDoRe achieved an overall average zero-shot score of 48.50%, outperforming RegMix (48.22%) and AutoScale (46.15%).
- Crucially, ByDoRe correctly identified Wikipedia as a primary knowledge source (weight 0.31), whereas automated baselines overfit to the validation distribution (e.g., RegMix assigned 0.87 to Pile-CC and neglected Wikipedia).
- Efficiency: ByDoRe required only 2.40 × 10¹⁶ FLOPs for the search phase, representing a 127x reduction in search overhead compared to RegMix and an 87x reduction compared to DoReMi.
Case 2: Specialized Targeting (GitHub and Enron Emails):
- In this scenario, standard scaling law assumptions were empirically violated (e.g., Enron Emails showed non-monotonic loss trajectories and rank reversals).
- Function-fitting baselines (RegMix, AutoScale, MDE) suffered significant performance collapse.
- ByDoRe achieved the highest average score (38.35%), surpassing both human heuristics and DoReMi, demonstrating robustness where rigid assumptions fail.
- It maintained peak efficiency with the lowest overhead among all automated baselines (1.44 × 10¹⁶ FLOPs).

Significance and Claims

The paper claims that ByDoRe revitalizes optimization-based domain weighting for large-scale applications by resolving the trade-off between stability and computational cost.

Mechanism: The authors assert that the key to success lies in the Gamma-Dirichlet hierarchical structure, which internalizes stable importance patterns while marginalizing out batch-level stochasticity.
Practicality: The method is presented as a practical solution for real-world LLM training, capable of adapting to domain weight oscillations in online one-pass learning without access to the full-scale corpus during the search phase.
Impact: The study suggests a promising direction for data optimization in the LLM era, specifically for data mixing and selection across all stages of training, offering a stable and agile framework that outperforms both manual tuning and existing automated strategies in both general and specialized contexts.