3.14 Data Mixing, Domain Weighting & Curriculum¶

Two teams are handed the same compute budget, the same architecture, and access to the same pile of raw tokens: a few trillion tokens of web text, a few hundred billion of code, tens of billions of math, some books, and a multilingual long tail. Team A throws everything into one shuffled stream in proportion to how much of each they happen to have on disk. Team B spends a week deciding what fraction of each batch should be web, code, math, books, and other languages — and then schedules those fractions to change as training proceeds. At the end, Team B’s model is meaningfully better at coding and reasoning, no worse at general language, and used the exact same number of FLOPs.

That gap is the subject of this chapter. Once you have cleaned and deduplicated your corpus (see Data Cleaning, Deduplication & Quality Filtering) and decided how big a model your budget supports (see Scaling Laws: Kaplan, Chinchilla & Beyond), you still face a deceptively deep optimization: the data mixture. What proportion of each domain do you sample? Do you upsample a small high-quality source or just deduplicate harder? Do you keep the mixture fixed, or change it over training — easy-to-hard curricula, context-window ramps, a high-quality “annealing” phase near the end? These choices routinely move benchmark scores by amounts that would otherwise cost you a 2x model-size increase.

We will develop the topic in four movements: (1) the mixing problem itself — what a mixture is, why it matters, and the upsampling-vs-dedup trade-off; (2) how to choose weights — manual ablations, and the proxy-model methods (DoReMi/Group-DRO, RegMix) that learn weights automatically; (3) a from-scratch DoReMi-style reweighting toy you can run, plus how a mixture is actually realized in a dataloader (and which libraries do it for you); and (4) scheduling over time — curriculum, context-length ramps, and mid-training annealing. The data pipeline that produces these domains is covered in Pretraining Data: Sources, Crawling & The Data Pipeline; here we assume the domains exist and ask how to blend them.

The Mixing Problem: What a Mixture Is and Why It Matters¶

Domains, weights, and the sampling distribution¶

Partition your corpus into \(k\) domains (also called groups or sources) \(\mathcal{D}_1, \dots, \mathcal{D}_k\) — for example web, code, math, books, encyclopedic, and a multilingual bucket. Each domain \(i\) has a natural size \(n_i\) (its token count after cleaning). A mixture is a probability vector

where \(w_i\) is the probability that the next token drawn for a training batch comes from domain \(i\). Equivalently, over a run of \(D\) total tokens, domain \(i\) contributes \(w_i D\) tokens of gradient signal. The mixture is a knob entirely separate from how much data you have: you can set \(w_{\text{code}} = 0.20\) even if code is only 5% of your corpus, by upsampling (revisiting code tokens multiple times), or set \(w_{\text{web}} = 0.40\) even though web is 80% of your corpus, by downsampling (skipping most web tokens).

Define the natural or proportional mixture as \(w_i^{\text{nat}} = n_i / \sum_j n_j\) — sample each token uniformly from the pool. Team A above used \(w^{\text{nat}}\). The central empirical fact of this chapter is that \(w^{\text{nat}}\) is almost never optimal. Why? Because the loss you actually care about is not “loss on the corpus you happen to have”; it is loss on a target distribution of downstream uses, which weights code, math, and reasoning far more heavily than their raw token counts.

The effective-epochs view¶

The cleanest way to think about a mixture is in terms of epochs per domain. If domain \(i\) has \(n_i\) unique tokens and you allocate it \(w_i D\) training tokens, then domain \(i\) is seen for

epochs. A small high-quality domain (say 20 B tokens of curated math) inside a 2 T-token run with \(w_{\text{math}} = 0.05\) is seen \(e_{\text{math}} = (0.05 \times 2{,}000\text{B}) / 20\text{B} = 5\) times. Meanwhile web, with 4 T unique tokens and \(w_{\text{web}} = 0.6\), is seen \(e_{\text{web}} = (0.6 \times 2{,}000) / 4{,}000 = 0.3\) epochs — less than once. This asymmetry is the crux of every mixing decision: upweighting a small domain forces you to repeat it, and repetition has a cost. Past roughly 4 epochs of a domain, returns to repeated data decay sharply, and past a dozen or so, repeated data can actively hurt (memorization, reduced generalization). The Muennighoff et al. “Scaling Data-Constrained Language Models” work quantified this: repeated tokens are worth progressively less than fresh ones, with the value of an epoch decaying roughly geometrically.

So a mixture choice is implicitly an epoch-budget choice. You are deciding, for each domain, “how many times is it worth re-reading this?”

Upsampling vs. deduplication: the trade-off¶

A subtle but important interaction: upsampling and deduplication push in opposite directions on the same axis — how many times a token is effectively seen.

• Deduplication (covered in Data Cleaning, Deduplication & Quality Filtering) removes near-duplicate documents so the model does not waste capacity memorizing boilerplate and does not silently see the same passage 50 times. It reduces effective epochs on duplicated content.
• Upsampling deliberately increases effective epochs on a chosen domain.

The trap is doing both blindly: if you under-deduplicate web and then also upsample a small domain, you can end up with a model that has seen the upsampled domain 5x (intended) but also seen the most common web boilerplate 50x (unintended). The principled recipe is: deduplicate aggressively first, to make “one epoch” mean one genuine pass over unique content; then upsample deliberately, with full visibility into the resulting epoch counts. Deduplication makes the epoch math honest; upsampling spends that honest budget where it helps.

There is also a quality-vs-quantity tension. Suppose you have a small, very high-quality math set. You can:

1. Upsample it (repeat it, raising \(e_{\text{math}}\)) — risks memorization but injects more high-quality gradient signal.
2. Leave it at one epoch and accept its small weight — safe but under-uses a great source.
3. Generate synthetic data in its style to enlarge the unique pool (see Synthetic Data for Pre- and Post-Training) — best of both worlds when feasible, since fresh unique tokens beat repeats.

A common modern practice is a hybrid: upsample high-quality small domains to 2–4 epochs (where repetition is still net-positive), and use synthetic generation rather than pushing past ~4–6 epochs.

Why the mixture is not separable from everything else¶

Two important couplings make mixing harder than “tune five numbers once”:

• Mixture interacts with model scale. The optimal mixture is not scale-invariant. Small models are capacity-limited and benefit from a “cleaner,” narrower diet; large models can absorb a more diverse, heavier-tailed mixture and turn it into capability. A mixture tuned on a 100 M-parameter proxy may be subtly wrong for a 70 B target. This is the central risk of all proxy-based methods, and we return to it below.
• Mixture interacts with the schedule. A fixed mixture is a special case of a schedule. As we will see in the curriculum section, the best results often come from changing \(w\) over time — e.g., more diverse/noisier data early, then concentrating high-quality math/code/instruction-like data in a final annealing phase.

Choosing the Weights I: Manual Ablations and the Target-Loss View¶

The objective: minimize a target-weighted loss¶

Make the goal explicit. You have a target distribution over domains \(p^{\star} = (p_1^{\star}, \dots, p_k^{\star})\) encoding how much you care about each — perhaps uniform over domains (every domain equally important), or skewed toward code/math if that is your product, or matched to a downstream eval suite. Your true objective is to minimize the target-weighted held-out loss:

where \(\ell_i\) is the per-domain held-out cross-entropy (in nats/token; see The Pretraining Objective & Loss). Crucially, the training mixture \(w\) is a separate object from the target weights \(p^{\star}\). You choose \(w\) — the sampling distribution — to minimize \(\mathcal{L}_{\text{target}}\). Setting \(w = p^{\star}\) is the naive guess, and it is usually wrong: domains differ in difficulty and in how much they transfer to one another, so the training weights that minimize a target loss are generally not equal to the target weights themselves. Easy domains need less weight to reach low loss; hard or high-transfer domains may deserve more.

Manual mixture ablations: the workhorse¶

Before any fancy method, the industry standard is mixture ablation at small scale: train many small models (say 100 M–1 B parameters, for a few billion tokens each) under different candidate mixtures, evaluate each on a fixed held-out suite, and pick the winner — then scale it up and hope the ranking holds. The Llama, GPT-3, and Gopher reports all describe variants of this. The procedure:

1. Fix a proxy scale small enough to run dozens of configs cheaply but large enough to be predictive (a few hundred million parameters is typical).
2. Define a small set of candidate mixtures — e.g., a grid or a few hand-designed points: “web-heavy,” “code-heavy,” “balanced,” “math-upsampled.”
3. Train one proxy per mixture for a fixed token budget.
4. Evaluate each on per-domain held-out loss and on a basket of downstream tasks — in practice EleutherAI’s lm-evaluation-harness run over a fixed task list, so every candidate mixture is scored by identical prompts and metrics (see Building Eval Harnesses). At proxy scale, per-domain held-out loss is usually the more sensitive signal: few-shot accuracies on a 300 M model are noisy and often near chance.
5. Pick the mixture optimizing your target metric, then scale, ideally re-checking the ranking at one intermediate scale.

This works and is robust, but it is expensive (cost grows linearly in the number of mixtures tried) and it explores only the handful of points you thought to try. It also bakes in the scale-transfer assumption. The methods in the next section automate the search and, in DoReMi’s case, find the weights with a single extra proxy run instead of a grid.

Choosing the Weights II: Proxy-Model Methods (DoReMi, Group-DRO & RegMix)¶

Manual ablation searches a few points. DoReMi (Domain Reweighting with Minimax Optimization; Xie et al., 2023) instead learns a mixture in one shot, using a small reference model and a small proxy model, then transfers the learned weights to the large target run. It is built on Group Distributionally Robust Optimization (Group-DRO), so we develop that first.

Distributionally Robust Optimization in one paragraph¶

Ordinary empirical risk minimization (ERM) minimizes the average loss over the training distribution. DRO instead minimizes the worst-case loss over a set of distributions — it is risk-averse, optimizing for the hardest reweighting an adversary could pick. In the group version, the adversary is restricted to reweighting the \(k\) domains:

where \(\Delta_k\) is the probability simplex over domains. The inner \(\max\) puts all weight on whichever domain currently has the highest loss; the outer \(\min\) trains \(\theta\) to bring that worst domain down. The fixed point is a model that is uniformly good across domains — no domain is left behind. This is exactly the property we want from a pretraining mixture: do not let math or low-resource languages collapse just because they are small.

The DoReMi trick: excess loss, not raw loss¶

A naive Group-DRO objective has a problem for pretraining: some domains are intrinsically harder (higher irreducible entropy) than others. The adversary would dump all weight on the hardest-to-model domain forever (e.g., noisy multilingual text), regardless of whether more weight actually helps. DoReMi’s key idea is to measure each domain not by raw loss but by excess loss relative to a fixed reference model trained once on the natural mixture:

Excess loss answers a sharper question: “On which domain is the proxy still far from what’s achievable?” A domain with high irreducible entropy will have high loss for both proxy and reference, so its excess is small once the proxy catches up — the adversary stops over-investing in it. A domain where the proxy lags the reference (lots of headroom) gets upweighted. DoReMi clamps excess at zero (you cannot do better than reference “for free”) and uses it to drive the weights.

The algorithm¶

DoReMi runs three steps:

1. Train a reference model \(\theta_{\text{ref}}\) (small, e.g. 280 M) on the natural mixture \(w^{\text{nat}}\). Record its per-domain losses \(\ell_i(\theta_{\text{ref}})\). (Used only to define the excess-loss baseline.)
2. Train a proxy model \(\theta\) of the same small size with online Group-DRO: at each step, evaluate the proxy’s per-domain excess loss, multiplicatively update domain weights toward high-excess domains, and use those weights to draw the next batch. Average the weights over all steps to get \(\bar{w}\).
3. Train the large target model on the fixed averaged mixture \(\bar{w}\) from step 2. The expensive run uses a static mixture; all the adaptivity happened cheaply in the proxy.

The online update in step 2 is exponentiated gradient ascent on the weights (multiplicative weights / Hedge). Let \(\lambda_i^{(t)}\) be the clamped excess loss of domain \(i\) at step \(t\). The weight update with step size \(\eta\) is

where the renormalization keeps \(w^{(t+1)}\) on the simplex and the small smoothing constant \(c\) (mixing in the uniform distribution) guarantees every domain keeps a floor of weight so it is never starved to zero. The proxy parameters \(\theta\) are trained normally (minimizing weighted loss) while the weights chase the excess — a two-player game whose averaged weights are the output.

The payoff: DoReMi reports faster convergence to a given loss and better downstream performance than the natural mixture, found with one small reference + one small proxy run — far cheaper than a grid of full-size ablations. The reported weights also transfer across an order of magnitude of scale reasonably well, though, as noted, transfer is not perfect and is the method’s main caveat.

The reference implementation is sangmichaelxie/doremi: a PyTorch/HF-Trainer codebase whose main additions are a domain-annotated streaming dataloader and a DoReMiTrainer that maintains the per-domain weight vector and reweights the per-token loss each step. The practical work in adopting it is almost entirely upstream of the algorithm — you must be able to tag every training example with its domain id and to draw batches at arbitrary domain proportions, which is the dataloader machinery we build in the next section.

Regression-based search: RegMix¶

DoReMi spends its budget on one reference plus one adaptive proxy run. RegMix (Liu et al., 2024) makes the opposite bet: train many tiny models — the paper uses hundreds of ~1 M-parameter models on ~1 B tokens each — on mixtures drawn randomly from the simplex (a Dirichlet prior over \(w\)), then fit a regression surrogate (the paper uses gradient-boosted trees, LightGBM) that predicts held-out loss on a target domain from the mixture vector \(w\). Because evaluating the surrogate is free, you can then search over millions of candidate mixtures and take the predicted argmin as your weights.

Three things make this attractive in practice. The proxy runs are embarrassingly parallel and individually trivial, so wall-clock is bounded by your cluster width rather than by a sequential schedule. The surrogate is reusable: re-optimize it for a different target domain, ask “what if I halve code?”, or read off which domains it treats as most influential, all without training anything new. And it is an explicit model of the mixture→loss map, which is exactly the object you want when you later ask whether that map is stable across scale. The cost is that you need many runs and a target metric the surrogate can regress on; RegMix reports matching or beating DoReMi at roughly a tenth of the compute. In 2026, RegMix-style surrogate search and DoReMi-style excess-loss reweighting are the two standard automated options, with manual ablation still the right tool when you only have three or four candidate mixtures in mind.

Online / adaptive mixing during the real run¶

DoReMi freezes the mixture for the target run. An alternative is to keep adapting during the large run — online data mixing. The appeal is that the optimal mixture genuinely changes over training (a model that has mastered easy web text may benefit from shifting weight to math later). The risk is instability and the cost of computing per-domain signals on the fly. Practical online schemes (e.g., Albalak et al.’s Online Data Mixing, and bandit-style approaches) treat each domain as an arm of a multi-armed bandit and use a reward signal — typically the rate of loss decrease on that domain (its learning velocity) — to shift weight toward domains where the model is currently learning fastest, while a smoothing/exploration term keeps every domain sampled. This connects directly to RL-style curriculum (see RL Data, Curriculum & Replay Management), where the same “train on what you’re learning from right now” intuition drives sample selection.

A DoReMi-Style Reweighting Toy You Can Run¶

Let us make all of this concrete with a small, self-contained simulation. We will not train real transformers — that would obscure the mechanism. Instead we model each domain’s loss with a realistic learning curve: per-domain loss falls as a power law in the number of tokens that domain has received, toward a domain-specific floor. This captures the two things that matter for mixing — domains have different floors (difficulty) and different learning rates (transfer/headroom) — and lets us watch the Group-DRO weights respond. The same multiplicative-weights loop transfers directly to a real proxy run; only the loss source changes.

import numpy as np

rng = np.random.default_rng(0)

# -----------------------------------------------------------------------------
# 1. Define k domains. Each has:
#    - floor:  irreducible loss (entropy of that domain's text), in nats/token
#    - rate:   power-law exponent for how fast loss falls with tokens seen
#    - scale:  loss = floor + scale * (tokens_seen + 1)^(-rate)
#    Domains differ in BOTH difficulty (floor) and headroom/speed (scale,rate).
# -----------------------------------------------------------------------------
domains = ["web", "code", "math", "books", "multi"]
floor = np.array([1.70, 1.10, 1.55, 1.80, 2.30])   # math & multi are "hard"
scale = np.array([2.0, 3.5, 4.0, 1.8, 2.2])        # code & math have big headroom
rate  = np.array([0.32, 0.28, 0.22, 0.30, 0.18])   # math & multi learn slowly
k = len(domains)

def domain_loss(tokens_seen):
    """Per-domain held-out loss given cumulative tokens seen in that domain."""
    return floor + scale * np.power(tokens_seen + 1.0, -rate)

# -----------------------------------------------------------------------------
# 2. Reference model: train once on the NATURAL mixture (proportional to pool
#    sizes), then read off its per-domain loss. This defines the excess-loss
#    baseline. We simulate "training" simply by accumulating tokens per domain.
# -----------------------------------------------------------------------------
pool = np.array([3000., 250., 30., 80., 400.])     # unique tokens (in B), illustrative
w_nat = pool / pool.sum()                          # natural mixture
REF_TOKENS = 5.0e4                                 # arbitrary proxy-scale token units

ref_tokens_per_domain = w_nat * REF_TOKENS
ref_loss = domain_loss(ref_tokens_per_domain)
print("natural mixture  :", np.round(w_nat, 3))
print("reference loss   :", np.round(ref_loss, 3))

# -----------------------------------------------------------------------------
# 3. Proxy model with online Group-DRO (DoReMi step 2).
#    - w:        current sampling weights over domains (the adversary's play)
#    - seen:     cumulative tokens per domain (the proxy's "knowledge")
#    Each step we draw a batch split by w, accumulate tokens, recompute the
#    proxy's per-domain loss, form CLAMPED excess loss vs the reference, and
#    apply an exponentiated-gradient (multiplicative-weights) update to w.
# -----------------------------------------------------------------------------
STEPS        = 4000
BATCH_TOKENS = 10.0          # tokens added per step (proxy-scale units)
ETA          = 1.0           # weight learning rate (exp-gradient step size)
SMOOTH_C     = 0.05          # uniform smoothing -> every domain keeps a floor weight

w    = w_nat.copy()          # start from the natural mixture
seen = np.zeros(k)           # proxy has seen nothing yet
w_history = []

for t in range(STEPS):
    # (a) Draw this step's batch according to current weights and "train":
    #     allocate BATCH_TOKENS across domains in proportion to w.
    seen += w * BATCH_TOKENS

    # (b) Proxy's current per-domain loss and CLAMPED excess loss vs reference.
    proxy_loss = domain_loss(seen)
    excess = np.maximum(proxy_loss - ref_loss, 0.0)   # cannot beat ref "for free"

    # (c) Exponentiated-gradient ascent on weights toward high-excess domains.
    w = w * np.exp(ETA * excess)
    w = w / w.sum()                                   # back onto the simplex
    w = (1.0 - SMOOTH_C) * w + SMOOTH_C / k           # uniform smoothing (floor)

    w_history.append(w.copy())

w_bar = np.mean(w_history, axis=0)   # DoReMi outputs the AVERAGED weights
print("\nfinal-step weights:", np.round(w, 3))
print("AVERAGED weights  :", np.round(w_bar, 3))
print("vs natural        :", np.round(w_nat, 3))
print("upweight factor   :", np.round(w_bar / w_nat, 2))

Running this prints something like:

natural mixture  : [0.798 0.066 0.008 0.021 0.106]
reference loss   : [1.767 1.461 2.621 2.022 2.77 ]
final-step weights: [0.305 0.174 0.174 0.174 0.174]
AVERAGED weights  : [0.361 0.172 0.133 0.136 0.198]
vs natural        : [0.798 0.066 0.008 0.021 0.106]
upweight factor   : [ 0.45  2.59 16.62  6.4   1.86]

Read the result. The natural mixture is 80% web; DoReMi-style reweighting collapses web from 80% to ~36% and dramatically upweights the small high-headroom domains — code ~2.6x, math ~17x, books ~6x. This is exactly the qualitative behavior reported for real DoReMi: it pulls weight out of the abundant, lower-headroom domain (web) and into domains where the proxy still has the most room to improve relative to the reference. The math domain, despite a high floor (it is genuinely hard), gets heavily upweighted because its excess — the gap the proxy can still close — is large. Note the multilingual domain, which has the highest floor and the slowest rate (small headroom relative to its difficulty), is upweighted only modestly: high raw loss alone does not earn weight; closeable loss does. That separation is the entire reason DoReMi uses excess loss instead of raw loss.

Two experiments to build intuition (left as exercises you can run in seconds):

• Set ref_loss = 0 (i.e., use raw loss instead of excess). You will see the adversary dump weight onto multi, the highest-floor domain, and refuse to leave — illustrating exactly the pathology excess loss was invented to fix.
• Sweep SMOOTH_C from 0.0 to 0.3. At 0.0, weights can drive a domain’s effective weight toward zero (starvation); larger values keep every domain alive but blur the signal. The DoReMi default lives near the small end.

Implementing a Mixture: The Dataloader and the Libraries That Ship It¶

A mixture is a number on a slide until the dataloader realizes it. Three implementation decisions turn \(w\) into actual batches, and getting them wrong quietly changes the mixture you think you are training on.

Sample at the sequence level, not the token level. Weights are defined over tokens, but you draw sequences. The standard construction is: for each sequence slot in the batch, draw a domain \(i \sim w\), then take that domain’s next context-length window. Because every packed sequence has exactly seq_len tokens, sequence-level multinomial sampling gives token proportions equal to \(w\) in expectation, with \(O(1/\sqrt{\text{batch} \times \text{steps}})\) noise. (A “stratified” alternative — fixing exactly \(w_i \cdot B\) sequences of each domain per batch — removes that noise but requires \(w_i B\) to be near-integral, and is worth it only for very small \(B\) or very small weights.)

Pack within a domain, mix across sequences. Packing (see Chat Templates, Data Formatting & Sequence Packing) concatenates documents to fill the context window. If you pack across domains, a single sequence contains web text followed by math, which corrupts per-domain loss accounting (you can no longer attribute a token’s loss to a domain) and — unless you use intra-document attention masking — lets tokens attend across a domain boundary that carries no real dependency. Build packed shards per domain, then mix at the sequence level. This also makes DoReMi/online-mixing feasible at all: you need per-domain loss, which requires per-domain sequences.

Make the draw resumable and audited. Seed the domain sampler from the global step, not once at run start, so a job that dies at step 40,000 and resumes draws the same domains it would have; and checkpoint each domain’s read position alongside the model (see Checkpointing, Fault Tolerance & Long-Running Jobs). Log realized weights and realized epochs continuously — a source that runs dry, a shard that fails to mount, or a filter that rejects harder in one domain will silently move you off the intended mixture.

import numpy as np

class DomainStream:
    """One domain's pre-tokenized, pre-PACKED token stream (e.g. a np.memmap over
    a .bin shard). Hands out fixed-length windows and cycles forever, tracking
    epochs so the realized epoch budget is a measurement, not an assumption."""
    def __init__(self, tokens, seq_len):
        self.tokens, self.seq_len = tokens, seq_len
        self.pos, self.epochs = 0, 0.0

    def next_seq(self):
        if self.pos + self.seq_len > len(self.tokens):
            self.pos = 0                      # wrapped: we are re-reading this domain
        seq = self.tokens[self.pos:self.pos + self.seq_len]
        self.pos += self.seq_len
        self.epochs += self.seq_len / len(self.tokens)
        return seq

    def state(self):                          # persist alongside the model checkpoint
        return {"pos": self.pos, "epochs": self.epochs}


def mixture_batches(streams, w, batch_size, steps, seed=1337, start_step=0):
    """Yield (step, batch, domain_ids). One domain draw PER SEQUENCE, so token
    proportions equal w in expectation when all sequences have equal length."""
    w = np.asarray(w, dtype=np.float64)
    w = w / w.sum()                           # defensive: must lie on the simplex
    for step in range(start_step, steps):
        # Seed per step, not once per run: restarting at step t reproduces the
        # exact same domain draws, so a resumed run sees the intended mixture.
        rng = np.random.default_rng([seed, step])
        ids = rng.choice(len(streams), size=batch_size, p=w)
        batch = np.stack([streams[i].next_seq() for i in ids])
        yield step, batch, ids


# --- demo: five fake domains sized like the worked example (unique tokens) ----
if __name__ == "__main__":
    SEQ_LEN, BATCH, STEPS = 1024, 32, 200
    names = ["web", "code", "math", "books", "multi"]
    sizes = [3_000_000, 250_000, 30_000, 80_000, 400_000]   # scaled-down pools
    w     = [0.456, 0.228, 0.073, 0.097, 0.146]             # from the worked example

    streams = [DomainStream(np.zeros(n, dtype=np.uint16), SEQ_LEN) for n in sizes]
    counts  = np.zeros(len(names))
    for step, batch, ids in mixture_batches(streams, w, BATCH, STEPS):
        counts += np.bincount(ids, minlength=len(names))
        assert batch.shape == (BATCH, SEQ_LEN)

    print("target   :", np.round(w, 3))
    print("realized :", np.round(counts / counts.sum(), 3))
    print("epochs   :", {n: round(s.epochs, 3) for n, s in zip(names, streams)})

The realized weights land within a fraction of a percent of the targets after only 6,400 sequences. The epoch line is the interesting output: with pools shrunk by \(10^6\) for the demo, math wraps roughly sixteen times in 200 steps. That is precisely the alarm this counter exists to raise — in a real run the same print tells you, at step 40,000 rather than at the post-mortem, that your best small domain has quietly crossed the repetition danger zone. The domain_ids the loop yields are also what you scatter losses into to compute the per-domain \(\ell_i\) that DoReMi and online mixing consume.

The libraries that do this for you¶

You will rarely ship the loop above, but every production stack has its equivalent, and knowing which knob is the mixture is the point:

• Hugging Face datasets — interleave_datasets([...], probabilities=w, seed=..., stopping_strategy=...) mixes streaming IterableDatasets document-by-document. The critical argument is stopping_strategy: the default "first_exhausted" halts the whole mix as soon as any source drains (a 5% math source over a small math corpus will truncate your run early), while "all_exhausted" re-cycles drained sources — i.e. it silently upsamples them. Neither is “the epoch budget you chose”; pick deliberately. Worked end to end in Capstone 14.2: Data — Sourcing, Filtering, Dedup, Tokenize & Pack.
• Megatron-LM / Megatron-Core — --data-path takes alternating weight/prefix pairs (--data-path 0.456 /data/web_text_document 0.228 /data/code_text_document ...); Megatron normalizes the weights and builds a blended index map over the per-domain binary datasets, so the blend is materialized as a deterministic sample-index array rather than sampled online. Because the blend is precomputed, the builder can tell you up front how many times each constituent dataset will be consumed — surface that number, it is your epoch audit for free. See Megatron-LM, DeepSpeed & Parallelism in Practice.
• MosaicML streaming — StreamingDataset(streams=[Stream(remote=..., proportion=0.456), ...]) mixes shards from object storage, with proportion (relative sampling weight), repeat (explicit epochs over a source), or choose (an absolute sample count) as three equivalent ways to say the same thing. Deterministic resumption mid-epoch is a first-class feature, which is exactly the resumability property the loop above hand-rolls.
• datatrove (Hugging Face) — not a mixer, but the pipeline that produces the per-domain, deduplicated, tokenized shards these mixers consume (Pretraining Data: Sources, Crawling & The Data Pipeline).

Whichever you use, the discipline is identical: express the mixture as weights in one config file, and emit realized-weight and realized-epoch metrics from the loader so target and reality can be compared at any point in the run.

Scheduling Over Time: Curriculum, Context Ramps & Annealing¶

So far \(w\) has been a single fixed vector. But the order in which a model sees data, and when it sees the best data, matters too. We now let the mixture be a function of training progress, \(w(t)\), and consider three orthogonal scheduling axes: difficulty (curriculum), sequence length (context ramp), and quality (annealing/mid-training).

Curriculum learning: easy-to-hard¶

Curriculum learning (Bengio et al., 2009) orders examples from easy to hard, mirroring how humans learn. For LLM pretraining the evidence is genuinely mixed — large transformers on shuffled data are remarkably robust, and a naive curriculum often yields little — but specific, targeted curricula do help. The practical recipes that survive contact with reality:

• Length-based difficulty: start with shorter or simpler documents, introduce longer/denser ones later. This overlaps with the context ramp below.
• Quality/complexity-based: introduce highly technical content (dense math, advanced code) after the model has basic linguistic competence, so it does not waste early capacity on tokens it cannot yet model. This is the curriculum intuition behind putting hard reasoning data later in training.
• Skill-staged for code/math: in domains with a natural difficulty gradient (e.g., basic syntax → algorithms → competition problems), staging by difficulty can outperform a uniform shuffle.

The mechanism, when it works: early in training the model’s gradients on very hard examples are high-variance and poorly aligned (it lacks the prerequisites), so those tokens are inefficiently used. Delaying them spends early compute on tokens with cleaner learning signal. The risk: a too-narrow early diet can cause the model to over-specialize and then struggle to adapt (a mild form of the loss-of-plasticity problem; see Continual & Domain-Adaptive Pretraining).

Context-window scheduling¶

Training directly at long context (e.g., 128 K tokens) from step one is wasteful: attention cost grows with sequence length (quadratically for vanilla attention), most early learning needs only local context, and long-document data is scarce. The dominant recipe is a context-length ramp: train the bulk of the run at a short context (e.g., 4 K), then extend to long context in a final phase with appropriate positional-encoding adjustments (RoPE base/theta scaling, etc.). This is a curriculum over sequence length, and it is one of the few curricula with near-universal adoption. Because it is a deep topic of its own — RoPE scaling, position interpolation, data selection for the long phase — we treat it fully in Long-Context Pretraining & Context Extension. For mixing purposes, the key point is that the mixture itself changes in the long-context phase: you upweight naturally long documents (books, repositories, multi-turn transcripts) because they are the only sources that exercise long-range dependencies.

Mid-training and annealing: save the best for last¶

The single most impactful scheduling idea in modern pretraining is the high-quality annealing phase (sometimes called mid-training or the cooldown phase). The recipe, popularized by the MiniCPM team’s analysis and adopted widely (Llama 3, OLMo, and others describe variants):

1. Train the vast majority of tokens on a broad, web-heavy mixture with a roughly constant (or slowly decaying) learning rate.
2. In the final fraction of training (often the last ~10–20% of tokens), simultaneously (a) decay the learning rate sharply toward zero and (b) shift the data mixture toward the highest-quality, most target-relevant data — curated math, code, textbooks, instruction-formatted and reasoning-heavy data.

The interaction between the two is the whole point. Data seen while the learning rate is high and falling fastest has the largest, most lasting effect on the final weights, because those late steps with a decaying LR are where the model settles into its final basin. Putting your best, most capability-dense data exactly there — when each gradient step still moves the weights but the model is no longer being yanked around — imprints those capabilities most strongly. The MiniCPM “Warmup-Stable-Decay” (WSD) schedule makes this explicit: a long stable-LR phase on the broad mixture, then a short decay phase on upweighted high-quality data. It also has a delicious practical benefit: because the stable phase uses a constant LR, you can branch multiple annealing experiments from a single stable checkpoint and try different final mixtures cheaply, instead of re-running from scratch. This is exactly how the book’s capstone model is built: Stack-100M runs ~18 B tokens on a broad web/code/math mix at constant LR, then branches a short decay phase on an upweighted high-quality mix — see Capstone 14.8: Mid-Training — Quality Annealing, Long-Context Extension & Capability Injection, with the mixture weights themselves set in Capstone 14.2.

This connects naturally to the boundary between pretraining and post-training: the annealing-phase mixture often resembles instruction/SFT data (see Supervised Fine-Tuning & Instruction Tuning), and a strong annealing phase reduces how much later fine-tuning is needed. It also overlaps with continual pretraining for domain adaptation (see Continual & Domain-Adaptive Pretraining).

Putting the schedule together¶

A modern end-to-end pretraining schedule, combining all three axes, looks like:

Each phase is a different \(w(t)\), a different context length, and a different point on the LR schedule (see Learning Rate Schedules, Warmup, Batch Size & Hyperparameters). The art of pretraining data is, in large part, the art of designing this table — and then validating it with the proxy and intermediate-scale checks from earlier in the chapter.

Bringing It Together: A Mixing & Scheduling Playbook¶

The decisions in this chapter compose into a repeatable workflow:

1. Partition and deduplicate. Define domains; deduplicate within each so “one epoch” is honest (Data Cleaning, Deduplication & Quality Filtering).
2. Set a target distribution \(p^{\star}\) reflecting what you care about downstream — not what you happen to have.
3. Find base weights. Either run manual mixture ablations at a proxy scale, run a DoReMi-style reference+proxy pass to get \(\bar{w}\) via excess-loss Group-DRO, or fit a RegMix-style regression surrogate over many tiny proxy runs and optimize it.
4. Convert to an epoch budget and check no domain exceeds the repetition danger zone (~4–6 epochs); upsample small high-value domains, downsample abundant low-value ones.
5. Validate at one intermediate scale before committing the full run.
6. Design the schedule \(w(t)\): broad stable phase, optional online velocity-based nudges, a long-context phase that upweights long documents, and a final annealing phase that concentrates the highest-quality data as the LR decays.
7. Account for annealing in the epoch math so your best small domains are not silently over-repeated.

Get these right and you buy capability gains that would otherwise cost a substantial model-size increase — at zero extra FLOPs. Data mixing is one of the highest-leverage, lowest-cost levers in the entire pretraining stack.

Further reading¶

• Xie, Pham, Dong, et al., DoReMi: Optimizing Data Mixtures Speeds Up Language Model Pretraining (2023) — the reference/proxy excess-loss Group-DRO method central to this chapter.
• Sagawa, Koh, Hashimoto, Liang, Distributionally Robust Neural Networks for Group Shifts (Group-DRO, 2020) — the optimization framework DoReMi builds on.
• Muennighoff, Rush, et al., Scaling Data-Constrained Language Models (2023) — the value of repeated tokens and the limits of upsampling.
• Liu et al., RegMix: Data Mixture as Regression for Language Model Pre-training (2024) — regression surrogate over many tiny proxy runs, optimized over the simplex.
• Albalak et al., Online Data Mixing for Language Model Pre-training — bandit-style adaptive mixing during the real run.
• Hu et al. (MiniCPM), MiniCPM: Unveiling the Potential of Small Language Models — the Warmup-Stable-Decay schedule and the high-quality annealing phase.
• Bengio, Louradour, Collobert, Weston, Curriculum Learning (2009) — the original easy-to-hard framing.
• Gao et al., The Pile — an early explicit, documented domain-weighted pretraining mixture.
• The Llama, Gopher, and OLMo technical reports — real-world descriptions of mixture ablations, annealing, and data-schedule design at frontier scale.

Exercises¶

1. A teammate has a small, high-value math set and wants to hit 3 effective epochs on it. They plan to upsample it 3x and then run deduplication on the combined corpus afterward “to clean things up.” Explain why the ordering is backwards, and describe what can go wrong quantitatively if the set already contains 4x internal near-duplication of some passages. What is the correct order of operations?

2. You have a 1.5 T-token training budget (\(D = 1500\) B). Your cleaned, deduplicated pools and chosen mixture weights are:

Compute the effective epochs \(e_i\) for each domain. Which domain is at the edge of the repetition danger zone, and what would you do about it?

3. Explain why DoReMi measures each domain by excess loss \(\text{excess}_i(\theta) = \ell_i(\theta) - \ell_i(\theta_{\text{ref}})\) (clamped at zero) rather than raw loss \(\ell_i(\theta)\). In the chapter’s toy run, the multi domain has the highest reference loss of all domains yet is upweighted only ~1.9x, while math (also high-loss) is upweighted ~17x. Reconcile these two facts.

4. You have a training budget of \(D = 800\) B tokens and these deduplicated pools: web 2000 B, code 200 B, math 20 B, books 60 B. You decide on target epoch counts based on quality and repetition tolerance: web 0.2, code 2, math 5, books 3. Following the chapter’s epoch-budget procedure, convert these into normalized mixture weights \(w_i\), and report the realized epochs after normalization. Does any domain exceed the danger zone?

5. Work one step of DoReMi’s exponentiated-gradient weight update by hand. There are \(k = 3\) domains with current weights \(w^{(t)} = (0.50, 0.30, 0.20)\). The clamped excess losses this step are \(\lambda = (0.20, 0.00, 0.60)\). Use step size \(\eta = 1.0\) and smoothing constant \(c = 0.10\). Compute \(w^{(t+1)}\) using the chapter’s update rule, and state in one sentence what smoothing bought you here.

6. Implementation. The chapter’s toy uses DoReMi-style excess loss (a level: “how far is this domain from achievable?”). The chapter contrasts this with online bandit mixing, which uses loss velocity (a derivative: “where am I improving fastest right now?”) and needs no reference model. Modify the toy’s proxy loop to implement velocity-based online mixing: replace the clamped excess signal with the per-step loss decrease per domain, and drop the reference model entirely. Then describe one qualitative way the resulting weights should differ from the excess-loss version.

import numpy as np

# Reuse domain_loss, domains, pool, w_nat, k from the chapter's toy.
# No reference model is trained: online bandit mixing needs none.

STEPS        = 4000
BATCH_TOKENS = 10.0
ETA_V        = 200.0        # larger: per-step velocities are small
SMOOTH_C     = 0.05

w         = w_nat.copy()
seen      = np.zeros(k)
prev_loss = domain_loss(seen)   # loss before any tokens are seen
w_history = []

for t in range(STEPS):
    # (a) Draw this step's batch by current weights and "train".
    seen += w * BATCH_TOKENS

    # (b) Loss VELOCITY: how much did each domain's loss fall this step?
    cur_loss = domain_loss(seen)
    velocity = np.maximum(prev_loss - cur_loss, 0.0)   # a derivative, not a level
    prev_loss = cur_loss

    # (c) Same multiplicative-weights update, now toward high-velocity domains.
    w = w * np.exp(ETA_V * velocity)
    w = w / w.sum()                                    # back onto the simplex
    w = (1.0 - SMOOTH_C) * w + SMOOTH_C / k            # uniform smoothing (floor)

    w_history.append(w.copy())

w_bar = np.mean(w_history, axis=0)
print("velocity-based AVERAGED weights:", np.round(w_bar, 3))
print("vs natural                     :", np.round(w_nat, 3))

7. You implement the mixture from Exercise 2 (\(D = 1500\) B; math has \(n_{\text{math}} = 25\) B unique tokens at \(w_{\text{math}} = 0.10\)) with Hugging Face interleave_datasets(streams, probabilities=w, stopping_strategy="first_exhausted"). Roughly how many tokens does your run actually produce, and why? What changes if you switch to "all_exhausted"? What must the loader report for you to know whether either behaviour matches your intended epoch budget?