Vol. XII · Mid-Training Plate 8 of 10checkpoint averaging, exponential moving average, model soups, model merging, weight averaging, noiseLD

Averaging and Merging.

Average the weights of a few recent checkpoints and the noise in each cancels — a better model for the cost of an addition.

At a constant or slowly decaying learning rate the weights wander around the minimum, pushed by each batch’s noise (plate XI.1). Each checkpoint is a noisy sample of where the weights could be; their average sits closer to the centre. Checkpoint averaging — or an exponential moving average (EMA) of the weights kept alongside training — exploits this, and Llama 3’s final pretrained model was an average of checkpoints taken during its annealing stage.

The same arithmetic works across runs if the models share a starting point and stay in the same basin. Wortsman and colleagues’ model soups (2022) averaged models fine-tuned from one checkpoint with different hyperparameters and outperformed the best single model; OLMo 2 averaged several annealing runs on different data orders. Merging models fine-tuned for different skills is a further step, with more ways to fail. The workshop measures averaging on the noisy quadratic model.

Below: left, the loss of the current weights and of their moving average through training at a constant learning rate. Right, the final loss of the last checkpoint, the average of the last K checkpoints, and the moving average.

What you will learn
Workshop
Noisy quadratic model, momentum SGD at constant learning rate 0.05 for 3 000 steps. Checkpoints every 10 steps; the window average is over the last K of them; the EMA has the chosen decay β. Left: loss of the current weights (grey) and of the EMA (blue), log scale. Right: final losses.

Noise cancels, signal stays

At a constant learning rate the last checkpoint ends at a loss of 8.78 × 10⁻¹: it is one noisy point among many. Averaging the last 20 checkpoints, spaced 10 steps apart, gives 3.51 × 10⁻², 25.0 times lower; averaging 5 gives 1.45 × 10⁻¹, and averaging 100 gives 1.31 × 10⁻². The window should be long enough for the noise to cancel and short enough that the weights have not moved on meanwhile. An EMA with β = 0.99 gets 4.26 × 10⁻² without storing checkpoints at all, which is why many training codes keep one.

Averaging does the job the learning-rate decay does — removing noise — without the decay, so it combines naturally with WSD schedules (plate XI.1) and is cheap enough to try on any run. Across runs the condition is linear connectivity: models fine-tuned from the same checkpoint usually lie in one basin where their average is also good, which is what model soups exploit. Models trained from different initialisations do not, and averaging them gives nonsense; merging models fine-tuned for different skills sits between, working when the changes are small and compatible and failing when they interfere.

Merging specialised models by adding their weight differences from a common base — task arithmetic — has become popular in open-model communities. It is cheap and sometimes strikingly effective, but it is not training: nothing checks that the combined model has kept either skill.

Terms

Checkpoint averaging
Averaging the weights of several checkpoints of one run.
Exponential moving average (EMA)
A running weighted average of the weights, updated every step.
Model soup
An average of models fine-tuned from one starting point with different settings.
Linear mode connectivity
When the straight line between two models stays at low loss, so averaging them works.
Averaging checkpoints cancels the noise each one carries; averaging across runs works when the models share a starting point and a basin.

Sources

  1. Izmailov et al., “Averaging Weights Leads to Wider Optima and Better Generalization”, UAI 2018.
  2. Wortsman et al., “Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time”, ICML 2022.
  3. Llama Team, AI @ Meta, arXiv:2407.21783, 2024, section 3.4.3.