Data Mixing › Data Mixing Experiments

Determining the optimal mixing recipe empirically requires a principled experimental methodology. You cannot afford to train full-scale models for every combination you want to test. Even at the scale of a few hundred experiments, the compute cost of evaluating each configuration with a full training run would exceed the budget of most projects. Instead, the field has converged on a workflow centered around proxy evaluations at reduced scale.

The scale of the search space makes careful experimental design essential. If you have 10 domain categories and want to test mixing proportions at 5% increments, you have combinatorially many candidate mixtures. Even with proxy models, evaluating more than a few dozen configurations is expensive. Researchers have developed several strategies for searching this space efficiently, from hand-crafted ablation designs to automated search algorithms.

The main insight allowing practical data mixing research is that the relative ranking of data mixtures tends to be consistent across model scales. A mixture that produces a better model at 1 billion parameters also tends to produce a better model at 70 billion parameters, even if the absolute performance numbers differ. This transfer property is what makes proxy experiments useful: you can identify good mixtures cheaply at small scale and trust that they will remain competitive at full scale.

This insight enables a much cheaper experimental workflow. You train small proxy models (typically 100 million to 1 billion parameters) on candidate data mixtures, then evaluate them on a benchmark suite covering the capabilities you care about, then rank the mixtures by proxy model performance, and finally use the top-ranked mixture for the full-scale training run. This workflow reduces the cost of each experiment by orders of magnitude compared to full-scale evaluation.

The evaluation suite should cover the capabilities you care about. If your target application is coding, your benchmarks should include coding tasks like HumanEval or MBPP. If you care about factual accuracy, include knowledge-intensive QA datasets like Natural Questions or TriviaQA. If reasoning matters, include mathematical benchmarks like GSM8K or MATH. Using a diverse benchmark suite reduces the risk of overfitting your mixture to a narrow set of tasks. A mixture that maximizes one benchmark while hurting others may not be the one you want.

The proxy model approach has important limits. The scale transfer assumption holds on average across many experiments, but individual exceptions exist. Some capabilities, particularly those associated with emergent abilities that appear suddenly at certain scale thresholds, simply do not manifest at proxy scale. At 100 million parameters, no mixing recipe triggers strong chain-of-thought reasoning or sophisticated multi-step problem solving. This makes it impossible to optimize for these capabilities through proxy experimentation alone. For emergent capabilities, you are essentially flying blind on mixing decisions and must rely on intuitions carried over from other design choices.

Even proxy model training has a cost. To test many mixture configurations efficiently, researchers typically run very short training runs, sometimes just 1 to 5% of the full compute budget, and use the resulting loss curves as an early signal of mixture quality.

The loss at intermediate training steps reflects how quickly the model is learning from each domain. A mixture where all domains improve steadily typically indicates a well-balanced recipe. A mixture where some domains stagnate while others improve rapidly may indicate imbalance: the model is spending too little time on the stagnating domains to make progress there. Short training runs can surface these differences early enough to be actionable without committing to the full compute cost of a complete proxy run.

The risk with extremely short runs is that some data mixing effects only manifest after extended training. A source that is helpful in the long run might not show its benefit within the first few billion tokens, either because the benefit compounds over many exposures or because the capability in question requires a complex integration of skills that only develops after sustained training. This is one reason practitioners supplement short-horizon ablations with a small number of longer validation runs to check that short-horizon rankings hold, particularly for their top few candidate recipes.

Another practical technique for sample-efficient exploration is variance reduction through ensembling. Instead of running a single proxy training run for each candidate recipe, running 3 to 5 independent runs with different random seeds and averaging their performance gives you lower-variance estimates of each mixture's quality. This allows you to distinguish better recipes from ones that just happened to get lucky in one run. The extra cost is usually worth it when you are comparing recipes that perform similarly in your initial screening.

DoReMi (Domain Reweighting with Minimax Optimization) is a more principled algorithm for data mixing that frames the problem as a distributionally robust optimization. Rather than manually testing mixing ratios and relying on human judgment to interpret the results, DoReMi uses a two-model approach that automatically finds mixing weights optimized against a formal objective.

The algorithm's goal is to minimize the worst-case excess loss across all domains. The intuition is similar to how a student should allocate study time: spend more effort on weak subjects, less on subjects already mastered. Domains where the proxy model performs poorly relative to the reference model receive higher weights, and domains where it already performs well receive lower weights. This continues iteratively until the proxy model's performance is as good as possible across all domains simultaneously, rather than being excellent on easy domains while struggling on hard ones.

The algorithm begins by training a reference model with uniform domain weights for some number of steps. This reference model establishes a baseline performance level for each domain under the null hypothesis of equal domain importance. Then a proxy model is trained with adaptive weights that are updated based on where the proxy is falling behind the reference.

The weight update at each step works by computing an excess loss signal for each domain. For domain $d$, the excess loss measures how much worse the proxy model is compared to the reference:

$$\Delta_d = \ell_d^{\text{proxy}} - \ell_d^{\text{ref}}$$

where:

• $\ell_d^{\text{proxy}}$: the average token-level cross-entropy loss of the proxy model on held-out data from domain $d$
• $\ell_d^{\text{ref}}$: the average token-level cross-entropy loss of the reference model on the same held-out data

A positive $\Delta_d$ means the proxy is struggling on domain $d$ relative to the reference. A negative value means the proxy has caught up or surpassed the reference on that domain. When the proxy has been trained with too little weight on a domain, its loss on that domain will remain high relative to the reference, generating a large positive $\Delta_d$ and triggering a weight increase.

The weight update then uses an exponentiated gradient step to shift weight toward domains with higher excess loss:

$$w_d^{(t+1)} \propto w_d^{(t)} \cdot \exp\left(\eta \cdot \Delta_d\right)$$

where:

• $w_d^{(t)}$: the weight assigned to domain $d$ at iteration $t$
• $\eta$: the step size controlling how aggressively weights shift based on the excess loss signal
• $\exp(\cdot)$: the exponential function, which ensures weights remain positive throughout the update

After the update, weights are renormalized to sum to 1 by dividing each by the total:

$$w_d^{(t+1)} = \frac{w_d^{(t)} \cdot \exp(\eta \cdot \Delta_d)}{\sum_{d'} w_{d'}^{(t)} \cdot \exp(\eta \cdot \Delta_{d'})}$$

This normalization step projects the updated weights back onto the probability simplex, the set of non-negative vectors that sum to 1, making sure the mixing recipe remains a valid probability distribution after each update step. Without this normalization, repeated exponentiation would cause weights to grow without bound.

The exponentiated gradient update has a nice property: it never sets any weight to exactly zero, so no domain is ever completely excluded. Domains that are consistently easy for the proxy model receive very small but nonzero weights, making sure the model retains at least some exposure to those domains throughout training. This softmax-like behavior is more desirable than an update rule that could collapse weights to zero, because it maintains the diversity of the training distribution even as it concentrates weight on the hardest domains.

The exponentiated gradient family of updates is well-studied in the online learning literature under the name "multiplicative weights" or "Hedge." DoReMi inherits theoretical guarantees about convergence toward the minimax-optimal weights under appropriate conditions. Specifically, the regret of the algorithm (how much worse it performs compared to the optimal static weight vector) grows as $O(\sqrt{T \log D})$ where $T$ is the number of update steps and $D$ is the number of domains. This guarantee holds regardless of how the domain losses evolve, making DoReMi robust to adversarial or unexpected loss dynamics.

The published DoReMi results showed meaningful improvements over manually designed mixing recipes on several benchmarks, particularly for domains that were systematically under-represented in human-designed recipes. The automated nature of the approach is also operationally appealing: it reduces the number of manual decisions you need to make and makes the mixing recipe more reproducible across different practitioners and projects.

Subsequent work extended the DoReMi framework and combined it with techniques from hyperparameter search. DoReMi-2 incorporates successive halving: rather than evaluating all domain candidates throughout training, the approach progressively eliminates the worst-performing configurations and reallocates compute to more promising ones. This is analogous to Hyperband or similar adaptive hyperparameter search methods. By discarding clearly inferior configurations early, successive halving concentrates the experimental budget on the part of the search space where the best answer is likely to live.

These techniques collectively shift data mixing from a manual art toward an automated optimization problem. You specify your domains, your evaluation benchmarks, and your compute budget, and the algorithm finds a near-optimal mixing recipe without requiring dozens of hand-designed experiments. The automation is not perfect: the proxy-to-full-scale transfer assumption still needs to hold, and the benchmarks you choose still encode your implicit preferences. But the process is substantially more systematic and less reliant on individual practitioner intuition than earlier approaches.