exponential moving average of weights
In plain wordsA smoothed copy of a model kept during training, which each update nudges only slightly toward the latest weights. The calmer copy often performs better than the jittery original.
A copy of a network's parameters maintained alongside training and updated after each step as a weighted blend of its previous value and the current weights, with decay β close to one (0.999 or 0.9999 are common), so that it averages roughly the last 1/(1 − β) iterates. Because stochastic gradient descentstochastic gradient descent, n.: A faster version of gradient descent that estimates which way is downhill from a small random sample of the training data at each step, instead of checking the whole dataset. iterates bounce around a low-loss region, the average usually sits nearer its center and gives better, more stable results for evaluation and release, without altering training itself. EMA weights are standard in diffusion modelsdiffusion model, n.: An AI that learns to create images, sound or video by practicing noise removal: shown pictures buried under static, it learns to clean them up, then builds new ones starting from pure static. and GANsgenerative adversarial network, n.: Two neural networks trained as rivals: one forges fake images, the other tries to spot the fakes. As both improve, the forger's output becomes realistic enough to fool people., whose raw iterates are noisy, and in self-supervised methods (Mean Teacher, MoCo, BYOL, DINO) in which a slowly moving 'teacher' or 'momentum encoder' supplies targets, much like a target networktarget network, n.: A slowly updated copy of an AI agent's learning network, used to compute the goals it learns toward, so those goals do not shift every time the network changes. in reinforcement learning. Stochastic weight averaging instead takes an equal-weight average of checkpoints along a cyclic or constant learning-rate trajectory. Averaging is a theoretical counterpart of learning-rate decay (see schedule-free optimizationschedule-free optimization, n.: A way of training that needs no plan for slowing learning down over time, so training can be stopped or extended at any point and still produce a well-settled model.).
Sample quality jumped when we switched evaluation to the EMA weights with decay 0.9999.
Origin
Iterate averaging comes from stochastic approximation: David Ruppert (1988) and Boris Polyak with Anatoli Juditsky (1992) showed that averaging SGD iterates attains optimal convergence rates. Exponential averaging of weights spread in deep learning in the mid-2010s through image-classification and GAN practice; Pavel Izmailov and colleagues introduced stochastic weight averaging in 2018.
Usage note
Not to be confused with the moving averages inside optimizers such as momentummomentum, n.: A training trick that lets each update keep part of the previous update's direction, like a ball rolling downhill that builds up speed and coasts over small bumps. and AdamAdam, prop. n.: The most widely used method for training neural networks. It smooths the direction of updates over time and automatically sets a separate step size for each setting based on how large its recent changes have been., which smooth gradients rather than weights. Evaluating raw weights when EMA weights were intended is a common bug.
Related
- stochastic gradient descentstochastic gradient descent, n.: A faster version of gradient descent that estimates which way is downhill from a small random sample of the training data at each step, instead of checking the whole dataset.
- schedule-free optimizationschedule-free optimization, n.: A way of training that needs no plan for slowing learning down over time, so training can be stopped or extended at any point and still produce a well-settled model.
- diffusion modeldiffusion model, n.: An AI that learns to create images, sound or video by practicing noise removal: shown pictures buried under static, it learns to clean them up, then builds new ones starting from pure static.
- target networktarget network, n.: A slowly updated copy of an AI agent's learning network, used to compute the goals it learns toward, so those goals do not shift every time the network changes.
- momentummomentum, n.: A training trick that lets each update keep part of the previous update's direction, like a ball rolling downhill that builds up speed and coasts over small bumps.
- checkpointcheckpoint, n.: A saved snapshot of a model partway through or at the end of training, like a save point in a video game: training can resume from it after a crash, and finished snapshots are what people download and run.
Sources
- Polyak, B. T. & Juditsky, A. B. (1992). Acceleration of Stochastic Approximation by Averaging. SIAM Journal on Control and Optimization 30(4), 838–855.
- Izmailov, P. et al. (2018). Averaging Weights Leads to Wider Optima and Better Generalization. UAI 2018.
Mentioned in
momentummomentum, n.: A training trick that lets each update keep part of the previous update's direction, like a ball rolling downhill that builds up speed and coasts over small bumps., DINODINO, prop. n.: A family of Meta vision models that learn to understand images from millions of unlabeled photos, by training one network to agree with a slowly updated copy of itself., exponential smoothingexponential smoothing, n.: A simple forecasting method that predicts the next value as a weighted average of past values, giving the most weight to recent ones and steadily less to older ones., mode connectivitymode connectivity, n.: The discovery that two separately trained versions of a network, though their settings look very different, can usually be joined by a path along which the model keeps working well., model soupmodel soup, n.: A single model made by averaging the internal numbers of several copies of a model that were fine-tuned in slightly different ways, often doing better than any one of them., schedule-free optimizationschedule-free optimization, n.: A way of training that needs no plan for slowing learning down over time, so training can be stopped or extended at any point and still produce a well-settled model.