Training Data Mixing

The mixture is a distribution over domains

Partition the corpus into k domains D_1, …, D_k. A data mixture is a weight vector α = (α_1, …, α_k) that lives on the probability simplex: every α_i ≥ 0 and Σ_i α_i = 1. To build a training batch you first sample a domain i ~ α, then sample an example uniformly from D_i. So α_i is literally the expected fraction of tokens the model will see from domain i.

This two-stage view matters because it decouples how big a domain is from how much you train on it. Let p_i = |D_i| / Σ_j |D_j| be the natural (token-proportional) fraction. Choosing α = p simply mirrors the corpus. But nothing forces that: α is a free knob, a design variable you optimize. Everything in mixture research is a method for choosing a good point on that simplex.

What the weights buy you: the training objective

The pretraining loss under mixture α is a weighted sum of per-domain losses:

L(θ; α) = Σ_i  α_i · L_i(θ)
L_i(θ) = E_{x ~ D_i} [ -log p_θ(x) ]   (avg NLL on domain i)

The gradient inherits the same shape: ∇L = Σ_i α_i ∇L_i. So the mixture weights are exactly the weights on each domain’s gradient contribution. Raising α_math tilts every step toward reducing math loss, at the expense of directions that would have lowered web loss.

The catch is that we do not actually want to minimize this weighted training loss — we want good downstream behavior. The mixture is a surrogate control: the game is choosing α so that minimizing L(θ; α) yields a model that generalizes well across the domains we care about, not just the biggest one.

Token-proportional is a choice, not a default

The tempting baseline is α = p: train on data in the proportion it naturally occurs. It feels neutral, but it is a strong and often bad choice. Real web-scale corpora are wildly imbalanced — common crawl text can be 90%+ of tokens while high-value math is a rounding error.

Under token-proportional sampling the loss is dominated by the majority domain: ∇L ≈ α_web ∇L_web when α_web is near one. The optimizer happily drives web perplexity down and barely moves on code or math, because those gradients are scaled by tiny weights. The model becomes fluent and shallow. This is why frontier recipes deliberately upweight scarce, high-signal domains far above their natural share — a small model has a limited token budget, and spending it to over-learn boilerplate web text is the costliest kind of waste.

Upsampling, downsampling, and effective epochs

Once α_i ≠ p_i, the mechanics are repetition and subsampling. Fix a token budget T. Domain i contributes α_i · T training tokens but only holds |D_i| unique tokens, so the number of times you pass over it is:

effective_epochs_i = (α_i · T) / |D_i|

Upsampling (α_i > p_i) means effective_epochs_i > 1 when the budget is large — the model sees those examples multiple times. Downsampling (α_i < p_i) means you use only a subset each pass. This is the practical cost hiding inside a mixture choice: aggressively upweighting a small domain forces many repeats of the same tokens, and repeated data yields diminishing returns and eventual memorization. A sane mixture keeps scarce domains’ effective epochs in a healthy range rather than rereading the same math problems dozens of times.

Temperature reweighting: one knob for the whole mix

A cheap, popular way to move off token-proportional without solving anything hard is a single temperature exponent τ ∈ (0, 1], borrowed from multilingual training:

α_i = p_i^τ / Σ_j p_j^τ

At τ = 1 you recover token-proportional. As τ → 0 the weights flatten toward uniform, boosting rare domains. Worked example with three domains at p = (0.90, 0.08, 0.02) for web, code, math. At τ = 0.5 take square roots: (0.949, 0.283, 0.141), sum 1.373, so α ≈ (0.69, 0.21, 0.10). Math jumped from 2% to 10% of tokens from one scalar. Temperature is a blunt instrument — it moves every domain along one curve and cannot say ‘more math but not more forum spam.’ That limitation is what motivates learning per-domain weights directly.

DoReMi: let a proxy model choose the weights

DoReMi (Domain Reweighting with Minimax Optimization) replaces hand-tuning with a small optimization that outputs the weight vector. It runs three cheap stages, all at small (proxy) scale before you commit the expensive big run:

(1) Train a reference model on some baseline mixture (e.g. token-proportional). (2) Train a proxy model with group distributionally-robust optimization, which simultaneously updates the model and a domain-weight vector α to minimize worst-case excess loss versus the reference. (3) Take the time-averaged α the proxy converged to and use it as the mixture for the real, large model.

The elegance is that no downstream benchmark is in the loop — DoReMi reweights using only the models’ own losses. It asks a purely internal question: which domains is the current model still far from mastering relative to a fair reference, and it pours weight there.

Excess loss: the signal DoReMi optimizes

The core quantity is excess loss — how much worse the proxy is doing on a domain than the reference already did:

λ_i = max( 0 ,  L_i(θ_proxy) - L_i(θ_ref) )

The clip at zero matters: once the proxy matches or beats the reference on a domain, its excess loss is zero and it stops pulling weight. That stops the method from dumping the budget into an intrinsically easy, low-entropy domain where low absolute loss is no signal to train more. By subtracting the reference, DoReMi measures headroom, not raw difficulty.

The minimax objective is then min_θ max_α Σ_i α_i λ_i over the simplex: the model tries to shrink excess loss everywhere, while α adversarially concentrates on whichever domains still have the most headroom — group DRO applied to domains.

The weight update: exponentiated gradient on the simplex

Because α must stay on the simplex, DoReMi updates it multiplicatively with an exponentiated-gradient (mirror-descent) step rather than plain additive gradient ascent:

α_i(t)  ∝  α_i(t-1) · exp( η · λ_i(t) )
then renormalize so  Σ_i α_i(t) = 1

A domain with high excess loss gets its weight multiplied up; a domain with zero excess loss is multiplied by exp(0) = 1 and only drifts down through renormalization. The learning rate η controls how sharply weight chases headroom. Crucially, the final mixture is the average of α(t) over all steps, not the last value — averaging turns the noisy per-step adversary into a stable distribution, the standard trick for extracting an equilibrium from mirror descent.

A worked reweighting step

Three domains, start from uniform α = (1/3, 1/3, 1/3) for web, code, math. Suppose the current excess losses are λ = (0.05, 0.20, 0.30) — the model is close to the reference on web but has real headroom on math. With η = 1:

unnormalized:  (1/3)·e^0.05, (1/3)·e^0.20, (1/3)·e^0.30
            =  0.350,          0.407,          0.450
sum = 1.207
α(t)      =  0.290,          0.337,          0.373

In one step web fell from 0.333 to 0.290 while math rose to 0.373. The update is gentle and proportional to headroom — nothing collapses to zero, but mass steadily migrates toward the domains the model has not caught up on. Iterate alongside the proxy’s training and average to get the learned mixture.