Weight Averaging, EMA, SWA, and Checkpoint Ensembling

They differ in storage, inference cost, and assumptions about parameter space.

The third card is not hypothetical. Take a zero-shot CLIP model and the model it became after fine-tuning, and blend their weights. That is WiSE-FT, published in 2022 by Wortsman and colleagues. One sentence of their abstract carries the whole method: “We address this tension by introducing a simple and effective method for improving robustness while fine-tuning: ensembling the weights of the zero-shot and fine-tuned models (WiSE-FT).” On ImageNet and five derived distribution shifts, the blend beat prior work by 4 to 6 percentage points under shift. ImageNet accuracy rose 1.6 pp. Six further shifts gained 2 to 23 pp. Fine-tuning and inference cost nothing extra.

What makes that arithmetic legal is the lineage, not the operation. The two endpoints are one model before and after a single fine-tuning run, so their coordinates already agree. Every other card on this page is a variation on whether that condition holds. The rest of the lesson is about how to check it rather than assume it.

An exponential moving average updates a shadow parameter vector after successful optimizer steps. A high decay coefficient changes slowly and emphasizes a longer history.

Averaging the weights, rather than the predictions, is what buys the gain. Mean Teacher is the canonical demonstration, at NIPS 2017. Tarvainen and Valpola state the design and the result together in their abstract: “To overcome this problem, we propose Mean Teacher, a method that averages model weights instead of label predictions. As an additional benefit, Mean Teacher improves test accuracy and enables training with fewer labels than Temporal Ensembling. Without changing the network architecture, Mean Teacher achieves an error rate of 4.35% on SVHN with 250 labels, outperforming Temporal Ensembling trained with 1000 labels.” The weight average beat the prediction average, and the prediction average had four times as many labels. With Residual Networks the same method improved CIFAR-10 with 4000 labels from 10.55% to 6.28%, and ImageNet 2012 with 10% of the labels from 35.24% to 9.11%.

The decay coefficient is a hyperparameter, not a default to inherit, and it is a nasty one. Karras and colleagues at NVIDIA open their 2024 paper with the practical problem: “Unfortunately, the EMA decay constant is a cumbersome hyperparameter to tune because the effects of small changes become apparent only when the training is nearly converged.” They then priced getting it wrong. Sweeping the EMA length of a single weight tensor, holding every other tensor at the global optimum, moved FID by as much as 10%, from 7.24 to about 6.5, in their CONFIG B. One tensor's memory horizon, and the headline metric moves. Their engineering answer is post-hoc EMA: store two averaged parameter vectors during the run, γ1 = 16.97 and γ2 = 6.94 — relative widths σrel of 0.05 and 0.10 — snapshotted once every 4096 training steps at batch size 2048, then reconstruct any EMA profile after training instead of guessing it before. The full paper improved the ImageNet-512 record FID from 2.41 to 1.81.

EMA can stabilize evaluation, especially under noisy updates, but it lags rapid adaptation. The update cadence must align with skipped steps, accumulation, and distributed synchronization. PyTorch pins that clock to optimizer updates explicitly: “For SWA and EMA, this call is usually done right after the optimizer step().”

Parameter arithmetic assumes compatible coordinates and functions.

The permutation problem has since been measured, and partly solved. The Git Re-Basin algorithms, published in 2022, permute one model into the coordinates of a reference model. That produced “zero-barrier linear mode connectivity between independently trained ResNet models on CIFAR-10”. Averaging two independently trained networks needs that alignment done first. The same paper reports a counterexample to its own single-basin claim.

The barrier itself has a measurement and a published threshold. An ICML 2020 paper built instability analysis around a single quantity, defining “the error barrier height of p as the maximum increase in error from that of W1 and W2 along path p”, evaluated along the straight line between two networks. Frankle and colleagues found that the barrier is a property of when the two runs diverged, not of how good either endpoint is: “All but the smallest MNIST network are unstable to SGD noise at initialization according to linear interpolation. However, by a point early in training (3% for ResNet-20 on CIFAR-10 and 20% for ResNet-50 on ImageNet), all networks become stable to SGD noise.” Three percent of a ResNet-20 run on CIFAR-10 is the entire shared prefix it took. That gives the practice a test with a number in it. Before averaging two checkpoints, ask not whether both score well, but whether they share a training prefix. If the run history cannot answer that, measure the error along the line between them.

SWA averages selected states rather than every early checkpoint.

The recalibration step is part of the method, not housekeeping bolted on afterwards. The SWA paper builds it in: “we run one additional pass over the data, as in Garipov et al. [2018], to compute the running mean and standard deviation of the activations for each layer of the network”. The reason is stated in the same sentence. Those statistics are not collected during training, because no training step ever ran with the averaged weights.

PyTorch ships that pass as its own API rather than folding it into the averaging call. “torch.optim.swa_utils.update_bn() is a utility function used to update SWA/EMA batch normalization statistics at the end of training”, and “update_bn() applies the swa_model to every element in the dataloader and computes the activation statistics for each batch normalization layer in the model”. It costs a full pass over the data and it must be invoked explicitly. Skip it and you deploy a model whose weights are an average and whose normalization buffers belong to one arbitrary point of the trajectory.

Separate trajectory smoothing from ensemble diversity.

Two of these steps are usually underestimated. Step 3 is a data pass, not a flag: update_bn() runs the averaged model over the dataloader before any evaluation number is believable. Step 4 has a fixed candidate list because the five objects are genuinely different, and their published results do not transfer between them. WiSE-FT's 1.6 pp on ImageNet is an interpolation result. Mean Teacher's 4.35% error on SVHN with 250 labels is a weight-averaging result. The deep-ensemble advantage that survived dataset shift at NeurIPS 2019 is a prediction-ensemble result about calibration rather than accuracy. Quoting one of them to justify another is the most common way this comparison goes wrong.

Step 5 then prices what step 4 chose. An ensemble of about five members — the size the shift benchmark found sufficient — costs five forward passes at every request, forever. An average costs one.

Averaging poor or incompatible states does not create a strong model automatically. The candidate set still needs a coherent training phase and validation rationale. Keep a simple strong checkpoint as a control. If averaging helps only after broad search, include that selection cost in the result.

Stochastic weight averaging was published at UAI in 2018. Averaging points along the SGD trajectory “leads to better generalization than conventional training”, its five authors write, and they describe SWA as “extremely easy to implement” with “almost no computational overhead”. Wortsman and colleagues carried the idea into fine-tuning at ICML in 2022. Their model soups averaged weights “without incurring any additional inference or memory costs”. The resulting ViT-G reached “90.94% top-1 accuracy on ImageNet”.

Notice what both claims are about: cost, measured against a trained model that already existed. Neither says averaging supplies quality that the candidate states did not have. The candidates in the SWA result come from one trajectory under a chosen schedule. The endpoints in the WiSE-FT result are a model and its own fine-tuned descendant. Reproduce the cheapness without reproducing the compatibility and there is nothing left of the finding.