# Stochastic Weight Averaging

Stochastic Weight Averaging (SWA) is a deep learning technique that averages model weights along the training trajectory to improve generalization, often yielding flatter minima and better performance than conventional training.

Stochastic Weight Averaging (SWA) is a model averaging technique in [deep-learning](https://www.wikiprompt.org/wiki/deep-learning) that improves generalization by averaging the weights of a [neural-network](https://www.wikiprompt.org/wiki/neural-network) over multiple points along its training trajectory. Unlike traditional training that retains only the final weights, SWA collects weights at regular intervals during the later phase of training and computes their arithmetic mean. This simple procedure often leads to flatter solutions in the loss landscape, which are associated with better test performance and robustness. SWA was introduced by Pavel Izmailov, Dmitrii Podoprikhin, Timur Garipov, Dmitry Vetrov, and Andrew Gordon Wilson in their 2018 paper "Averaging Weights Leads to Wider Optima and Better Generalization." The method requires no changes to the underlying model architecture or loss function, making it easy to integrate into existing training pipelines.

The core idea behind SWA is that during training with stochastic gradient descent (SGD) or its variants, the model parameters oscillate around a region of low loss. Averaging these parameters can smooth out the oscillations and yield a solution that lies in a wide, flat minimum. Such flat minima are hypothesized to generalize better than sharp minima because they are less sensitive to perturbations and shifts between training and test distributions. SWA has been shown to improve accuracy and calibration across various architectures and tasks, including image classification, object detection, and language modeling.

## Mechanism and Algorithm

SWA operates in two phases. In the first phase, the model is trained with a standard learning rate schedule (e.g., cosine annealing or step decay) to reach a region near a minimum. In the second phase, training continues with a constant or cyclical learning rate, and the weights are collected at the end of each epoch or after a fixed number of iterations. The SWA weights are updated as a running average: n_swa = n_swa + 1; w_swa = (w_swa * (n_swa - 1) + w) / n_swa, where w is the current weights and w_swa is the averaged weights. The final model uses w_swa.

A key detail is that batch normalization statistics must be recomputed after averaging, because the running mean and variance of batch normalization layers are not averaged along with the weights. In practice, after the SWA weights are computed, a forward pass over the training data is performed to update the batch normalization statistics. This step is crucial for models that use batch normalization, such as [ResNets](https://www.wikiprompt.org/wiki/residual-network).

## Theoretical Motivation

The success of SWA is often explained through the lens of loss landscape geometry. Research by Izmailov et al. and subsequent work has shown that SGD tends to converge to points near the boundary of a low-loss region, whereas averaging over multiple points moves the solution toward the center of that region. This central point typically lies in a flatter basin, which is associated with better generalization. The connection between flat minima and generalization has been studied extensively, with connections to PAC-Bayes bounds and the sharpness of the loss function.

Another perspective relates SWA to Bayesian inference. Under certain assumptions, the trajectory of SGD can be viewed as a Markov chain that samples from a posterior distribution. Averaging along this trajectory approximates the posterior mean, which is the Bayes-optimal prediction for squared error loss. This interpretation links SWA to stochastic gradient Markov chain Monte Carlo methods, though SWA is simpler and does not require careful tuning of noise scales.

## Variants and Extensions

Several variants of SWA have been proposed to further improve performance. One notable extension is Stochastic Weight Averaging Gaussian (SWAG), introduced by Wesley Maddox, Timur Garipov, Pavel Izmailov, Dmitry Vetrov, and Andrew Gordon Wilson in 2019. SWAG approximates the posterior distribution of the weights by fitting a Gaussian using the first and second moments of the collected weights. This allows for uncertainty estimation and improved calibration, and it can be used for Bayesian deep learning.

Another variant is Fast Geometric Ensembling (FGE), which uses a cyclical learning rate schedule to collect weights from different modes of the loss landscape. SWA can be combined with FGE to achieve even better performance. Additionally, SWA has been integrated with techniques like [learning rate schedules](https://www.wikiprompt.org/wiki/learning-rate-schedule), [data augmentation](https://www.wikiprompt.org/wiki/data-augmentation), and [pruning](https://www.wikiprompt.org/wiki/model-pruning) to enhance robustness and efficiency.

## Applications

SWA has been applied successfully across a wide range of domains. In computer vision, it improves the accuracy of image classifiers on datasets like CIFAR-10, CIFAR-100, and ImageNet. For example, SWA has been shown to boost the test accuracy of a PreAct ResNet-164 on CIFAR-100 from 81.15% to 82.90% without any architectural changes. In object detection, SWA helps stabilize training and improve mean average precision.

In natural language processing, SWA has been used to train [Transformers](https://www.wikiprompt.org/wiki/transformer) and [large language models](https://www.wikiprompt.org/wiki/large-language-model) more effectively. It has been shown to reduce overfitting and improve perplexity in language modeling tasks. SWA also benefits [generative models](https://www.wikiprompt.org/wiki/generative-ai) such as GANs and diffusion models by providing smoother weight updates, which can lead to more stable training and higher quality samples.

Beyond standard supervised learning, SWA has been adapted for semi-supervised learning, transfer learning, and domain adaptation. Its simplicity and generality make it a valuable tool in the [machine-learning](https://www.wikiprompt.org/wiki/machine-learning) practitioner's toolkit.

## Comparison with Other Techniques

SWA is often compared with other regularization and ensembling methods. Unlike traditional ensembles that average predictions from multiple independently trained models, SWA averages weights, which is computationally cheaper and requires only a single training run. This makes it particularly attractive when model storage or inference cost is a concern.

Compared to techniques like [dropout](https://www.wikiprompt.org/wiki/dropout) or [batch normalization](https://www.wikiprompt.org/wiki/batch-normalization), SWA does not modify the training dynamics but rather post-processes the weights. It can be used in conjunction with these methods. SWA also differs from [Adam](https://www.wikiprompt.org/wiki/adam-optimizer) and other adaptive optimizers, which maintain per-parameter learning rates; SWA is typically applied on top of SGD or Adam trajectories.

One limitation of SWA is that it requires a schedule to determine when to start collecting weights. If the collection starts too early, the average may be dominated by suboptimal weights; if too late, the benefits diminish. However, in practice, a simple rule of thumb is to start collecting after the learning rate has decayed to a low value.

## Practical Considerations

Implementing SWA is straightforward. Most deep learning frameworks, such as PyTorch and TensorFlow, provide utilities or examples for SWA. The main steps are: (1) train the model with a standard schedule, (2) continue training with a constant or cyclical learning rate, (3) collect weights at intervals, (4) compute the running average, and (5) recompute batch normalization statistics.

SWA has been shown to be robust to hyperparameter choices. The learning rate for the second phase is typically set to the minimum value of the initial schedule or a small constant. The frequency of weight collection can be every epoch or every few hundred iterations. In practice, collecting weights for 5 to 20 epochs is often sufficient.

## Impact and Reception

The introduction of SWA has had a significant impact on the deep learning community. It provided a simple yet effective method to improve generalization without additional computational cost during inference. The paper has been widely cited and has inspired numerous follow-up works on weight averaging and flat minima. SWA is now a standard technique in many training pipelines, especially in academic research and competitions.

SWA has also influenced the development of other averaging methods, such as Exponential Moving Average (EMA), which is commonly used in training [generative models](https://www.wikiprompt.org/wiki/generative-ai) and [large language models](https://www.wikiprompt.org/wiki/large-language-model). EMA is a variant that gives more weight to recent weights, and it is often used in practice due to its simplicity and effectiveness. SWA and EMA are complementary, and some frameworks allow switching between them.

## Limitations and Future Directions

Despite its benefits, SWA has limitations. It assumes that the loss landscape is relatively smooth and that averaging weights is meaningful, which may not hold for all architectures or tasks. For example, models with batch normalization require careful handling, and models with complex parameter spaces (e.g., recurrent networks) may not benefit as much. Additionally, SWA does not provide uncertainty estimates unless extended (e.g., SWAG).

Future research directions include developing adaptive averaging schedules, combining SWA with Bayesian methods, and applying SWA to new domains such as reinforcement learning and federated learning. As deep learning continues to evolve, weight averaging techniques like SWA remain a fundamental tool for improving model performance and reliability.

---
Source: https://www.wikiprompt.org/wiki/swa-stochastic-weight-averaging
License: CC BY-SA 4.0 (https://creativecommons.org/licenses/by-sa/4.0/)
Last updated: 2026-09-13T03:59:32.079754+00:00