stochastic weight averaging (SWA)

Late in training, SGD run with a cyclical or high constant learning rate does not converge to a single point; it bounces around the rim of a wide, flat basin, visiting many good but slightly different weight vectors. Stochastic weight averaging makes a simple bet: if all those points sit on the rim, their average should sit nearer the center of the basin, and centers of wide flat basins tend to generalize better than any individual rim point.

Concretely, after a normal training phase you switch to a high constant or cyclical learning rate and record the weights at the end of each cycle, then take their running average theta_SWA = (1/n) sum theta_i. One subtlety: batch-normalization statistics are not parameters being averaged, so after forming the averaged weights you must do a single extra forward pass over the data to recompute the activation statistics for the averaged model.

SWA gives a nearly free boost to generalization on top of ordinary SGD, and it connects directly to the flat-minima reasoning behind sharpness-aware methods. Its Bayesian extension, SWAG, fits a Gaussian over the collected weights to estimate uncertainty. Crucially the end product is a single set of weights, so unlike a true ensemble it adds no extra cost at inference time.