Vol. XII · Mid-Training Plate 2 of 10annealing, learning-rate decay, high-quality data, data mix, mathematics, Llama 3, targeted improvementLD

Annealing on Good Data.

The last stretch of the decay is where the model settles — so choose what it settles on.

Plate XI.1 showed that the loss drops sharply while the learning rate decays: the weights stop bouncing and settle. Which minimum they settle into depends on the data seen during the decay. If that data is enriched with high-quality examples of a skill, the model settles closer to what that skill needs. Llama 3’s team reported that annealing a pretrained 8B model on a small amount of high-quality code and mathematics raised its GSM8k score by 24.0% and its MATH score by 6.4%; for the 405B model the gain was negligible, as the larger model already did well.

The noisy quadratic model of Vol. XI can show the mechanism. Give it two objectives that share the same weights — a general one and a “quality” one whose optimum lies a little away — and train mostly on the general objective at a constant learning rate. Then decay the learning rate over the final fifth, either on the same mix or on a mix with more of the quality objective. The workshop compares the two.

Below: loss on the general and the quality objectives through training, for the chosen share of quality data in the decay phase, with the stable phase’s 10% for comparison.

What you will learn
Workshop
Noisy quadratic model with two objectives sharing 20 weights; the quality objective’s optimum is displaced. Momentum SGD, 3 000 steps, learning rate 0.05; stable phase draws the quality objective 10% of the time; the final 20% decays the learning rate linearly to zero with the chosen quality share. Lines: general loss (blue) and quality loss (red); dashed, annealing at 10%.

Choosing where the model settles

Until the decay begins the two lines are identical, because the data have not changed. During the decay, a 50% quality share brings the quality loss to 3.65 × 10⁻², 2.8 times lower than annealing on the stable 10% mix (1.01 × 10⁻¹); the general loss ends 3.3 times higher. With 30% quality the gain is 1.4 times; with the whole decay on quality data, 8.4 times, at a general cost of 9.5 times. The same compute, spent in the same final fifth, gives very different models depending on what they read while settling.

Published results follow the same pattern. Llama 3’s annealing on a small set of high-quality code and mathematics improved the 8B model’s GSM8k score by 24% and MATH by 6.4%, while the 405B model gained little: a larger model has less left to settle. Labs also use annealing as a cheap probe of data quality — anneal a checkpoint on a candidate dataset and see how much a benchmark moves — which lets a new source be judged without training a model from scratch.

Because annealing data has such leverage, it is also where benchmark contamination is most dangerous: a few thousand leaked test items in the annealing mix can move a score far more than the same items spread through pretraining (plate IX.10).

Terms

Annealing
Decaying the learning rate to zero at the end of training.
Annealing mix
The data used during the decay, often enriched with high-quality or targeted text.
Targeted improvement
A gain on specific skills from data chosen for them.
Annealing probe
Using a short anneal on candidate data to judge its quality.
Data seen while the learning rate decays decides where the model settles; annealing on high-quality data buys large targeted gains cheaply.

Sources

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