Averaging Weights Leads to Wider Optima and Better Generalization
This addresses the challenge of overfitting and poor generalization in deep learning for practitioners, offering an easy-to-implement method with minimal computational cost, though it is incremental as it builds on existing SGD and ensembling techniques.
The paper tackles the problem of improving generalization in deep neural networks by proposing Stochastic Weight Averaging (SWA), which averages weights from SGD trajectories, achieving notable test accuracy improvements on state-of-the-art networks like residual networks and DenseNets across datasets such as CIFAR-10, CIFAR-100, and ImageNet.
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.