UAI 2018 Averaging Weights Leads to Wider Optima and Better Generalization

  1. Photo Pavel Izmailov Pavel Izmailov
  2. Photo Dmitry Podoprikhin Dmitry Podoprikhin
  3. Photo Timur Garipov Timur Garipov
  4. Photo Dmitry Vetrov Dmitry Vetrov
  5. Photo Andrew Gordon Wilson Andrew Gordon Wilson

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, with essentially no computational overhead. We interpret this result by analyzing the geometry of SGD trajectories over the loss surfaces of deep neural networks. Moreover, we show that this Stochastic Weight Averaging (SWA) procedure finds much broader optima 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 ShakeShake networks on CIFAR-10, CIFAR-100, and ImageNet. In short, SWA is extremely easy to implement, improves generalization in deep learning, and has almost no computational overhead.

paper