Vol. XII · Mid-Training Plate 3 of 10upsampling, domain mixture, late upsampling, code, mathematics, curriculum, data orderLD

Upsampling Domains.

The same extra tokens of a scarce domain, spread through the run or concentrated in the decay — and the decay wins.

A data mix is a budget (plate IX.8), and the scarce domains — code, mathematics, scientific reasoning — are often upsampled beyond their natural share. Mid-training adds a choice of timing: the extra tokens can be spread evenly through the run, or concentrated at the end, during the learning-rate decay, when the model is settling (plate XII.2). Recipes increasingly do the latter, and Llama 3’s report describes upsampling mathematical data and code in its final stages.

The workshop trains a softmax bigram model — one table of logits for which character follows which, trained by Adam — on two domains of the course corpus: plates from the computing courses and from the mathematics and physics courses. Three runs spend the same 2 000 steps: one keeps science at 10% throughout; one raises it to the chosen share in the final 400 decaying steps; and one spreads the same number of science tokens evenly through the run.

Below: held-out loss on the science text through training, for the three schedules. Right, the final losses on both domains.

What you will learn

Describe upsampling and why scarce domains need it.
Compare spreading extra tokens with concentrating them late.
Measure the effect on both domains.
Relate the result to published recipes.

Workshop

Softmax bigram model on characters, Adam, batches of 256 positions, 2 000 steps; learning rate 0.05, decayed linearly to zero over the last 400. Domains: course plates from computing (C, C++, graphics, command line, compilers, CUDA, databases) and from mathematics and physics (calculus, classical mechanics, control). Held-out loss on 10% of each domain.

Where the extra tokens go

Raising science to 50% in the final 400 steps brings its held-out loss to 2.691 nats per character. The same number of science tokens spread evenly — 18% of every batch — gives 2.703, and keeping science at 10% gives 2.719. The late schedule wins by 0.012 nats on science, and costs 0.002 on computing. The effect is small here because a bigram model has little to learn, but its direction is the one larger models show: tokens seen during the decay count for more. At 90% late, the science loss is 2.681.

This is the logic of mid-training mixes. Llama 3’s report describes upsampling mathematics and code in its final stages, and OLMo 2’s second stage concentrates curated mathematics, code and reference text. The trade-off is the same as in plate IX.8, now in time as well as in proportion: what is upsampled at the end is learned well, at a small cost to what is down-weighted, and data too scarce to fill a whole run can still shape the final model.

Late upsampling also limits repetition. A scarce source read only during the final stage can be repeated a few times there without approaching the epochs at which repetition stops paying (plate X.6).

Terms

Upsampling: Giving a domain a larger share of training than its share of the data.
Late upsampling: Raising a domain’s share only in the final stage of training.
Bigram model: A model predicting each character from the previous one only.
Domain: A body of text with its own subject and style.

Extra tokens of a scarce domain buy more during the decay than spread through the run; mid-training mixes are built on that.

Sources

Llama Team, AI @ Meta, “The Llama 3 Herd of Models”, arXiv:2407.21783, 2024, sections 3.1.3 and 3.4.3.
Team OLMo, “2 OLMo 2 Furious”, arXiv:2501.00656, 2024.