Suppose you have done everything in the last two sections. You have 15 trillion clean, deduplicated tokens of English web text, a trillion tokens of code, a hundred billion of math, a few tens of billions of books and papers, five billion of Wikipedia, and text in a hundred other languages ranging from a trillion tokens to a few million. You plan to train on 15 trillion tokens. What fraction of each batch should come from each source?

This one decision shapes the model as much as almost any architectural choice. Too little code and the model cannot program, too much and its prose gets worse. Sample languages in proportion to their size and the model barely learns Swahili; sample them equally and it sees the same few million Swahili tokens a hundred times. This section is about how labs set those proportions, change them during training, decide how often data may repeat, and when it is safe to train on text that a model wrote.

Domain weights, tokens, and epochs

Call the sources domains $i = 1, \dots, m$, with $N_i$ tokens available in domain $i$. A mixture is a vector of weights $\mathbf{w}$ on the probability simplex, $w_i \ge 0$ and $\sum_i w_i = 1$, and each training token is drawn from domain $i$ with probability $w_i$. With a total budget of $D$ tokens, domain $i$ contributes $w_i D$ tokens, so it is seen for

In plain words, a domain's epoch count is how much of the budget it gets divided by how much of it exists. Upweighting a small domain does not create data, it repeats it.

GPT-3 is the classic worked example (Brown et al. 2020). It trained on 300 billion tokens with the mixture below. Common Crawl was 60 percent of the batches but only 0.44 epochs, so most of it was never seen, while Wikipedia was 3 percent of the batches and was seen 3.4 times.

The weights say what the GPT-3 team believed. Small, high-quality sources were worth repeating, and the huge web crawl was worth sampling sparsely. Modern mixtures express the same belief in a different vocabulary. LLaMA 3's final pretraining mix was roughly 50 percent general knowledge, 25 percent mathematics and reasoning, 17 percent code, and 8 percent multilingual text (Dubey et al. 2024), with the categories assigned by classifiers rather than by source.

Recomputing from the table gives 3.5 epochs for WebText2, 2.0 for Books1, and 3.0 for Wikipedia rather than the paper's 2.9, 1.9, and 3.4, because the published weights and token counts are rounded (the weights even add to 1.01). The picture is the same. Compared with sampling in proportion to size, the mixture moves mass from the web toward the curated sources, by a factor of about 6 for WebText2, 5 for Wikipedia, and 3 for Books1.

Temperature sampling across languages

The same question has a clean, much-studied answer for languages. Let $q_i = N_i/\sum_j N_j$ be language $i$'s natural share of the data. Temperature sampling draws from

sometimes written with a temperature $\tau = 1/\alpha$, as $p_i \propto q_i^{1/\tau}$. In plain words, raise every language's share to a power below one and renormalize, which shrinks the gaps between big and small languages.

The two ends are easy to read. With $\alpha = 1$ we sample in proportion to size, so every language gets the same number of epochs, $e_i = p_i D/N_i = D/\sum_j N_j$, and small languages get almost no tokens. With $\alpha = 0$ every language gets the same share $1/m$, so small languages get many epochs. In between, taking logs shows what the power does,

so a language 1,000 times smaller than English in the data is $1000^{\alpha}$ times smaller in the batches, about 8 times for $\alpha = 0.3$. The ratio of epochs between them is $e_j/e_i = (q_i/q_j)^{1 - \alpha}$, so the small language is repeated $1000^{0.7} \approx 126$ times more often per token. That is the trade-off in one formula. Lowering $\alpha$ buys a small language more exposure at the price of repetition.

Multilingual models settled on low values. mBERT used $\alpha = 0.7$, XLM used 0.5, and XLM-R and mT5 both used 0.3 (Conneau et al. 2020; Xue et al. 2021). XLM-R also named the trade-off it fights against, the curse of multilinguality. For a fixed model size, adding more languages improves low-resource languages at first, through transfer from related ones, and then starts to hurt every language, because capacity is shared.

With $\alpha = 0.3$, Yoruba's 100 million tokens would be repeated 146 times in a trillion-token run, far past the point where repetition helps (we will see where that point is at the end of this section). That is what UniMax (Chung et al. 2023) fixes. It allocates the budget as uniformly as possible across languages, but caps every language at a maximum number of epochs and passes the excess on to the larger ones. Processing languages from smallest to largest, each gets the smaller of its uniform share of the remaining budget and its cap.

import numpy as np

def temperature_mixture(sizes, alpha, budget):
    sizes = np.asarray(sizes, dtype=float)
    p = sizes / sizes.sum()                       # placeholder: ignores alpha
    return p, p * budget / sizes

def alpha_for_share(sizes, i, target):

import numpy as np

def unimax(sizes, budget, max_epochs):
    sizes = np.asarray(sizes, dtype=float)
    return np.full(len(sizes), budget / len(sizes))       # placeholder: ignores the cap

Learning the mixture with DoReMi

Hand-set weights encode intuition. Can we learn them instead? Training the full model once per candidate mixture is out of the question, so every practical method uses small proxy models. DoReMi (Xie et al. 2023) is the best-known, and its idea is worth knowing in detail because it shows how a minimax objective becomes a few lines of code.

DoReMi first trains a small reference model on some default mixture. It then trains a second small proxy model, while adjusting the domain weights $\boldsymbol\alpha$ to emphasize the domains where the proxy is furthest behind the reference. Formally, it solves the group distributionally robust problem

where $\mathcal{L}_i$ is the loss on domain $i$ and $\Delta^m$ is the simplex. In plain words, the proxy model tries to do well on every domain, and an adversary keeps pointing it at the domain where it has the most excess loss compared with the reference. Excess loss is the point. A domain that is intrinsically hard (random-looking text) has high loss for both models and no excess, so it does not soak up weight, while a domain that is learnable but not yet learned does.

The adversary's update is an exponentiated gradient step. At each step, compute the per-domain excess loss on the current batch, clipped at zero, $\lambda_i = \max\{\mathcal{L}_i(\theta) - \mathcal{L}_i(\theta_{\text{ref}}), 0\}$, and set

with step size $\eta = 1$ and smoothing $c = 10^{-3}$ in the paper. The weights DoReMi outputs are the average of $\boldsymbol\alpha$ over the whole proxy run, not the final value.

The results were striking. Weights found with a 280M-parameter proxy were used to train an 8B model, 30 times larger, on The Pile's 22 domains. The 8B model improved its average one-shot accuracy on generative tasks by 6.5 points and reached the baseline's final accuracy in 75,000 steps instead of 200,000, 2.6 times faster. The learned weights also say something. The web crawl, Pile-CC, went from 11 percent of the default mixture to 61 percent, a vote for broad web text over the hand-curated domains.

import numpy as np

def doremi_update(alpha, proxy_loss, ref_loss, eta=1.0, c=1e-3):
    raise NotImplementedError

def doremi_average(alpha0, proxy_losses, ref_losses, eta=1.0, c=1e-3):
    raise NotImplementedError

DoReMi is one of several approaches, and none has become standard. Others fit a regression from mixture weights to validation loss across many tiny runs and pick the predicted best mixture, or fit data mixing laws that predict loss as a function of the weights. LLaMA 3's team ran scaling-law experiments on small models trained on candidate mixtures and extrapolated (Dubey et al. 2024). The honest summary is that mixtures at frontier labs are set by a combination of these proxy experiments and many careful ablations, and that the proxy's verdict does not always survive at full scale.

Curricula and annealing

Should the order of the data matter? Classic curriculum learning, easy examples first, has a weak record in language-model pretraining. The model sees each token about once, and the shuffled mixture works about as well as hand-designed orderings in most published comparisons. One kind of ordering has become universal, though, and it is about the end of training.

As the learning rate decays at the end of a run, the loss drops sharply, and the data the model sees during that drop has an outsized effect on the final model (§8.7 explains why, through the noise in the parameters). So labs anneal. During the last part of training, they switch to a mixture richer in their best data. LLaMA 3 annealed its learning rate to zero over the final 40 million tokens with high-quality sources upsampled, then averaged checkpoints from that phase (Dubey et al. 2024). OLMo 2 made it an explicit second stage, mid-training, on a separate mix called Dolmino of high-quality web text, math, and curated sources, for 50 billion tokens on the 7B model. They ran it three times with different data orders and averaged the three resulting models, which they call souping (Team OLMo 2024).

Annealing also became a tool for measuring data. To ask whether a new dataset helps, LLaMA 3's team took an 8B model halfway through training and annealed it to zero over 40 billion tokens with 30 percent of the mixture from the new dataset, then compared benchmarks. This is far cheaper than training a model from scratch per dataset. In one such test, annealing on high-quality math and code improved the 8B model's GSM8K score by 24.0 percent and MATH by 6.4 percent, but that test mixture included those benchmarks' own training sets, which the final run left out. On the 405B model, the gains from annealing were negligible. Bigger models seem to need less of this last-minute help.

Code and math

Code and math take a share of the mixture far larger than their share of the web, for two reasons. They are skills people pay for, and practitioners believe they improve reasoning in general. The second belief has some evidence behind it. Aryabumi et al. (2024) ablated code in pretraining and found that including it improved not only coding but also natural-language reasoning and world-knowledge tasks, an empirical finding at the scales they tested, up to 2.8B parameters.

Both need their own pipelines, because the filters of §8.2 are tuned for prose and destroy code and math (C4's curly-brace rule, Gopher's alphabetic-word rule). Code comes from repositories, deduplicated at the file level, filtered by license, by file statistics (line length, fraction of alphanumeric characters, autogenerated headers), and often by a quality classifier. Math is harder to find. The best results come from training a classifier to recognize math pages and iterating, recalling more math pages from the crawl with each round. DeepSeekMath built a 120-billion-token math corpus from Common Crawl this way, and math-specific text extraction must preserve LaTeX and MathML that general extractors throw away.

Synthetic data, and the risk of model collapse

The web is finite and its best parts are small. A strong model can write new training data, textbook-style explanations, exercises with solutions, rephrasings of web pages in a cleaner style. This works surprisingly well. The phi-1 model (Gunasekar et al. 2023) had 1.3B parameters and was trained on about 7 billion tokens, about 6 billion of code from The Stack and StackOverflow, filtered for educational value with a classifier trained on GPT-4's judgments, plus under a billion tokens of synthetic textbooks written by GPT-3.5. After about eight passes over that data, and a short fine-tuning on about 180 million tokens of synthetic exercises, it reached 50.6 percent on the HumanEval coding benchmark, ahead of many models more than ten times its size.

What is the risk? If every generation of models trains on text written by the previous generation, errors compound. Shumailov et al. (2024) named the result model collapse. The tails of the original distribution disappear first, then the distribution narrows toward a few modes. The simplest possible model shows why, and it is worth deriving because the mechanism is general.

Let generation 0 be a Gaussian, $\mathcal{N}(\mu_0, \sigma_0^2)$. Generation $t + 1$ draws $M$ samples from generation $t$ and fits a Gaussian by maximum likelihood, so its variance is the sample variance $s^2 = \frac1M\sum_k (x_k - \bar x)^2$. For Gaussian samples, $Ms^2/\sigma_t^2$ follows a chi-square distribution with $M - 1$ degrees of freedom, which has mean $M - 1$, so

In plain words, each generation's fit is slightly too narrow on average, because a finite sample rarely includes the extremes, and the shrinkage compounds. It is worse than the mean suggests. The logarithm of the variance is a sum of independent terms $\log(\chi^2_{M-1}/M)$, each with mean about $-2/M$ (since $\mathbb{E}\log\chi^2_{k} = \psi(k/2) + \log 2 \approx \log k - 1/k$, where $\psi$ is the digamma function). So a typical run's variance decays like $e^{-2t/M}$, twice as fast as the mean, which is held up by rare lucky runs. Meanwhile the mean $\mu_t$ does a random walk. Both effects show up in the cell.

The standard deviation collapses toward zero within a few hundred generations, at about the rate the drift predicts. A categorical distribution over a vocabulary collapses in an even more blunt way, which the widget below simulates. Once a rare token draws zero samples in some generation, its fitted probability is zero and it can never come back.

The good news is that collapse is a property of replacing data, not of using synthetic data. Gerstgrasser et al. (2024) showed that if each generation trains on all the real and synthetic data accumulated so far, the error stays bounded, and in a linear model they proved a bound independent of the number of generations. Real pipelines are closer to accumulation. They keep the human data, add a limited amount of synthetic data, filter it, and draw it from strong models rather than from the model being trained.

import numpy as np

def collapse(p0, n, generations, accumulate, seed=0):
    raise NotImplementedError

Two cautions about synthetic data apply even without collapse. It inherits the generator's mistakes and blind spots, so it needs the same filtering and verification as web data (for code and math, executing tests or checking answers is the strongest filter). And it can carry benchmark contamination in disguise, if the generator memorized benchmark questions and paraphrases them into the training set, so the decontamination of §8.3 should run on synthetic data too. Where synthetic data shines most is not pretraining but post-training, where a small amount of targeted data shapes behavior (§14.2), and in self-improvement loops that are a field of their own (training on a model's own filtered outputs).

Epochs and repetition

If good data is scarce, how many times can you repeat it? The best evidence comes from Muennighoff et al. (2023), who trained hundreds of models on a fixed amount of unique data repeated for different numbers of epochs. Up to about 4 epochs, repeated data was almost as good as fresh data. Beyond that, each extra pass was worth less, and around 40 epochs more repetition added almost nothing. They summarized this as an effective data formula, in which $U$ unique tokens repeated $R$ extra times are worth $D' = U\big(1 + R^\star(1 - e^{-R/R^\star})\big)$ fresh tokens, with a fitted constant $R^\star \approx 15.4$, so infinite repetition is worth at most about 16 unique epochs. §9.4 derives what that formula means for compute-optimal training when data runs out, and implements it (Exercise 9.4c).

For a mixture, the rule of thumb is what matters. Compute every domain's epochs with $e_i = w_iD/N_i$ from the start of this section, and keep the small, precious domains to a few epochs. That is why the web is trained on in under one epoch at most frontier labs, while small curated sources (Wikipedia, books, the best math) are repeated a few times, as GPT-3's mixture already did. Be careful with the distribution of repetition, too, not only its average. Hernandez et al. (2022) showed that repeating a small subset many times is far more harmful than repeating everything a few times (§8.3).

In PyTorch

In a real data loader, the mixture is a sampler over several datasets. Here is a minimal one with PyTorch's WeightedRandomSampler, plus DoReMi's weight update written with tensors, run on a toy problem with three domains of different difficulty.

Domain 2 has only 10 examples, so its 20 percent share means it is repeated about 100 times in 5,000 draws, the trap of the first section in miniature. In the DoReMi toy, domain 1 starts behind the reference and catches up, and its weight rises and stays high even after its excess loss has faded. That is the multiplicative update at work. Ignoring the smoothing, $\alpha_{T,i} \propto \alpha_{0,i}\exp\big(\eta\sum_{t\le T}\lambda_{t,i}\big)$, so the weights remember the cumulative excess of each domain, and a domain loses weight only when others build up excess of their own. Domain 2 is hard for both models and never gains weight, which is the point of using excess loss.

Read more

• Xie et al. (2023). DoReMi, Optimizing Data Mixtures Speeds Up Language Model Pretraining. The minimax objective, the update rule, and the transfer from a 280M proxy to an 8B model.
• Muennighoff et al. (2023). Scaling Data-Constrained Language Models. How much repeated data is worth, with the effective-data formula.
• Chung et al. (2023). UniMax, Fairer and More Effective Language Sampling for Large-Scale Multilingual Pretraining. The epoch-capped alternative to temperature sampling, with a clear comparison table.
• Conneau et al. (2020). Unsupervised Cross-lingual Representation Learning at Scale. XLM-R, the $\alpha = 0.3$ sampling choice, and the curse of multilinguality.
• Shumailov et al. (2024). AI models collapse when trained on recursively generated data, with Gerstgrasser et al. (2024), Is Model Collapse Inevitable?. The problem and the condition under which it goes away.
• Gunasekar et al. (2023). Textbooks Are All You Need. What a small model can learn from filtered and synthetic textbook-quality data.
• Dubey et al. (2024). The Llama 3 Herd of Models. Section 3.1 describes the data mix, the annealing experiments, and the classifiers behind them.

Watch next

▶

▶

▶