Source-linked AI summary
Averaging Weights Leads to Wider Optima and Better Generalization
Pavel Izmailov, Dmitrii Podoprikhin, Timur Garipov, Dmitry Vetrov, Andrew Gordon Wilson
TL;DR
Deep neural-network training commonly relies on SGD with a decaying learning rate, while the geometry of useful solutions and their relation to generalization remains important. The paper introduces SWA, which averages weights along SGD trajectories under cyclical or constant learning rates, and reports flatter solutions and improved performance across benchmarks with nearly no overhead.
Problem
The paper addresses how loss-surface geometry and the choice of points reached during SGD relate to generalization in deep neural networks.
Method
SWA averages multiple weights collected along SGD trajectories using cyclical or constant learning rates, producing one model rather than averaging ensemble predictions.
Results
SWA improves generalization across architectures and benchmarks, finds wider solutions than SGD, and approximates FGE with a single model.
Takeaways & Limitations
SWA is architecture-agnostic, easy to implement, and improves generalization at virtually no additional cost over conventional training.
Takeaways & Limitations
For practical use, the authors recommend initializing SWA from a conventionally pretrained model because this gives faster and more stable convergence than starting SWA from scratch.
Abstract
from arXiv · showhide
Deep neural networks are typically trained by optimizing a loss function with an SGD variant, in conjunction with a decaying learning rate, until convergence. We show that simple averaging of multiple points along the trajectory of SGD, with a cyclical or constant learning rate, leads to better generalization than conventional training. We also show that this Stochastic Weight Averaging (SWA) procedure finds much flatter solutions than SGD, and approximates the recent Fast Geometric Ensembling (FGE) approach with a single model. Using SWA we achieve notable improvement in test accuracy over conventional SGD training on a range of state-of-the-art residual networks, PyramidNets, DenseNets, and Shake-Shake networks on CIFAR-10, CIFAR-100, and ImageNet. In short, SWA is extremely easy to implement, improves generalization, and has almost no computational overhead.
1 INTRODUCTION
The paper introduces Stochastic Weight Averaging (SWA), which averages weights visited by SGD under cyclical or constant learning rates. SWA finds wider, flatter solutions, approximates FGE with one model, and improves test performance across architectures and benchmarks with negligible overhead.
- Core method: SWA averages weights proposed over SGD iterations while using cyclical or constant learning rates to reach more desirable regions of weight space.The procedure is presented as a drop-in replacement for standard training.
- Connection to FGE: SWA approximates Fast Geometric Ensembling while retaining the test-time computation, convenience, and interpretability of a single model.FGE ensembles require k times more test-time computation for k models.
- Geometry and generalization: SWA finds wider, more centered solutions than SGD, often with slightly worse train loss but substantially better test error.SGD tends to remain near the boundary of a wide flat region, whereas SWA moves toward its center.
- Geometry and generalization: The loss function is asymmetric between SWA and SGD, with SGD near sharp ascent and SWA in flatter regions along that direction.The paper identifies this geometry as part of why SWA improves generalization.
- Empirical results: SWA improves test accuracy across residual, PyramidNet, DenseNet, and Shake-Shake architectures on CIFAR-10, CIFAR-100, and ImageNet.Reported gains include 0.8% for ResNet-50 and DenseNet-161, 0.6% for ResNet-150, over 1.3% on CIFAR-100, and over 0.4% on CIFAR-10.
- Practicality: SWA is extremely easy to implement and has virtually no computational overhead compared with conventional training.The paper provides an implementation and describes the method as requiring practically the same computation as SGD.
2 RELATED WORK
Related work connects SWA to loss-surface geometry, stochastic optimization, FGE, and regularization. SWA extends weight averaging and trajectory exploration into a neural-network training procedure that can approximate ensembles with one model.
- Loss-surface geometry: The paper frames SWA within efforts to understand how loss-surface geometry and solution width relate to deep-learning generalization.Prior work offers competing views on whether broad or sharp optima explain generalization.
- Optimization: SWA builds on running averages of SGD weights, extending ideas from convex optimization that are not typically used to train neural networks.Exponentially decaying averages with decaying learning rates are described as smoothing SGD while performing comparably.
- Optimization: Constant-learning-rate SGD can be interpreted as sampling around a loss minimum, motivating SWA's exploration of high-performing weight regions.The cited analysis links the sampling covariance to the learning rate under simplifying assumptions.
- Fast Geometric Ensembling: FGE uses cyclical learning rates to collect nearby models with diverse predictions and ensemble them without extra training-time cost.SWA follows these trajectories to construct a single model approximating an FGE ensemble.
- Regularization: SWA and dropout are both training procedures and regularizers motivated by ensemble approximation, and the paper reports that they can be combined for improved performance.Their implementations differ substantially: dropout varies architectures across minibatches, whereas SWA averages weights along an SGD trajectory.
3 STOCHASTIC WEIGHT AVERAGING
Stochastic Weight Averaging (SWA) averages SGD iterates obtained with cyclical or constant learning rates to reach central, flatter regions of high-performing solutions. It improves generalization, approximates FGE with one model, and adds negligible computation.
- SWA procedure: SWA averages SGD proposals while cyclical or constant learning-rate schedules explore high-performing regions of weight space.The procedure continues training from a pretrained model and averages weights along the resulting trajectory.
- Trajectory geometry: SGD trajectories explore near the periphery of high-performing networks, while averaging produces a more central point with higher test performance.Train-loss and test-error surfaces are shifted, so the train-loss minimizer need not be optimal on test.
- Solution width: SWA converges to a wider solution than SGD: its train-loss and test-error curves remain broad when moving along random directions.On CIFAR-100, SGD's converged test error is lower by 1.5% than the SWA solution in the referenced comparison, while SWA's curves are considerably wider.
- Solution width: SWA finds a flatter region in the same basin as SGD, including along the direction connecting the two solutions, where SGD can lie near steep ascent.The loss is asymmetric in this direction, and SWA's flatter location is associated with better generalization.
- Connection to ensembling: Averaging weights approximates the FGE ensemble with a single model because the difference between averaged predictions and averaged weights is second-order for nearby proposals.The SWA and FGE probabilities have norm difference 0.079, and identically labeled objects comprise 95.26% of cases.
4 EXPERIMENTS
Experiments compare SWA with conventional SGD and FGE across CIFAR and ImageNet architectures, while also testing learning-rate schedules and fixed-rate training. SWA consistently improves performance, can match FGE with one model, and benefits from pretrained initialization for faster, more stable convergence.
- SWA substantially outperforms SGD within one training budget across CIFAR experiments and improves further with additional epochs.
- 0.5% improvement on CIFAR-10 and 0.75-1.5% improvement on CIFAR-100 are consistent across architectures, excluding Shake-Shake on CIFAR-10.
- SWA achieves comparable or better performance than FGE ensembles using a single model, although CIFAR-100 usually requires more than one budget.
- 0.6-0.9% improvement over pretrained models is achieved by SWA for all three ImageNet architectures.
- More aggressive constant learning rates generally accelerate SWA convergence, while intermediate rates often perform best.
- 81.7 final test accuracy is achieved by fixed-rate SWA on Wide ResNet-28-10 trained from scratch on CIFAR-100.
- Pretraining SWA with conventional training is recommended because it yields faster and more stable convergence than starting SWA from scratch.
5 DISCUSSION
The discussion presents SWA as a simple, architecture-agnostic training modification that improves generalization with virtually no additional cost. It also frames broader optima and loss-surface geometry as directions for further research.
- SWA is architecture-agnostic and improves generalization performance at virtually no additional cost over conventional training.
- SWA may support faster convergence and combination with large batch sizes because it discovers broader optima than conventional SGD.
- The paper connects SWA to better understanding of loss-surface geometry and generalization in deep learning.
A.1 EXPERIMENTAL DETAILS
The experiments use model-specific SWA learning-rate schedules, including constant and cyclic schedules, across CIFAR and ImageNet settings. Conventional SGD uses momentum with annealed schedules and architecture-specific alternatives.
- Implementations: The CIFAR experiments use implementations including Preactivation-ResNet-164, while ImageNet models and pretrained networks come from specified external sources.
- SWA learning rates: SWA uses cyclic learning rates for PyramidNet and Shake-Shake Net, while VGG, Wide ResNet, and ResNet use constant learning rates.PyramidNet uses α1 = 0.05 and α2 = 0.001 with cycle length 3; VGG and Wide ResNet use α1 = 0.01, while ResNet uses 0.01 on CIFAR-10 and 0.05 on CIFAR-100.
- SWA learning rates: Shake-Shake cycles reproduce the learning rates from standard-training epochs 1600–1700, using cycle length c = 100 epochs.
- SWA learning rates: For ImageNet experiments, SWA uses a cyclic learning-rate schedule with α1 = 0.001, α2 = 10^-5, and cycle length c = 1.
- SGD learning rates: Conventional SGD uses momentum 0.9 with annealed learning rates, including a fixed first half, linear decay over the next 40% of epochs, and a constant final 10%.VGG uses α1 = 0.05, while Preactivation ResNet and Wide ResNet use α1 = 0.1; Shake-Shake Net and PyramidNets use architecture-specific schedules.
A.2 TRAINING RESNET WITH A CONSTANT LEARNING RATE
The experiment trains Preactivation ResNet-164 with a constant learning rate and begins averaging after epoch 200. Figure 9 compares test-error trajectories under constant and decaying schedules, including averaged SGD points.
- Experimental setup: The Preactivation ResNet-164 experiment sets the constant learning rate to α1 = 0.1 and starts averaging after epoch 200.
- Results visualization: Figure 9 plots test error across training epochs for constant and decaying learning-rate schedules, with averaged constant-rate SGD points shown from epoch 200.