Data Mixtures and Domain Weighting

GPT-3's pretraining corpus was roughly 60% Common Crawl, 22% WebText2, 8% Books1, 8% Books2, and 3% Wikipedia. Those percentages were set by hand, informed by intuition and ablation runs on smaller models. For years, this "best guess with post-hoc ablations" approach was the industry norm, even as training runs crossed tens of thousands of GPU-hours. The question of how much web text versus code versus curated books a model should see turns out to be non-trivial: it shapes vocabulary breadth, factual accuracy, reasoning ability, and toxicity in ways that interact with model scale in surprising ways.

What a data mixture actually is

Before training begins, a pretraining corpus is partitioned into labelled domains. A domain can be as coarse as "web text" or as fine-grained as "StackOverflow Python answers". Each domain \(d_i\) is assigned a weight \(w_i \geq 0\) with \(\sum_i w_i = 1\). During training, at each step the dataloader samples a domain according to those weights, then draws a random batch from that domain's shard.

The simplest weighting scheme is proportional: set \(w_i = |D_i| / \sum_j |D_j|\) where \(|D_i|\) is the token count of domain \(i\). This is how you would train if you simply concatenated everything and shuffled. Proportional weighting means a 10-trillion-token web crawl naturally overwhelms a 50-billion-token books corpus, which might be fine or catastrophic depending on what you want the model to do.

Oversampling is common for high-quality but small sources. If books are 0.5% of raw tokens but you want the model to reason carefully, you might upsample them to 5-10% of training steps, effectively training on each book token 10-20 times. The Chinchilla scaling laws (Hoffmann et al., 2022) showed that for a fixed compute budget, the optimal model trains on far more tokens than earlier practice suggested -- around 20 tokens per parameter. Once you adopt that budget, you often have to repeat high-quality data while using web text only once.

# Simplified sampling logic used by most training frameworks
domain_weights = {"web": 0.67, "code": 0.15, "books": 0.10, "wiki": 0.08}

def sample_batch(domain_weights, domain_shards, batch_size):
    domain = random.choices(
        list(domain_weights.keys()),
        weights=list(domain_weights.values()),
        k=1
    )[0]
    return next(domain_shards[domain])  # pre-shuffled shard iterator

The weights are a hyperparameter. Like learning rate, they interact with scale in ways that make small-model experiments imperfect proxies for large-model behaviour.

Why uniform or proportional weighting is often wrong

Proportional weighting optimises for reducing average loss across the corpus. But average loss is dominated by the largest domain. If web text is 80% of tokens and Wikipedia is 1%, the model can ignore Wikipedia almost entirely and still see great average loss numbers. The result is a model that writes fluently but hallucinates basic facts.

More formally, suppose you care about a downstream task whose data distribution is concentrated in domain \(d_k\). Proportional weighting allocates compute to \(d_k\) roughly proportional to \(|D_k| / \sum_j |D_j|\). If \(d_k\) is small, the model gets limited exposure to it, and generalisation to that task suffers regardless of how many total tokens you train on.

There is also a subtler quality-quantity tension. The Pile (Gao et al., 2021) aggregates 22 diverse subsets from academic papers (Pile-CC, PubMed, GitHub, ArXiv, and so on) and demonstrated that diversity improves cross-domain generalisation. But not all 22 domains are equal quality. Proportional weighting gives more weight to large noisy sources by construction. Upsampling small high-quality sources like academic papers is almost always beneficial, but quantifying "how much" requires experiments.

Automated domain weight optimisation: DoReMi

The most principled published method for automated weight optimisation is DoReMi (Xie et al., 2023). The algorithm has two phases:

1. Train a small "proxy" model (280M parameters in the paper) using group distributionally robust optimisation (group DRO) over domains. Group DRO minimises the worst-case loss across domains rather than average loss, which forces the model to allocate capacity to every domain, not just the dominant ones.

2. Extract the domain weights implied by the proxy model's loss surface, then use those weights to train the full-scale target model (8B parameters in the paper).

The key insight is that the proxy model is cheap to run, so you can afford to search over domain weights at proxy scale and transfer the result to full scale. On the Pile, DoReMi's discovered weights differed sharply from the default ones -- books were upsampled substantially, while certain web-crawl subsets were downsampled. The result was a 2.6x reduction in training steps to reach baseline perplexity, and roughly 6.5 percentage-point gains on downstream tasks.

A practical limitation: DoReMi still requires running a full proxy training loop, which for a 280M model on 200B tokens costs on the order of thousands of GPU-hours. For smaller teams, the overhead is non-trivial.

Mixing laws: predicting performance before you train

A more recent direction asks whether there is a mathematical function \(f(w_1, ..., w_n, N, T)\) that predicts model loss given mixture weights \(w_i\), model size \(N\), and token count \(T\), without running the full training job. If such a function exists and generalises, you can grid-search weights at small scale and extrapolate.

Ye et al. (2024, ICLR 2025) showed empirically that for each domain \(d_i\), the per-domain loss follows a power-law in the number of tokens from that domain:

where \(a_i\), \(b_i\), \(c_i\) are domain-specific constants fitted from small experiments. The overall loss is then a weighted sum \(\mathcal{L} = \sum_i w_i \mathcal{L}_i\). Because this separates across domains, you can optimise \(w_i\) analytically or with cheap search. Applied to a 1B-parameter model on 100B tokens (RedPajama data), their optimal weights matched training 48% longer using the default mixture, at no additional compute cost.

Ge et al. (2024) extended this with BiMix, a bivariate law that jointly models domain proportions and total data volume, achieving under 0.2% mean relative error in loss prediction. The practical implication is that mixing is becoming a numerical optimisation problem rather than an art form.

When it falls down

Proxy-target mismatch. DoReMi and mixing-law approaches assume that weight rankings learned at small scale transfer to large scale. This holds roughly for loss, but less reliably for capability-specific benchmarks. A 280M-parameter proxy may not correctly predict that upsampling mathematical text benefits a 70B model on complex reasoning tasks; the benefit may only emerge at scale where certain capabilities crystallise.

Domain label granularity. Both DoReMi and mixing laws operate on predefined domain boundaries. If you label all web text as one domain, the optimiser cannot distinguish between high-quality news prose and low-quality forum spam within that domain. Coarse labels obscure quality variation. But using very fine-grained domains increases the number of weights to optimise and reduces sample size per domain, making the fitting harder.

Dynamic curricula versus static weights. Most methods assume a fixed weight schedule for the full training run. There is evidence that the optimal weights shift over training: early in training the model benefits from broad coverage; later in training, upsampling high-quality or task-specific data becomes more valuable. Static weights are a compromise. Curriculum learning with dynamically adjusted weights is an active research area without a widely adopted solution.

Small-domain starvation. If a domain has very few tokens and a high weight, the model will see repeated data from it. Repeating data more than roughly 4 epochs yields diminishing returns and eventually hurts generalisation (Muennighoff et al., 2022). If your high-quality domain is genuinely small, upsampling has a ceiling.

Evaluation contamination interacts with weighting. If your benchmark evaluation set leaks into a domain you heavily oversample, the resulting model performance is inflated. Domain-level decontamination and weight optimisation need to be coordinated, but most pipelines treat them as separate steps.

Further reading

• Xie et al. (2023). "DoReMi: Optimising Data Mixtures Speeds Up Language Model Pretraining." NeurIPS 2023. https://arxiv.org/abs/2305.10429
• Ye et al. (2024). "Data Mixing Laws: Optimising Data Mixtures by Predicting Language Modelling Performance." ICLR 2025. https://arxiv.org/abs/2403.16952
• Gao et al. (2021). "The Pile: An 800GB Dataset of Diverse Text for Language Modelling." https://arxiv.org/abs/2101.00027
• Muennighoff et al. (2023). "Scaling Data-Constrained Language Models." https://arxiv.org/abs/2305.16264