Pretraining and scaling laws

Problem

A decoder-only language model is a stack of transformer blocks with a next-token head. Its architecture is settled in a page; what makes it useful is pretraining: running it over a huge unlabelled corpus with a self-supervised objective, so that the targets come from the text itself and no human labels are needed. Everything later, fine-tuning, preference training and prompting, starts from the weights pretraining produces.

Pretraining is also by far the most expensive step: months on thousands of accelerators. Before spending that, a team must decide how big the model should be, how many tokens to train it on, and what text those tokens come from. Getting the split between parameters and tokens wrong wastes a large share of the budget, and the only way to decide it in advance is to measure how loss falls on small runs and extrapolate. Those measurements are the scaling laws.

Intuition

Predicting the next word of arbitrary text well requires knowing grammar, facts, arithmetic, code and how arguments go, so a model that gets good at it learns all of those as a side effect. The loss is just cross-entropy on the next token, averaged over every position of every document.

That loss falls smoothly and predictably as the model, the data and the compute grow, following a power law: each tenfold increase buys roughly the same fractional improvement. Because it is so regular, a few small training runs can predict the loss of a run a thousand times larger. The open question is how to spend a fixed budget. A bigger model learns more per token but costs more per token; at a fixed number of FLOPs, more parameters means fewer tokens. There is a best size for every budget, and the two landmark papers on it disagree about where it is.

Mechanics

Pretraining objectives

Causal language modelling is what GPT-style models use. Under a causal attention mask, position i sees tokens 1…i and predicts token i + 1, so one forward pass over n tokens yields n − 1 training targets, every token but the first. The loss is the mean cross-entropy over them; its exponential is the perplexity. Because the model only ever looks left, the same network generates text one token at a time with no change.

Masked language modelling is BERT's objective for encoder-only models. 15% of positions are picked; of those, 80% are replaced by a [MASK] token, 10% by a random token and 10% left unchanged, so that the model cannot rely on [MASK] appearing at fine-tuning time. The loss covers only the picked positions. The model sees both sides of every gap, which gives better representations for classification and retrieval, but it learns from only 15% of the tokens it reads and cannot generate text left to right.

Span corruption is T5's objective for encoder-decoder models. Contiguous spans covering about 15% of tokens are each replaced by one sentinel token; the decoder outputs each sentinel followed by the tokens it hid. With an average span length of 3, inputs and targets are both shorter than the original text, which makes the objective cheap. Later work such as UL2 mixes several such denoisers with causal LM; today's largest models almost all use plain causal LM because one objective then serves both pretraining and generation.

Data mixtures and deduplication

A pretraining corpus is a mixture of sources: filtered web crawl, code, books, scientific papers, encyclopedias, forums. Each gets a sampling weight, chosen by small ablation runs, and the weight times the token budget sets how many times each source is read. Small high-quality sources are typically upweighted and so repeated for several epochs while the web is not even read once. Muennighoff and colleagues found that repeating data up to about 4 epochs is almost as good as fresh data, with returns falling quickly after that.

Web text is full of duplicates: mirrors, boilerplate, templated pages. Duplicates waste compute, increase memorization of the repeated text, and leak test sets into training. Exact deduplication hashes each document, or each long substring, and drops repeats. Near-duplicate detection uses MinHash: split each document into overlapping word n-grams (shingles) and measure similarity as the Jaccard index |A ∩ B| / |A ∪ B| of the shingle sets. For a random hash function, the probability that two sets have the same minimum hash equals their Jaccard index, so k hash functions give a k-number signature whose agreement rate estimates the similarity. Comparing every pair is still quadratic, so locality-sensitive hashing cuts each signature into b bands of r rows and only compares documents that match exactly on some band. A pair with similarity s becomes a candidate with probability 1 − (1 − s^r)^b, an S-shaped curve whose threshold sits near (1/b)^1/r. Quality filters (language ID, heuristics, classifiers trained to spot reference-quality text) and removal of personal data and benchmark text run alongside.

Compute: C ≈ 6 · N · D, derived

Almost all of a transformer's arithmetic is matrix multiplication by its weights. Take one linear layer with a weight W of shape f_in × f_out applied to D tokens. Multiplying an (D × f_in) input by W takes D · f_in · f_out multiply-adds, which is 2 · D · f_in · f_out FLOPs: 2 FLOPs per parameter per token in the forward pass. Backpropagation needs two more products of the same size: the gradient with respect to the input, ∂L/∂X = ∂L/∂Y · Wᵀ, to pass back to the previous layer, and the gradient with respect to the weight, ∂L/∂W = Xᵀ · ∂L/∂Y, for the update. That is 4 more FLOPs per parameter per token, 6 in total. Summed over every weight matrix:

Kaplan and colleagues add the attention scores, which multiply activations by activations rather than by weights: about 2 · n_layer · n_ctx · d_model extra forward FLOPs per token, small next to 2N unless the context is much longer than d_model. Embedding lookups are not multiplications and are usually excluded from N, as is the optimizer's elementwise update. The formula ignores activation recomputation (which adds another forward pass) and counts useful work only; dividing it by hardware peak gives a lower bound on time, and the achieved fraction of peak is the model FLOPs utilization, or MFU.

For a GPT-style block with model width d, attention has four d × d projections (Q, K, V, output), 4d² parameters, and an MLP with hidden width 4d has 8d²: about 12 · n_layer · d² non-embedding parameters in all.

The Kaplan scaling laws

Kaplan and colleagues (OpenAI, 2020) trained models from under a million to over a billion parameters and found that test loss follows a power law in each resource when the others are not the bottleneck:

with N counting non-embedding parameters, D in tokens and L in nats per token on their WebText2 data. The shape of the architecture (depth versus width) mattered far less than N. Fitting how the best model for each compute budget grows, they found that model size should grow much faster than data, N_opt ∝ C^0.73, so most of any extra compute should go into a bigger model trained on relatively few more tokens. GPT-3, with 175B parameters trained on about 300B tokens, followed that advice.

The Chinchilla scaling laws

Hoffmann and colleagues (DeepMind, 2022) trained over 400 models from 70M to 16B parameters and reached a different answer. One key difference: Kaplan's runs used a learning-rate schedule of the same length regardless of how many tokens a run would actually use, which understates what small models reach on long training. Chinchilla tuned the cosine schedule to each run's length, and used three methods to locate the best size. The third fits one formula for loss as a function of both N and D:

E is the irreducible loss of the data, the entropy of text that no model can remove; the other two terms are the penalties for a finite model and finite data. To find the best size for a budget, substitute D = C/(6N) and set the derivative with respect to N to zero. That gives αA/N^α = βB/D^β, and solving:

With the fitted values, a = 0.45 and b = 0.55; the other two methods gave a and b both close to 0.5. Either way, parameters and tokens should grow at about the same rate, roughly equally for each tenfold more compute. The widely used summary is about 20 tokens per parameter, which turns C = 6 · N · 20N into N = √(C / 120). By that measure GPT-3 and Gopher were far too large for their data, and Chinchilla, at 70B parameters on 1.4T tokens with the budget of 280B-parameter Gopher, beat it on almost every benchmark.

Beyond compute-optimal

Chinchilla-optimal minimizes training compute for a target loss, but a deployed model's cost is dominated by inference, which scales with N and the number of tokens served. A smaller model trained on far more tokens than 20 per parameter costs more to train but less to run, so current open models are deliberately "over-trained": Llama 3's 8B model was trained on over 15T tokens, close to 2,000 tokens per parameter. Loss keeps falling with more data, just more slowly than the compute-optimal recipe would buy.

Emergent abilities, and caveats

Wei and colleagues catalogued tasks where small models score at chance and, past some scale, accuracy jumps: multi-digit arithmetic, some reasoning benchmarks. They called these emergent abilities, abilities that cannot be predicted by extrapolating smaller models. Schaeffer, Miranda and Koyejo argued that many of these jumps come from the metric, not the model. Exact-match accuracy gives credit only if every token of the answer is right, so with per-token accuracy p and a k-token answer it behaves like p^k: a smooth rise in p looks like a sudden jump. Measured with a continuous metric such as per-token log-likelihood, the same models improve smoothly. Some capabilities do appear abruptly on any metric, and small benchmarks add noise, so the honest summary is: the loss is predictable, specific downstream abilities much less so, and a claimed jump should be checked against a smoother metric before it is believed.

PyTorch v2.15.0 supplies the loss and the sampling operations; data pipelines, deduplication and scaling-law fits live in separate tools.

Worked example

Implementation

The causal-LM loss with padding ignored, masked-LM corruption, and the sizing arithmetic. model is any decoder that returns logits of shape (B, T, V).

import math
import torch
import torch.nn.functional as F

def causal_lm_loss(model, tokens, pad_id):
    logits = model(tokens[:, :-1])                       # position i predicts token i + 1
    targets = tokens[:, 1:].masked_fill(tokens[:, 1:] == pad_id, -100)
    return F.cross_entropy(logits.reshape(-1, logits.size(-1)), targets.reshape(-1), ignore_index=-100)

def mask_for_mlm(tokens, mask_id, vocab_size, rate=0.15, generator=None):
    labels = tokens.clone()
    picked = torch.bernoulli(torch.full(tokens.shape, rate), generator=generator).bool()
    labels[~picked] = -100                               # loss only on picked positions
    roll = torch.rand(tokens.shape, generator=generator)
    inputs = tokens.clone()
    inputs[picked & (roll < 0.8)] = mask_id
    swap = picked & (roll >= 0.8) & (roll < 0.9)
    inputs[swap] = torch.randint(vocab_size, tokens.shape, generator=generator)[swap]
    return inputs, labels                                # the last 10% keep their token

def chinchilla_size(compute, tokens_per_param=20):
    n = math.sqrt(compute / (6 * tokens_per_param))
    return n, tokens_per_param * n

print(chinchilla_size(6 * 70e9 * 1.4e12))               # (7e10, 1.4e12)

Real pretraining code packs documents end to end into fixed-length sequences separated by an end-of-text token rather than padding, so almost no positions are wasted; the attention mask may also stop tokens attending across document boundaries. The loss is usually computed in float32 even under mixed precision, because the softmax over a large vocabulary is sensitive to rounding. A scaling-law study trains a ladder of small models at several budgets, fits the curves, and only then commits to the large run.

Tradeoffs

• Causal vs masked vs span objectives. Causal LM learns from every token and generates natively; masked LM gives bidirectional representations suited to classification and retrieval but learns from 15% of tokens; span corruption is cheap per token and suits encoder-decoder models for translation and summarization.
• Compute-optimal vs inference-optimal. Chinchilla sizing minimizes training cost for a given loss; a model that will serve billions of tokens is cheaper overall if it is smaller and trained far longer.
• Repeating data. High-quality data is finite. A few epochs over it costs little; many epochs mean diminishing returns and more memorization, which pushes teams toward better filtering and synthetic data.
• Deduplication strength. Aggressive near-dedup removes boilerplate and test leakage but can also remove legitimately similar documents such as licences and standard code; the band and row settings choose where to cut.
• Trusting a fit. Scaling laws are fits on one data distribution, tokenizer and architecture; the constants do not transfer, the exponents only roughly, and Kaplan and Chinchilla reached different allocations from different training choices.

Connections

References

• Kaplan, McCandlish, Henighan, Brown, Chess, Child, Gray, Radford, Wu and Amodei. Scaling Laws for Neural Language Models. 2020. arXiv:2001.08361
• Hoffmann, Borgeaud, Mensch et al. Training Compute-Optimal Large Language Models. NeurIPS 2022. arXiv:2203.15556
• Devlin, Chang, Lee and Toutanova. BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding. NAACL 2019. arXiv:1810.04805
• Raffel, Shazeer, Roberts et al. Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer. JMLR 2020. arXiv:1910.10683
• Muennighoff, Rush, Barak et al. Scaling Data-Constrained Language Models. NeurIPS 2023. arXiv:2305.16264
• Wei, Tay, Bommasani et al. Emergent Abilities of Large Language Models. TMLR 2022. arXiv:2206.07682
• Schaeffer, Miranda and Koyejo. Are Emergent Abilities of Large Language Models a Mirage? NeurIPS 2023. arXiv:2304.15004
• Broder. On the Resemblance and Containment of Documents. Compression and Complexity of Sequences, 1997.
• PyTorch v2.15.0: nn.CrossEntropyLoss, F.cross_entropy, torch.bernoulli, torch.multinomial.