Averaging Weights Leads to Wider Optima and Better Generalization

Pavel Izmailov∗1

Dmitrii Podoprikhin∗2,3 Timur Garipov∗2,3 Dmitry Vetrov3 Andrew Gordon Wilson1
Cornell University, 2 Moscow State University, 3 Higher School of Economics

arXiv:1803.05407v1 [cs.LG] 14 Mar 2018

1

Abstract
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 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,
and has almost no computational overhead.

1

Introduction

With a better understanding of the loss surfaces for multilayer networks, we can accelerate the convergence, stability, and accuracy of training procedures in deep learning. Recent work [Garipov et al., 2018, Draxler et al.,
2018] shows that local optima found by SGD can be connected by simple curves of near constant loss. Building
upon this insight, Garipov et al. [2018] also developed
Fast Geometric Ensembling (FGE) to sample multiple
nearby points in weight space to create high performing
ensembles in the time required to train a single DNN.
FGE uses a high frequency cyclical learning rate with
SGD to select networks to ensemble. In Figure 1 (left)
∗

Equal contribution.

we see that the weights of the networks ensembled by
FGE are on the periphery of the most desirable solutions.
This observation suggests it is promising to average these
points in weight space, rather than model space. Although the general idea of maintaining a running average of weights traversed by SGD dates back to Ruppert
[1988], this procedure is not typically used to train neural
networks. It is sometimes applied as an exponentially decaying running average in combination with a decaying
learning rate (where it is called an exponential moving
average), which smooths the trajectory of conventional
SGD but does not perform very differently. However, we
show that an equally weighted average of the points traversed by SGD with a cyclical or constant learning rate,
which we refer to as Stochastic Weight Averaging (SWA),
has many surprising and promising features for training
deep neural networks, leading to a better understanding
of the geometry of their loss surfaces. Indeed, SWA with
cyclical or constant learning rates can be used as a dropin replacement for standard SGD training of multilayer
networks — but with improved generalization and essentially no overhead. In particular:
• We show that SGD with cyclic [e.g., Loshchilov
and Hutter, 2016] and constant learning rates traverses regions of weight space corresponding to
high-performing networks. We find that while these
models are moving around this optimal set they
never reach its central points. We show that we can
move into this more desirable space of points by averaging the weights proposed over SGD iterations.
• While FGE ensembles [Garipov et al., 2018] can
be trained in the same time as a single model, test
predictions for an ensemble of k models requires k
times more computation. We show that SWA can
be interpreted as an approximation to FGE ensembles but with the test-time, convenience, and interpretability of a single model.
• We demonstrate that SWA leads to solutions inside

Test error (%)

Test error (%)
30

W2

20

WSWA

10

> 50

> 50

50

50
35.11

28.49

27.52
5

23.65

22.38

W3

0

21.24

WSWA

0

20.67

0

10

20

30

40

50

0.4391
0.2206
5

0.1131
0.06024

WSWA

0

0.03422

epoch 125

epoch 125

0.02142

20.15

19.95
−10

0.8832

21.67

20.64
−10

> 0.8832

WSGD

10

35.97

24.5

W1

WSGD

10

Train loss

19.62
−5

0

5

10

15

20

25

0.00903
−5

0

5

10

15

20

25

Figure 1: Illustrations of SWA and SGD with a Preactivation ResNet-110 on CIFAR-100. Left: test error surface
for three FGE samples and the corresponding SWA solution (averaging in weight space). Middle and Right: test
error and train loss surfaces showing the weights proposed by SGD (at convergence) and SWA, starting from the same
initialization of SGD after 125 training epochs.
substantially wider optima than SGD. Keskar et al.
[2017] conjectures that the width of the optima is
critically related to generalization. We illustrate that
the loss on the train is shifted with respect to the
test error (Figure 1, middle and right panels, and
sections 3, 4). We show that SGD generally converges to a point near the boundary of the wide flat
region of optimal points. SWA on the other hand
is able to find a point centered in this region, often
with slightly worse train loss but with substantially
better test error.
• SWA achieves notable improvement for training
a broad range of architectures over several consequential benchmarks. In particular, running
SWA for just 10 epochs on ImageNet we are able
to achieve approximately 0.8% improvement for
ResNet-50 and DenseNet-161, and 0.6% improvement for ResNet-150. We achieve improvement
of over 1.3% on CIFAR-100 and of over 0.4% on
CIFAR-10 with Preactivation ResNet-110, VGG-16
and Wide ResNet-28-10. We also achieve substantial improvement for the recent Shake-Shake Networks and PyramidNets.
• SWA is extremely easy to implement and has virtually no computational overhead compared to the
conventional training schemes.

2

Related Work

This paper is fundamentally about better understanding
the geometry of loss surfaces and generalization in deep
learning. We follow the trajectory of weights traversed
by SGD, leading to new geometric insights and the intuition that SWA will lead to better results than standard
training. Empirically, we make the discovery that SWA
notably improves training of many state-of-the-art deep
neural networks over a range of consequential benchmarks, with essentially no overhead.

The procedures for training neural networks are constantly being improved. New methods are being proposed for architecture design, regularization and optimization. The SWA approach is related to work in both
optimization and regularization.
In optimization, there is great interest in how different
types of local optima affect generalization in deep learning. Keskar et al. [2017] claim that SGD is more likely to
converge to broad local optima than batch gradient methods, which tend to converge to sharp optima. Moreover,
they argue that the broad optima found by SGD are more
likely to have good test performance, even if the training
loss is worse than for the sharp optima. On the other hand
Dinh et al. [2017] argue that all the known definitions of
sharpness are unsatisfactory and cannot on their own explain generalization. Chaudhari et al. [2016] propose the
Entropy-SGD method that explicitly forces optimization
towards wide valleys. They report that although the optima found by Entropy-SGD are wider than those found
by conventional SGD, the generalization performance is
still comparable.
The SWA method is based on averaging multiple points
along the trajectory of SGD with cyclical or constant
learning rates. The general idea of maintaining a running
average of weights proposed by SGD was first considered in convex optimization by Ruppert [1988] and later
by Polyak and Juditsky [1992]. However, this procedure
is not typically used to train neural networks. Practitioners instead sometimes use an exponentially decaying running average of the weights found by SGD with
a decaying learning rate, which smooths the trajectory of
SGD but performs comparably.
SWA is making use of multiple samples gathered through
exploration of the set of points corresponding to high performing networks. To enforce exploration we run SGD
with constant or cyclical learning rates. Mandt et al.
[2017] show that under several simplifying assumptions

running SGD with a constant learning rate is equivalent
to sampling from a Gaussian distribution centered at the
minimum of the loss, and the covariance of this Gaussian
is controlled by the learning rate. Following this explanation from [Mandt et al., 2017], we can interpret points
proposed by SGD as being constrained to the surface of
a sphere, since they come from a high dimensional Gaussian distribution. SWA effectively allows us to go inside
the sphere to find higher density solutions.
In a procedure called Fast Geometric Ensembling (FGE),
Garipov et al. [2018] showed that using a cyclical learning rate it is possible to gather models that are spatially
close to each other but produce diverse predictions. They
used the gathered models to train ensembles with no
computational overhead compared to training a single
DNN model. In recent work Neklyudov et al. [2018]
also discuss an efficient approach for model averaging
of Bayesian neural networks. SWA was inspired by following the trajectories of FGE proposals, in order to find
a single model that would approximate an FGE ensemble, but provide greater interpretability, convenience, and
test-time scalability.
Dropout [Srivastava et al., 2014] is an extremely popular approach to regularizing DNNs. Across each minibatch used for SGD, a different architecture is created
by randomly dropping out neurons. The authors make
analogies between dropout, ensembling, and Bayesian
model averaging. At test time, an ensemble approach
is proposed, but then approximated with similar results
by multiplying each connection by the dropout rate. At a
high level, SWA and Dropout are both at once regularizers and training procedures, motivated to approximate an
ensemble. Each approach implements these high level
ideas quite differently, and as we show in our experiments, can be combined for improved performance.

3

Stochastic Weight Averaging

We present Stochastic Weight Averaging (SWA) and analyze its properties. In section 3.1, we consider trajectories of SGD with a constant and cyclical learning rate,
which helps understand the geometry of SGD training
for neural networks, and motivates the SWA procedure.
Then in section 3.2 we present the SWA algorithm in
detail, in section 3.3 we derive its complexity, and in
section 3.4 we analyze the width of optima found by
SWA versus conventional SGD training. In section 3.5
we then examine the relationship between SWA and the
recently proposed Fast Geometric Ensembling [Garipov
et al., 2018]. Finally, in section 3.6 we consider SWA
from the perspective of stochastic convex optimization.
We note the name SWA has two meanings: on the one

hand, it is an average of SGD weights. On the other,
with a cyclical or constant learning rate, SGD proposals
are approximately sampling from the loss surface of the
DNN, leading to stochastic weights.
3.1

Analysis of SGD Trajectories

SWA is based on averaging the samples proposed by
SGD using a learning rate schedule that allows exploration of the region of weight space corresponding to
high-performing networks. In particular we consider
cyclical and constant learning rate schedules.
The cyclical learning rate schedule that we adopt is inspired by Garipov et al. [2018] and Smith and Topin
[2017]. In each cycle we linearly decrease the learning
rate from α1 to α2 . The formula for the learning rate at
iteration i is given by
α(i) = (1 − t(i))α1 + t(i)α2 ,
1
t(i) = (mod(i − 1, c) + 1) .
c
The base learning rates α1 ≥ α2 and the cycle length c
are the hyper-parameters of the method. Here by iteration we assume the processing of one batch of data. Figure 2 illustrates the cyclical learning rate schedule and
the test error of the corresponding points. Note that unlike the cyclical learning rate schedule of Garipov et al.
[2018] and Smith and Topin [2017], here we propose to
use a discontinuous schedule that jumps directly from
the minimum to maximum learning rates, and does not
steadily increase the learning rate as part of the cycle.
We use this more abrupt cycle because for our purposes
exploration is more important than the accuracy of individual proposals. For even greater exploration, we also
consider constant learning rates α(i) = α1 .
We run SGD with cyclical and constant learning rate
schedules starting from a pretrained point for a Preactivation ResNet-110 on CIFAR-100. We then use the first,
middle and last point of each of the trajectories to define a 2-dimensional plane in the weight space containing all affine combinations of these points. In Figure 3
we plot the loss on train and error on test for points in
these planes. We then project the other points of the trajectory to the plane of the plot. Note that the trajectories
do not generally lie in the plane of the plot, except for the
first, last and middle points, showed by black crosses in
the figure. Therefore for other points of the trajectories it
is not possible to tell the value of train loss and test error
from the plots.
The key insight from Figure 3 is that both methods explore points close to the periphery of the set of highperforming networks. The visualizations suggest that

Learning rate

Algorithm 1 Stochastic Weight Averaging
α1

α2
n

Test error (%)

50
45
40
35
30
25
20
0

1c

2c

3c

iteration number

Figure 2: Top: cyclical learning rate as a function of
iteration. Bottom: test error as a function of iteration
for cyclical learning rate schedule with PreactivationResNet-110 on CIFAR-100. Circles indicate iterations
corresponding to the minimum learning rates.

both methods are doing exploration in the region of space
corresponding to DNNs with high accuracy. The main
difference between the two approaches is that the individual proposals of SGD with a cyclical learning rate
schedule are in general much more accurate than the proposals of a fixed-learning rate SGD. After making a large
step, SGD with a cyclical learning rate spends several
epochs fine-tuning the resulting point with a decreasing
learning rate. SGD with a fixed learning rate on the other
hand is always making steps of relatively large sizes, exploring more efficiently than with a cyclical learning rate,
but the individual proposals are worse.
Another important insight we can get from Figure 3 is
that while the train loss and test error surfaces are qualitatively similar, they are not perfectly aligned. The shift
between train and test suggests that more robust central
points in the set of high-performing networks can lead to
better generalization. Indeed, if we average several proposals from the optimization trajectories, we get a more
robust point that has a substantially higher test performance than the individual proposals of SGD, and is essentially centered on the shifted mode for test error. We
further discuss the reasons for this behaviour in sections
3.4, 3.5, 3.6.
3.2

SWA Algorithm

We now present the details of the Stochastic Weight Averaging algorithm, a simple but effective modification for
training neural networks, motivated by our observations
in section 3.1.

4c

Require:
weights ŵ, LR bounds α1 , α2 ,
cycle length c (for constant learning rate c = 1), number of iterations n
Ensure: wSWA
w ← ŵ {Initialize weights with ŵ}
wSWA ← w
for i ← 1, 2, . . . , n do
α ← α(i) {Calculate LR for the iteration}
w ← w − α∇Li (w) {Stochastic gradient update}
if mod(i, c) = 0 then
nmodels ← i/c {Number of models}
·nmodels +w
{Update average}
wSWA ← wSWA
nmodels +1
end if
end for

Following Garipov et al. [2018], we start with a pretrained model ŵ. We will refer to the number of epochs
required to train a given DNN with the conventional
training procedure as its training budget and will denote
it by B. The pretrained model ŵ can be trained with the
conventional training procedure for full training budget
or reduced number of epochs (e.g. 0.75B). In the latter case we just stop the training early without modifying the learning rate schedule. Starting from ŵ we continue training, using a cyclical or constant learning rate
schedule. When using a cyclical learning rate we capture
the models wi that correspond to the minimum values of
the learning rate (see Figure 2), following Garipov et al.
[2018]. For constant learning rates we capture models
at each epoch. Next, we average the weights of all the
captured networks wi to get our final model wSWA . The
procedure is summarized in Algorithm 1.
Note that for cyclical learning rate schedule, the SWA
algorithm is related to FGE [Garipov et al., 2018], except
that instead of averaging the predictions of the models,
we average their weights, and we use a different type of
learning rate cycle. In section 3.5 we show how SWA
can approximate FGE, but with a single model.
3.3

Computational Complexity

The time and memory overhead of SWA compared to
conventional training is negligible. During training, we
need to maintain a copy of the DNN that aggregates the
average of the weights in memory. The memory consumption of training is dominated by storing the activations of the DNN. The size of the weights is relatively
small even for large DNNs. After the training is complete we only need to store the model that aggregates the
average.

Test error (%)

Train loss

> 50

> 1.7
1.7

30

50

30

35.97

0.7899
20

0.3758
0.1874

10

20

28.49
24.5

10

0.1017
0

22.38

0.06269

0

21.24

0.04494
−10

0.03013
−10

0

10

20

30

40

20.64
−10

50

19.95
−10

0

10

20

30

40

50

Test error (%)

Train loss

> 50

> 1.7

50

1.7

30

30

37.03

0.883
0.5062

20

30.04

20

0.3324
0.2522

10

26.28
24.26

10

0.2152
0

0.1981

23.17
0

22.58

0.1835
0

10

20

30

40

50

21.9
0

10

20

30

40

50

Figure 3: (Left) The L2 -regularized cross-entropy train loss and (Right) test error surfaces of a Preactivation ResNet110 on CIFAR-100 in the plane containing the first, middle and last points (indicated by black crosses) in the trajectories with (Top) cyclical and (Bottom) constant learning rate schedules.
During training extra time is only spent to update the aggregated weight average. This operation is of the form
wSWA ←

wSWA · nmodels + w
,
nmodels + 1,

and it only requires computing a weighted sum of the
weights of two DNNs. As we apply this operation at
most once per epoch, the computational overhead is negligible. Indeed, a similar operation is performed as a part
of each gradient step, and each epoch consists of hundreds of gradient steps.
If the DNN uses batch normalization [Ioffe and Szegedy,
2015], we need to compute the running mean and standard deviation of the activations for each layer of the network found by SWA after the training is finished. To do
so we run one additional pass over data in the same way
as in Garipov et al. [2018].
3.4

Optima Width

Keskar et al. [2017] and Chaudhari et al. [2016] conjecture that the width of a local optimum is related to generalization. The general explanation for the importance
of width is that the surfaces of train loss and test error

are shifted with respect to each other and it is thus desirable to converge to the modes of broad optima, which
stay approximately optimal under small perturbations. In
this section we compare the solutions found by SWA and
SGD and show that SWA generally leads to much wider
optima.
Let wSWA and wSGD denote the weights of DNNs trained
using SWA and conventional SGD, respectively. Consider the rays
wSWA (t, d) = wSWA + t · d,

wSGD (t, d) = wSGD + t · d,

which follow a direction vector d on the unit sphere,
starting at wSWA and wSGD , respectively. In Figure
4 we plot train loss and test error of wSWA (t, di ) and
wSGD (t, di ) as a function of t for 10 random directions
di , i = 1, 2, . . . , 10 drawn from a uniform distribution
on the unit sphere. For this visualization we use a Preactivation ResNet-110 on CIFAR-100.
First, while the loss values on train for wSGD and wSWA
are quite similar (and in fact wSGD has a slightly lower
train loss), the test error for wSGD is lower by 1.5% (at
the converged value corresponding to t = 0). Further,
the shapes of both train loss and test error curves are con-

0.20

30
28

0.15

26

0.10
24

0.05

22
20
0

5

10

15

0.00

20

0

5

10

15

20

Figure 4: (Left) Test error and (Right) L2 -regularized cross-entropy train loss as a function of a point on a random ray
starting at SWA (blue) and SGD (green) solutions for Preactivation ResNet-110 on CIFAR-100. Each line corresponds
to a different random ray.
2.5

25.0

1.5

22.5

1.0

20.0

0.5

17.5

0.0
−80

−60

−40

−20

0

20

Test error
SWA
SGD

39.5
2.0

Test error (%)

Test error (%)

27.5

Train loss
SWA
SGD

Train loss

Test error
SWA
SGD

Train loss
SWA
SGD

1.5

34.5

1.0

29.5

0.5

24.5

0.0
−60

40

Distance

Train loss

30.0

−40

−20

0

20

40

Distance

Figure 5: L2 -regularized cross-entropy train loss and test error as a function of a point on the line connecting SWA
and SGD solutions on CIFAR-100. Left: Preactivation ResNet-110. Right: VGG-16.
siderably wider for wSWA than for wSGD , suggesting that
SWA indeed converges to a wider optimum: we have to
step much further away from the solution found by wSWA
to increase error by a given amount. We even see the
error curve for SGD has an inflection point that is not
present for these distances with SWA.
Notice that in Figure 4 any of the random directions from
wSGD increase test error. However, we know that the direction from wSGD to wSWA would decrease test error,
since wSWA has considerably lower test error than wSGD .
In other words, the path from wSGD to wSWA is qualitatively different from all directions shown in Figure 4, because along this direction wSGD is far from optimal. We
therefore consider the line segment connecting wSGD and
wSWA :
w(t) = t · wSGD + (1 − t) · wSWA .
In Figure 5 we plot the train loss and test error of w(t)
as a function of signed distance from wSWA for Preactivation ResNet-110 and VGG-16 on CIFAR-100.
We can extract several key insights about wSWA and wSGD
from Figure 5. First, the train loss and test error plots
are indeed substantially shifted, and the point obtained
by minimizing the train loss is far from optimal on test.

Second, wSGD lies near the boundary of a wide flat region
of the train loss. Further, the loss is very steep near wSGD .
Keskar et al. [2017] argue that the loss near sharp optima found by SGD with very large batches are actually
flat in most directions, but there exist directions in which
the optima are extremely steep. They conjecture that because of this sharpness the generalization performance
of large batch optimization is substantially worse than
that of solutions found by small batch SGD. Remarkably, in our experiments in this section we observe that
there exist directions of steep ascent even for small batch
optima, and that SWA provides even wider solutions (at
least along random directions) with better generalization.
3.5

Connection to Ensembling

Garipov et al. [2018] proposed the Fast Geometric Ensembling (FGE) procedure for training ensembles in the
time required to train a single model. Using a cyclical
learning rate, FGE generates a sequence of points that
are close to each other in the weight space, but produce
diverse predictions. In SWA instead of averaging the predictions of the models we average their weights. However, the predictions proposed by FGE ensembles and

SWA models have similar properties.
Let f (·) denote the predictions of a neural network
parametrized by weights w. We will assume that f is
a scalar (e.g. the probability for a particular class) twice
continuously differentiable function with respect to w.
Consider points wi proposed by FGE. These points are
close in the weight space by design,
Pn and concentrated
1
around their average wSWA =
Pnn i=1 wi . We denote
∆i = wi − wSWA . Note that i=1 ∆i = 0. Ensembling
the networks wi corresponds to averaging the function
values
n

1X
f (wi ).
f¯ =
n i=1
Consider the linearization of f at wSWA .
f (wj ) = f (wSWA ) + h∇f (wSWA ), ∆j i + O(k∆j k2 ),
where h·, ·i denotes the dot product. Thus, the difference
between averaging the weights and averaging the predictions
n


1X
f¯ − f (wSWA ) =
h∇f (wSWA ), ∆i i + O(k∆i k2 )
n i=1
+
*
n
1X
∆i + O(∆2 ) = O(∆2 ),
= ∇f (wSWA ),
n i=1
where ∆ = maxni=1 k∆i k. Note that the difference between the predictions of different perturbed networks is
f (wi ) − f (wj ) = h∇f (wSWA ), ∆i − ∆j i + O(∆2 ),
and is thus of the first order of smallness, while the
difference between averaging predictions and averaging
weights is of the second order of smallness. Note that for
the points proposed by FGE the distances between proposals are relatively small by design, which justifies the
local analysis.
To analyze the difference between ensembling and averaging the weights of FGE proposals in practice, we
run FGE for 20 epochs and compare the predictions of
different models on the test dataset with a Preactivation
ResNet-110 [He et al., 2016] on CIFAR-100. The norm
of the difference between the class probabilities of consecutive FGE proposals averaged over the test dataset is
0.126. We then average the weights of the proposals
and compute the class probabilities on the test dataset.
The norm of difference of the probabilities for the SWA
model and the FGE ensemble is 0.079, which is substantially smaller than the difference between the probabilities of consecutive FGE proposals. Further, the fraction
of objects for which consecutive FGE proposals output

the same labels is not greater than 87.33%. For FGE
and SWA the fraction of identically labeled objects is
95.26%.
The theoretical considerations and empirical results presented in this section suggest that SWA can approximate
the FGE ensemble with a single model.

3.6

Connection to Convex Minimization

Mandt et al. [2017] showed that under strong simplifying assumptions SGD with a fixed learning rate approximately samples from a Gaussian distribution centered
at the minimum of the loss. Suppose this is the case
when we run SGD with a fixed learning rate for training a DNN.
Let us denote the dimensionality of the weight space of
the neural network by d. Denote the samples produced
by SGD by wi , i = 1, 2, . . . , k. Assume the points wi
are concentrated around the local optimum ŵ. The SWA
Pk
solution is given by wSWA = n1 i=1 wi . The points wi
are samples from a multidimensional Gaussian N (ŵ, Σ)
for some covariance matrix Σ defined by the curvature of
the loss, batch size and the learning rate. Note that the
samples from a multidimensional Gaussian are concentrated on the ellipsoid
n
√ o
1
z ∈ Rd | kΣ− 2 (z − ŵ)k = d ,
and the probability mass for a sample to end up inside the
ellipsoid near ŵ is negligible On the other hand, wSWA is
guaranteed to converge to ŵ as k → ∞.
Moreover, Polyak and Juditsky [1992] showed that averaging SGD proposals achieves the best possible convergence rate among all stochastic gradient algorithms. The
proof relies on the convexity of the underlying problem
and in general there are no convergence guarantees if the
loss function is non-convex [see e.g. Ghadimi and Lan,
2013]. The DNN loss functions are known to be nonconvex [e.g. Choromanska et al., 2015]. However, the
trajectory of SGD the loss surfaces of DNNs are approximately convex [e.g. Goodfellow et al., 2015].
Finally, there is not a direct correspondence between
minimizing the standard training loss functions (e.g. L2
regularized cross-entropy) and choosing parameters that
are best for generalization.
In Figure 5 we can see that SWA converges to a central point, ignoring the local non-convexities of the DNN
loss.

Table 1: Accuracies (%) of SWA, SGD and FGE methods on CIFAR-100 and CIFAR-10 datasets for different training
budgets. Accuracies for the FGE ensemble are from Garipov et al. [2018].

DNN (Budget)

SGD

VGG-16 (200)
ResNet-110 (150)
WRN-28-10 (200)
PyramidNet-272 (300)

72.55 ± 0.10
78.49 ± 0.36
80.82 ± 0.23
83.41 ± 0.21

VGG-16 (200)
ResNet-110 (150)
WRN-28-10 (200)
ShakeShake-2x64d (1800)

93.25 ± 0.16
95.28 ± 0.10
96.18 ± 0.11
96.93 ± 0.10

FGE (1 Budget)

SWA
1.25 Budgets

1.5 Budgets

73.91 ± 0.12
79.77 ± 0.17
81.46 ± 0.23
–

74.17 ± 0.15
80.18 ± 0.23
81.91 ± 0.27
83.93 ± 0.18

74.27 ± 0.25
80.35 ± 0.16
82.15 ± 0.27
84.16 ± 0.15

93.59 ± 0.16
95.56 ± 0.11
96.45 ± 0.11
–

93.70 ± 0.22
95.77 ± 0.04
96.64 ± 0.08
97.16 ± 0.10

93.64 ± 0.18
95.83 ± 0.03
96.79 ± 0.05
97.12 ± 0.06

1 Budget

CIFAR-100

4

74.26
79.84
82.27
–
CIFAR-10

93.52
95.45
96.36
–

Experiments

We compare SWA against conventional SGD training
on CIFAR-10, CIFAR-100 and ImageNet ILSVRC-2012
[Russakovsky et al., 2012]. We also compare to Fast Geometric Ensembling (FGE) [Garipov et al., 2018], but
we note that FGE is an ensemble whereas SWA corresponds to a single model. Conventional SGD training
uses a standard decaying learning rate schedule (details
in the Appendix) until convergence. We found an exponentially decaying average of SGD to perform comparably to conventional SGD at convergence.
4.1

CIFAR Datasets

For the experiments on CIFAR datasets we use VGG16 [Simonyan and Zisserman, 2014], a 110-layer
Preactivation-ResNet [He et al., 2016] and Wide ResNet28-10 [Zagoruyko and Komodakis, 2016] models. Additionally, we experiment with the recent Shake-Shake2x64d [Gastaldi, 2017] on CIFAR-10 and PyramidNet272 (bottleneck, α = 200) [Han et al., 2016] on CIFAR100. All models are trained using L2 -regularization, and
VGG-16 also uses dropout.
For each model we define budget as the number of
epochs required to train the model until convergence with
conventional SGD training, such that we do not see improvement with SGD beyond this budget. We use the
same budgets for VGG, Preactivation ResNet and Wide
ResNet models as Garipov et al. [2018]. For ShakeShake and PyramidNets we use the budgets indicated by
the papers that proposed these models [Gastaldi, 2017,
Han et al., 2016]. We report the results of SWA training
within 1, 1.25 and 1.5 budgets of epochs.
For VGG, Wide ResNet and Preactivation-ResNet models we first run standard SGD training for ≈ 75% of the

training budget, and then use the weights at the last epoch
as an initialization for SWA with a fixed learning rate
schedule. We ran SWA for 0.25, 0.5 and 0.75 budget
to complete the training within 1, 1.25 and 1.5 budgets
respectively.
For Shake-Shake and PyramidNet architectures we do
not report the results in one budget. For these models
we use a full budget to get an initialization for the procedure, and then train with a cyclical learning rate schedule
for 0.25 and 0.5 budgets. We used long cycles of small
learning rates for Shake-Shake, because this architecture
already involves many stochastic components.
We present the details of the learning rate schedules for
each of these models in the Appendix.
For each model we also report the results of conventional
SGD training, which we denote by SGD. For VGG, Preactivation ResNet and Wide ResNet we also provide the
results of the FGE method with one budget reported in
Garipov et al. [2018]. Note that for FGE we report the
accuracy of an ensemble of 6 to 12 networks, while for
SWA we report the accuracy of a single model.
We summarize the experimental results in Table 1. For
all models we report the mean and standard deviation
of test accuracy over 3 runs. In all conducted experiments SWA substantially outperforms SGD in one budget, and improves further, as we allow more training
epochs. Across different architectures we see consistent improvement by ≈ 0.5% on CIFAR-10 (excluding
Shake-Shake, for which SGD performance is already extremely high) and by 0.75-1.5% on CIFAR-100. Amazingly, SWA is able to achieve comparable or better performance than FGE ensembles with just one model. On
CIFAR-100 SWA usually needs more than one budget
to get results comparable with FGE ensembles, but on
CIFAR-10 even with 1 budget SWA outperforms FGE.

4.2

ImageNet

On ImageNet we experimented with ResNet-50, ResNet152 [He et al., 2016] and DenseNet-161 [Huang et al.,
2017]. For these architectures we used pretrained models from PyTorch.torchvision. For each of the
models we ran SWA for 10 epochs with a cyclical learning rate schedule with the same parameters for all models
(the details can be found in the Appendix), and report the
mean and standard deviation of test error averaged over
3 runs. The results are shown in Table 2.
Table 2: Accuracies (%) on ImageNet dataset for SWA
and SGD with different architectures.

DNN
ResNet-50
ResNet-152
DenseNet-161

SGD
76.15
78.31
77.65

SWA
5 epochs
10 epochs
76.83 ± 0.01 76.97 ± 0.05
78.82 ± 0.01 78.94 ± 0.07
78.26 ± 0.09 78.44 ± 0.06

rate parameters (α1 , α2 ) ∈ {(10−1 , 10−3 ), (5 · 10−2 , 5 ·
10−4 ), (10−2 , 10−4 ), (5 · 10−3 , 5 · 10−5 )}. Among the
constant learning rates we consider α1 ∈ {10−1 , 5 ·
10−2 , 10−2 , 10−3 }.
We plot the test error of the SWA procedure for different
learning rate schedules as a function of the number of
training epochs in Figure 6.
We find that in general the more aggressive constant
learning rate schedule leads to faster convergence of
SWA. In our experiments we found that setting the learning rate to some intermediate value between the largest
and the smallest learning rate used in the annealing
scheme in conventional training usually gave us the best
results. The approach is however universal and can work
well with different learning rate schedules tailored for
particular tasks.
4.4

DNN training with a fixed learning rate

50

For all 3 architectures SWA provides consistent improvement by 0.6-0.9% over the pretrained models.
Effect of the Learning Rate Schedule

26

Baseline
LR = 0.1
LR = 0.05
LR = 0.01
LR = 0.001

25

Test error (%)

24
23

CLR(0.01, 0.0001)
CLR(0.005, 0.00005)
CLR(0.05, 0.0005)
CLR(0.1, 0.001)

30
25

SGD
Const LR SGD
Const LR SWA

20
15

50

100

150

200

250

300

Epochs

21
20
19
160

35

0

22

140

40

Test error (%)

4.3

45

180

200

220

Figure 7: Test error as a function of training epoch for
constant (green) and decaying (blue) learning rate schedules for a Wide ResNet-28-10 on CIFAR-100. In red we
average the points along the trajectory of SGD with constant learning rate starting at epoch 140.

Epochs

Figure 6: Test error as a function of training epoch for
SWA with different learning rate schedules with a Preactivation ResNet-110 on CIFAR-100.
In this section we explore how the learning rate schedule
affects the performance of SWA. We run experiments on
Preactivation ResNet-110 on CIFAR-100. For all schedules we use the same initialization from a model trained
for 125 epochs using the conventional SGD training. As
a baseline we use a fully-trained model trained with conventional SGD for 150 epochs.
We consider a range of constant and cyclical learning
rate schedules. For cyclical learning rates we fix the cycle length to 5, and consider the pairs of base learning

In this section we show that it is possible to train DNNs
from scratch with a fixed learning rate using SWA. We
run SGD with a fixed learning rate of 0.05 on a Wide
ResNet-28-10 [Zagoruyko and Komodakis, 2016] for
300 epochs from a random initialization on CIFAR-100.
We then averaged the weights at the end of each epoch
from epoch 140 and until the end of training. The final
test accuracy of this SWA model was 81.7.
Figure 7 illustrates the test error as a function of the number of training epochs for SWA and conventional training. The accuracy of the individual models with weights
averaged by SWA stays at the level of ≈ 65% which is
16% less than the accuracy of the SWA model. These results correspond to our intuition presented in section 3.6
that SGD with a constant learning rate oscillates around

the optimum, but SWA converges.
While being able to train DNN with fixed learning rate is
a surprising property of SWA, for practical purposes we
recommend initializing SWA from a model pretrained
with conventional training (possibly for a reduced number of epochs), as it leads to faster and more stable convergence than running SWA from scratch.

5

Discussion

We have presented Stochastic Weight Averaging (SWA)
for training neural networks. SWA is extremely easy to
implement, architecture-agnostic, and improves generalization performance at virtually no additional cost over
conventional training.
There are so many exciting directions for future research.
SWA does not require each weight in its average to correspond to a good solution, due to the geometry of weights
traversed by the algorithm. It therefore may be possible
to develop SWA for much faster convergence than standard SGD. One may also be able to combine SWA with
large batch sizes while preserving generalization performance, since SWA discovers much broader optima than
conventional SGD training. Furthermore, a cyclic learning rate enables SWA to explore regions of high posterior density over neural network weights. Such learning
rate schedules could be developed in conjunction with
stochastic MCMC approaches, to encourage exploration
while still providing high quality samples. One could
also develop SWA to average whole regions of good
solutions, using the high-accuracy curves discovered in
Garipov et al. [2018].
A better understanding of the loss surfaces for multilayer
networks will help continue to unlock the potential of
these rich models. We hope that SWA will inspire further
progress in this area.
Acknowledgements. This work was supported by
NSF IIS-1563887, Samsung Research, Samsung Electronics and Russian Science Foundation grant 17-1101027. We also thank Vadim Bereznyuk for helpful comments.

References
Pratik Chaudhari, Anna Choromanska, Stefano Soatto,
and Yann LeCun.
Entropy-sgd: Biasing gradient descent into wide valleys.
arXiv preprint
arXiv:1611.01838, 2016.
Anna Choromanska, Mikael Henaff, Michael Mathieu,
Gérard Ben Arous, and Yann LeCun. The loss surfaces

of multilayer networks. In Artificial Intelligence and
Statistics, pages 192–204, 2015.
Laurent Dinh, Razvan Pascanu, Samy Bengio, and
Yoshua Bengio. Sharp minima can generalize for deep
nets. arXiv preprint arXiv:1703.04933, 2017.
F. Draxler, K. Veschgini, M. Salmhofer, and F. A. Hamprecht. Essentially No Barriers in Neural Network Energy Landscape. arXiv preprint, 2018.
Timur Garipov, Pavel Izmailov, Dmitrii Podoprikhin,
Dmitry P Vetrov, and Andrew Gordon Wilson. Loss
surfaces, mode connectivity, and fast ensembling of
DNNs. arXiv preprint arXiv:1802.10026, 2018.
Xavier Gastaldi. Shake-shake regularization.
preprint arXiv:1705.07485, 2017.

arXiv

Saeed Ghadimi and Guanghui Lan. Stochastic first-and
zeroth-order methods for nonconvex stochastic programming. SIAM Journal on Optimization, 23(4):
2341–2368, 2013.
Ian J Goodfellow, Oriol Vinyals, and Andrew M Saxe.
Qualitatively characterizing neural network optimization problems. International Conference on Learning
Representations, 2015.
Dongyoon Han, Jiwhan Kim, and Junmo Kim.
Deep pyramidal residual networks. arXiv preprint
arXiv:1610.02915, 2016.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian
Sun. Deep residual learning for image recognition.
In Proceedings of the IEEE conference on computer
vision and pattern recognition, pages 770–778, 2016.
Gao Huang, Zhuang Liu, Kilian Q Weinberger, and Laurens van der Maaten. Densely connected convolutional
networks. In Proceedings of the IEEE conference on
computer vision and pattern recognition, volume 1,
page 3, 2017.
Sergey Ioffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing
internal covariate shift. In International Conference
on Machine Learning, pages 448–456, 2015.
Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang.
On large-batch training for deep learning: Generalization gap and sharp minima. International Conference
on Learning Representations, 2017.
Ilya Loshchilov and Frank Hutter. Sgdr: stochastic gradient descent with restarts. arxiv preprint
arxiv:1608.03983, 2016.
Stephan Mandt, Matthew D Hoffman, and David M Blei.
Stochastic gradient descent as approximate bayesian
inference. arXiv preprint arXiv:1704.04289, 2017.

Boris T Polyak and Anatoli B Juditsky. Acceleration of
stochastic approximation by averaging. SIAM Journal
on Control and Optimization, 30(4):838–855, 1992.
David Ruppert. Efficient estimations from a slowly
convergent robbins-monro process. Technical report,
Cornell University Operations Research and Industrial
Engineering, 1988.
Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause,
Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej
Karpathy, Aditya Khosla, Michael Bernstein, et al.
Imagenet large scale visual recognition challenge. International Journal of Computer Vision, 115(3):211–
252, 2012.
Karen Simonyan and Andrew Zisserman. Very deep convolutional networks for large-scale image recognition.
arXiv preprint arXiv:1409.1556, 2014.
Leslie N Smith and Nicholay Topin. Exploring loss
function topology with cyclical learning rates. arXiv
preprint arXiv:1702.04283, 2017.
Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky,
Ilya Sutskever, and Ruslan Salakhutdinov. Dropout:
A simple way to prevent neural networks from overfitting. The Journal of Machine Learning Research, 15
(1):1929–1958, 2014.
Sergey Zagoruyko and Nikos Komodakis. Wide residual
networks. arXiv preprint arXiv:1605.07146, 2016.

A

Appendix

A.1

Experimental details

α1

Learning rate

K. Neklyudov, D. Molchanov, A. Ashukha, and
D. Vetrov. Variance Networks: When Expectation
Does Not Meet Your Expectations. ArXiv e-prints,
2018.

α2
0

100

200

300

400

Epochs

Figure 8: Cyclical learning rate used for Shake-Shake as
a function of iteration.

SWA learning rates. For PyramidNet SWA uses a
cyclic learning rate with α1 = 0.05 and α2 = 0.001
and cycle length 3. For VGG and Wide ResNet we used
constant learning α1 = 0.01. For ResNet we used constant learning rates α1 = 0.01 on CIFAR-10 and 0.05 on
CIFAR-100.
For Shake-Shake Net we used a custom cyclic learning rate based on the cosine annealing used when training Shake-Shake with SGD. Each of the cycles replicate
the learning rates corresponding to epochs 1600 − 1700
of the standard training and the cycle length c = 100
epochs. The learning rate schedule is depicted in Figure
8 and follows the formula



1600 + epoch(i) mod 100)
α(i) = 0.1 · 1 + cos π ·
,
1800
where epoch(i) is the number of data passes completed
before iteration i.
For all experiments with ImageNet we used cyclic learning rate schedule with the same hyperparameters α1 =
0.001, α2 = 10−5 and c = 1.

For the experiments on CIFAR datasets (section 4.1) we
used the following implementations (embedded links):
• Shake-Shake-2x64d
• PyramidNet-272
• VGG-16
• Preactivation-ResNet-110
• Wide ResNet-28-10
Models for ImageNet are from here. Pretrained networks
can be found here.

SGD learning rates. For conventional SGD training
we used SGD with momentum 0.9 and with an annealed
learning rate schedule. For VGG, Wide ResNet and Preactivation ResNet we fixed the learning rate to α1 for the
first half of epochs (0B–0.5B), then linearly decreased
the learning rate to 0.01α1 for the next 40% of epochs
(0.5B–0.9B), and then kept it constant for the last 10%
of epochs (0.9B – 1B). For VGG we set α1 = 0.05,
and for Preactivation ResNet and Wide ResNet we set
α1 = 0.1. For Shake-Shake Net and PyramidNets
we used the cosine and piecewise-constant learning rate
schedules described in Gastaldi [2017] and Han et al.
[2016] respectively.

A.2

Training ResNet with a Constant Learning
Rate

In this section we present the experiment on training
Preactivation ResNet-110 using a constant learning rate.
The experimental setup is the same as in section 4.4. We
set the learning rate to α1 = 0.1 and start averaging after
epoch 200. The results are presented in Figure 9.
50

SGD
Const LR SGD
Const LR SWA

Test error (%)

45

40

35

30

25

20
0

50

100

150

200

250

300

Epochs

Figure 9: Test error as a function of training epoch for
constant (green) and decaying (blue) learning rate schedules for a Preactivation ResNet-110 on CIFAR-100. In
red we average the points along the trajectory of SGD
with constant learning rate starting at epoch 200.

