# Training trajectories, mini-batch losses and the curious role of the
learning rate.

Mark Sandler  Andrey Zhmoginov  Max Vladymyrov  Nolan Miller
Google Research
<EMAIL_ADDRESS>
###### Abstract

Stochastic gradient descent plays a fundamental role in nearly all
applications of deep learning. However its ability to converge to a global
minimum remains shrouded in mystery. In this paper we propose to study the
behavior of the loss function on fixed mini-batches along SGD trajectories. We
show that the loss function on a fixed batch appears to be remarkably convex-
like. In particular for ResNet the loss for any fixed mini-batch can be
accurately modeled by a quadratic function and a very low loss value can be
reached in just one step of gradient descent with sufficiently large learning
rate. We propose a simple model that allows to analyze the relationship
between the gradients of stochastic mini-batches and the full batch. Our
analysis allows us to discover the equivalency between iterate aggregates and
specific learning rate schedules. In particular, for Exponential Moving
Average (EMA) and Stochastic Weight Averaging we show that our proposed model
matches the observed training trajectories on ImageNet. Our theoretical model
predicts that an even simpler averaging technique, averaging just two points a
many steps apart, significantly improves accuracy compared to the baseline. We
validated our findings on ImageNet and other datasets using ResNet
architecture.

##  1 Introduction

Stochastic Gradient Descent (SGD; Robbins & Monro, 1951) played a fundamental
role in the rise and success of deep learning. Similarly, the learning rate,
the multiplier controlling the magnitude of the weight update, is perhaps the
most crucial hyper-parameter that controls the training trajectory (Goodfellow
et al., 2016; Smith, 2015; Bengio, 2012; Smith, 2015). Despite the advances in
stochastic optimization methods (Kingma & Ba, 2014) designed to reduce the
need for tuning, and in particular dependence on learning rate, novel learning
rate schedules such as cosine (Loshchilov & Hutter, 2016) and cyclical
learning rate schedule (Smith, 2015) continue to be the subject of active
research and are far from being rigorously understood. On the other hand, the
impact of stochasticity of the optimization process caused by changing input
mini-batches also known as ”batch noise”, (Ziyin et al., 2021) has been
relatively poorly understood. In

Figure 1: Weight averaging v.s. equivalent learning rate schedule. The dotted
vertical red lines show the set of independent trajectories “sidetrips” with
appropriate learning rate schedule that start at corresponding point in the
main trajectory.

#### Contributions

In this paper we show that the loss as a function of learning rate of the
training batch is dramatically different than that on the held-out batch, when
measured on SGD trajectory. This connection to learning rate leads to a simple
theoretical model describing observed results. We believe this connection of
the learning rate to the loss as a function of held out and training match has
been largely overlooked in the literature.

One practical result that we demonstrate that various popular averaging
techniques have equivalent learning rate schedules. To the best of our
knowledge such connection has not been observed in the literature. We
demonstrate the equivalency in two different scenarios: a simple scenario
where we assume that the gradient for the same input changes little between
steps, and the other where we show that this equivalency holds in the limit
case of an arbitrary long trajectory. We demonstrate (e.g. Figure 1) that
these results hold empirically for a general training of ImageNet (Russakovsky
et al., 2015) using ResNet architecture (He et al., 2015).

We propose a simple model that captures these behaviors. The model we consider
is similar to that of Mandt et al. (2017), but our analysis is noticeably
simpler. We further show that our simple model is suitable for studying
training trajectories with a spectrum of time scales ranging from fast to slow
and corresponding to long trenches in the loss landscape. We show how model
weights quickly converge towards a quasi-stationary distribution in some
degrees of freedom, while slowly drifting towards the global minimum in others
and derive the expression for the auto-correlation of the averaged trajectory
as a function of the averaging kernel (Appendix A).

Using this model we develop a simple geometric explanation of why averaging
improves the accuracy in the first place. We show that even on large datasets
such as ImageNet and for deep architectures using the fixed learning rate the
trajectory traverses an ellipsoid centered around the target minimum and the
averaging moves the trajectory to the inside of that ellipsoid. Further, we
show that for two independently trained trajectories that share a common burn-
in period the _distance_ in parameter space along one trajectory to the final
solution of another monotonically decreases as training progresses and
asymptotically stabilizes at a fixed value. This suggests that most of the
probability mass of the training outcomes is concentrated on a sphere-like
volume.

Our final insight is the dramatic difference in the behavior of the loss
function for the batch that it is being optimized v.s. a new (training) batch.
We are not aware of any study in the literature that explores this difference.
While this is reminiscent of the generalization gap, there is a crucial
distinction: we analyze the loss on a different _training_ batch, and not on a
validation batch. To some degree our work also illuminates the approach of
Amortized Proximal Optimization (Bae et al., 2022) that used the loss on the
unseen batch referred as _Functional Space Discrepancy_ , as a negative
regularization force for learning rate selection. Curiously their motivation
was to reduce the change of the loss on the unseen batch as a way to minimize
the predictions on unseen data. However, as we demonstrate here, such type of
adaptive schedule tries to apply asymptotically largest update, that prevents
loss on unseen batches to _increase_ and thus the training to diverge.

Interestingly, such difference in gradients and behavior of two training
batches, also indicates that a simple intuition (Ziyin et al., 2021; Mandt et
al., 2017; Sato & Nakagawa, 2014b) that stochastic gradient descent is an
approximation to batch gradient could be misleading and thus should be used
judiciously, despite it being unbiased estimate of the full gradient.

#### Paper organization

The paper is organized as follows. In the next section we describe some of the
ImageNet experiments that led us to our theoretical model described in section
Figure 4. In Section 5 we show some additional experiments on both real data
and synthetic model that we introduced in Figure 4. Finally in Section 6 we
outline some of the open questions and future directions.

##  2 Related work

Averaging intermediate SGD iterates in the weight space can be traced back to
Ruppert (1988). In more recent years it become a common tool for speeding up
convergence (Kaddour, 2022) and improving accuracy (Izmailov et al., 2018). It
was first directly analyzed in Polyak & Juditsky (1992) and has been studied
in the literature since. For example in Mandt et al. (2017) the connection
between SGD and Bayesian inference was established. In particular they
demonstrated that SGD training that initialized within common zone of
attraction can be seen as Markov chain with Gaussian stationary distribution
for constant learning rate, which is related to our observation about
trajectory stabilizing at a fixed distance from the global minimum. However in
present work we mainly analyze the behavior of a _s_ ingle trajectory. This
has immediate advantage that it allows us to reason about finite trajectory
lengths, as well as the connection between learning rate schedules and
averaging.

Another related line of inquiry is _Stochastic Weight Averaging_ (SWA) by
Izmailov et al. (2018), where a trajectory with cyclic learning rates and
averaging every $c$ steps is usedas a way to improve accuracy. This in theory
is slightly different from Ruppert-Polyak averaging, since not every point on
the trajectory is averaged, however, many of the experiments presented in that
paper use $c=1$ and we suspect this difference is insignificant. In that work
it was also observed that averaging between trajectory allows to move inside
the surface, but no further discussion has followed. Additionally, the
connection between SWA and generalization was further explored in He et al.
(2019), where it was showed that averaging introduces bias towards flatter
part of the basin and this results in better generalization. In contrast here
our goal is to illuminate the mechanisms by which SWA changes the behavior of
empirical loss in the context of stochastic (rather than full) gradient
descent. Connecting current work and the generalization performance is an
interesting open direction. The quadratic theoretical model we consider is
reminiscent of that in Schaul et al. (2013), but our analysis is both simpler
than former and more general than latter. We note that analysis only applies
to stochastic gradient descent. In case of full gradient descent there have
been several recent works showing that quadratic approximation model might be
toosimplistic (Ma et al., 2022; Damian et al., 2022; Cohen et al., 2021).

Recent work on basins and local minima connectivity (Ainsworth et al., 2022;
Benzing et al., 2022; Frankle et al., 2019) have demonstrated that there is a
lot of apparent structure in the loss space, in particular it was shown that
different initialization lead to different equivalence classes, which can then
be mapped between each other using simple permutation of the neurons. That
work is complementary to ours in that they show that there are simple ways of
moving basins, while we concentrate on properties of a single basin.

##  3 Empirical observations

Figure 2: Loss as a function of a step size in a single step gradient descent
in the middle of a trajectory. The step shown here is step $75\,000$ out of
$100\,000$ run. The behavior is typical for other steps as well. The small
rectangle on the left graph approximately shows the location of the zoomed-in
right graph. The unseen batches is computed over $10$ batches, with dark shade
showing the standard deviation from the mean, while light shade shows the
max/min values observed. Figure 3: Loss for a fixed batch on the interpolation
between two points of the original training trajectory. We use two points
along the ImageNet training trajectory with $t=75\,000$, and $t=85\,000$. From
each point we perform 9 steps of gradient descent on a fixed batch and measure
the loss on interpolation between corresponding pairs of points of each
trajectory. The left graph shows the loss on the training batch, and note how
within just 3 steps it reaches nearly 0. The middle graph shows the loss on
held out. At step 0 the training and held out batch loss profiles are very
similar. Rightmost graph shows the visualization of the process.

In this section our goal is to explore some of the properties of SGD that will
lead us to the theoretical model of Figure 4.

#### Experimental setup

For all our experiments with non-synthettic data we use ResNet34 (He et al.,
2015) architecture and ImageNet (Russakovsky et al., 2015). We use standard
SGD with momentum $0.9$, that produces good accuracy on ImageNet. Most of our
experiments on real data involve starting with partially trained
architectures, and running a separate _side-trip_ with different learning rate
characteristics and/or with a fixed _training_ mini-batch to demonstrate
certain behaviors. We also will use the notion of _held-out_ mini-batch which
is simply a _fixed_ batch different from the training mini-batch. Throughout
this section we will be referring to the training trajectory as the _main_
trajectory to differentiate from side-trips.

#### Single step experiments

We begin with a simple question: how does a loss landscape look for a single
batch applied with a varying learning rate? Specifically, does the loss behave
on that batch v.s. a separate held-out batch? In Figure 2 we plot the loss
v.s. the learning rate when applied to a single step using either the batch
for which the direction was computed or a separate held-out batch. Here we
used step 75K, out of 100K, and batch size 512, however very similar results
hold elsewhere in the trajectory and for different batch sizes. See Figure 12
for details.

As can be seen from Figure 2 the loss for both curves is smooth and can be
well approximated by a low-degree polynomial. More importantly. the loss on
the training batch reaches loss value below the average loss of a fully
trained model in a single step. Typical loss value for fully trained ImageNet
is close to $1$, while a single step starting from the loss value $1.8$
results in a final loss of $0.5$. The loss on the held-out batch behaves very
differently: only sufficiently small learning rate leads to any improvement to
the loss on the full distribution. Therefore any analytical model purporting
to describe SGD trajectories should have this property. Remarkably this
suggests that full gradient trajectory is not a good approximation of the SGD
trajectory, if we want to understan the dynamics of learning rate selection.

#### Multi-step side-trip

As we just observed, a single step along the training batch direction can
bring the loss down to values considerably lower than the best average loss of
a fully trained model. Naturally one can ask – does the loss basin for this
specific batch remains the same regardless of the starting point? In other
words, if we apply gradient descent with a fixed batch at one point of the
main trajectory, would it be the same minima basin as if started from
different point of the main trajectory? Note: this question is different from
the notion of the global basin connectivity explored in Frankle et al. (2019).
Instead, we look at the behavior of the loss for a fixed mini-batch as we
branch off different parts of the training trajectory. Following the analysis
of Frankle et al. (2019) we use interpolation to check that we are in the same
basin. In Figure 3 we show that the side-trip using the same fixed batch along
the main training trajectory leads to the same basin. Remarkably, the held-out
batch loss while increasing dramatically also stays in the same basin as the
main trajectory.

##  4 Analytical model

(a)

(b)

Figure 4: Stochastic gradient descent. Figure 4(a) shows how stochastic
gradient direction decomposes into a drift and a drag component. The magnitude
of the drag component changes as $\sqrt{\lambda}$, while the magnitude of the
drift component is proportional to $\lambda$. Further, the drift component has
a quadratic bias to increase loss, thus enabling us to pick learning rate that
allows to reduce the loss. Figure 4(b) illustrates how stationary trajectory
traverses the sphere of a fixed loss, while taking a midpoint enables to get
_inside_ the sphere. Best viewed in color.

In this section we propose a model describing the behavior of the high
dimensional weight vector $\theta$, as it moves along the SGD trajectory, in
such a way that it can explaining the phenomena that we described in Section
3: the loss for the individual batches behaves like a low-degree polynomial
and reaches remarkably low value in a single step. Thus we can use quadratic
function that models the loss function for individual samples:

$L(\theta,x)=\|A_{x}\theta+c_{x}\|^{2}/2,$

where $A_{x}$ is a matrix and $c_{x}$ is a vector describing the loss function
for a random sample $x\sim D$. This model is similar to the one proposed in
Schaul et al. (2013), however it analyzed only the special case of a constant
and diagonal $A_{x}$, while we consider a general scenario. Following the
mini-batch gradient where individual batches are sampled from some
distribution we then will be moving in a direction whose expectation is
$E_{x\sim{\mathcal{D}}}\left[\frac{\partial{L(\theta,x)}}{\partial\theta}\right]=\frac{\partial{E_{x\sim{\mathcal{D}}}\left[L(\theta,x)\right]}}{\partial\theta}.$
Without loss of generality we will assume that
$\frac{\partial{E_{x\sim{\mathcal{D}}}\left[L(0,x)\right]}}{\partial\theta}=0$.

The global loss over the sample distribution ${\mathcal{D}}$ is

$\displaystyle L(\theta)$
$\displaystyle=E_{x\sim{\mathcal{D}}}\left[L(\theta,x)\right]$
$\displaystyle=\theta^{T}E_{x\sim{\mathcal{D}}}\left[A_{x}^{T}A_{x}\right]\theta+E_{x\sim{\mathcal{D}}}\left[c_{x}^{T}A_{x}\right]\theta+$
$\displaystyle\phantom{=}+E_{x\sim{\mathcal{D}}}\left[c_{x}^{T}c_{x}\right]$
$\displaystyle=\theta^{T}E_{x\sim{\mathcal{D}}}\left[A_{x}^{T}A_{x}\right]\theta+E_{x\sim{\mathcal{D}}}\left[c_{x}^{T}c_{x}\right]$
(1)

where we take expectation over all possible samples. The last transition holds
since we assumed that the loss function achieves minimum at $0$, and thus
$A_{x}$ and $c_{x}$ satisfy $E\left[A^{T}_{x}c_{x}\right]=0$.

The stochastic gradient and the full gradient can then be written respectively
as:

$\frac{\partial L(\theta,x)}{\partial\theta}={A^{T}_{x}(A_{x}\theta+c_{x})}.$
(2) $\displaystyle\frac{\partial L(\theta)}{\partial\theta}$
$\displaystyle=\frac{\partial
E_{x\sim{\mathcal{D}}}\left[\|A_{x}\theta+c_{x}\|^{2}/2\right]}{\partial\theta}$
(3)
$\displaystyle=E_{x\sim{\mathcal{D}}}\left[A^{T}_{x}(A_{x}\theta+c_{x})\right].$
(4)

Now we estimate how close we can get to $0$ for a fixed learning rate. Note,
even though the individual steps gradients can be assumed to be unbiased
estimators of global loss gradient, there is generally no guarantee that
trajectory will converge to a minimum. Instead we can show that the trajectory
will be traversing an ellipsoid around the minimum.

###### Lemma 1

If $E_{x\sim{\mathcal{D}}}\left[A^{T}_{x}A_{x}\right]=I$ the trajectory will
stabilize at an ellipsoid of a fixed size proportional to the square root of
the learning rate $\sqrt{\lambda}$.

#### Proof:

Consider a fixed sample $x$ at step $t$. Let $b_{x}:=c_{x}^{T}A$.

$\theta_{t+1}=\theta_{t}-\lambda A^{T}_{x}A_{x}\theta_{t}-\lambda b_{x}.$ (5)

We would like to describe the behavior of $\|\theta_{t+1}\|-\|\theta_{t}\|$ as
a function of $\lambda$. First, we observe that
$\lambda\theta^{T}_{t}A^{T}_{x}A_{x}\theta_{t}\geq 0$, with equality achieved
only when $A_{x}\theta_{t}=0$. Thus, the second term in Equation 5 can be
decomposed into a projection onto $\theta_{t}$ and its orthogonal component
using the inner product with $\theta_{t}$:

$\lambda
A^{T}_{x}A_{x}\theta_{t}=\lambda[\gamma_{x}\theta_{t}+\theta^{\perp}_{t,A_{x}}]$

where
$\gamma_{x}:=\frac{\theta^{T}_{t}A^{T}_{x}A_{x}\theta_{t}}{\|\theta_{t}\|^{2}}$,
and $\theta^{\perp}_{t,A_{x}}$ is a vector orthogonal to $\theta_{t}$.
Therefore the update Equation 5 can be rewritten as:

$\theta_{t+1}=(1-\lambda\gamma_{x})\theta_{t}-\lambda[\theta^{\perp}_{t,A_{x}}+b_{x}]=(1-\lambda\gamma_{x})\theta_{t}-\lambda\xi_{x}$
(6)

where

$\xi_{x}:=\theta^{\perp}_{t,A_{x}}+b_{x}.$ (7)

The update to $\theta_{t}$ consists of two terms: (a) the component
$-\lambda\gamma_{x}\theta_{t},$ whose magnitude depends on $A_{x}$ that
reduces norm of $\theta$ and (b) an orthogonal drifting term $\xi_{x}$. The
term $\xi_{x}$ also consists of two components: $\theta^{\perp}_{T,A_{x}}$
which is orthogonal to $\theta_{t}$ by construction, and
$b_{x}=A^{T}_{x}c_{x}$ which we can assume to be orthogonal to $\theta_{t}$
since it is independent of the orientation of $\theta_{t}$. Thus, $\xi_{x}$ is
orthogonal to $\theta_{t}$ and is determined by $A_{x},c_{x}$ and
$\theta_{t}$. These components are illustrated on Figure 4(a). Most
importantly the drift component $\xi_{x}$ is approximately orthogonal to
$\theta$ and its contribution to $\|\theta_{t+1}\|^{2}$ is quadratic in
$\lambda$, while the drag component, has a linear contribution. Let us compute
the change in norms given $x$:

$\displaystyle\|\theta_{t+1}\|^{2}-\|\theta_{t}\|^{2}$
$\displaystyle=(1-\lambda\gamma_{x})^{2}\|\theta_{t}\|^{2}+\lambda^{2}\|\xi_{x}\|^{2}-\|\theta_{t}\|^{2}$
$\displaystyle=-\lambda\gamma_{x}(2-\lambda\gamma_{x})\|\theta_{t}\|^{2}+\lambda^{2}\|\xi_{x}\|^{2},$
(8)

therefore the change in norm becomes zero when

$\lambda=\lambda_{x}:=\frac{2\gamma_{x}\|\theta_{t}\|^{2}}{\|\xi_{x}\|^{2}+\gamma_{x}^{2}\|\theta_{t}\|^{2}}.$
(9)

Now we can estimate the _expected_ change in norm for a fixed $\theta_{t}$
using Equation 8:

$\displaystyle
E_{x\sim{\mathcal{D}}}\left[\|\theta_{t+1}\|^{2}-\|\theta_{t}\|^{2}|\theta_{t}\right]=-2\lambda\|\theta_{t}\|^{2}E_{x\sim{\mathcal{D}}}\left[\gamma_{x}|\theta_{t}\right]$
$\displaystyle\qquad+\lambda^{2}E_{x\sim{\mathcal{D}}}\left[\|\xi_{x}\|^{2}+\gamma_{x}^{2}\|\theta_{t}\|^{2}|\theta_{t}\right].$
(10)

The expected change in norm becomes zero when

$\lambda=2\|\theta_{t}\|^{2}\frac{E_{x\sim{\mathcal{D}}}\left[\gamma_{x}|\theta_{t}\right]}{E_{x\sim{\mathcal{D}}}\left[\|\xi_{x}\|^{2}\right]+\|\theta_{t}\|^{2}E_{x\sim{\mathcal{D}}}\left[\gamma^{2}_{x}|\theta_{t}\right]},$
(11)

or, equivalently, for fixed $\lambda$ the norm has zero expected change

$\displaystyle\|\theta_{t}\|^{2}$ $\displaystyle=\frac{\lambda
E_{x\sim{\mathcal{D}}}\left[\|\xi_{x}\|^{2}|\theta_{t}\right]}{2E_{x\sim{\mathcal{D}}}\left[\gamma_{x}|\theta_{t}\right]-\lambda
E_{x\sim{\mathcal{D}}}\left[\gamma_{x}^{2}|\theta_{t}\right]}$
$\displaystyle=\frac{\lambda
E_{x\sim{\mathcal{D}}}\left[\|\xi_{x}\|^{2}|\theta_{t}\right]}{2E_{x\sim{\mathcal{D}}}\left[\gamma_{x}|\theta_{t}\right]}+O(\lambda^{2}).$
(12)

For the special case when
$E_{x\sim{\mathcal{D}}}\left[A^{T}_{x}A_{x}\right]=I$, we have
$E_{x\sim{\mathcal{D}}}\left[\gamma_{x}|\theta_{t}\right]=I$, and since
$A_{x}$ is independent of $\theta_{t}$ we can assume that
$\theta_{t,A_{x}}^{\perp}$ is orthogonal to $b_{x}$ and its magnitude is
independent of $\theta_{t}$, therefore
$E_{x\sim{\mathcal{D}}}\left[\|\xi\|^{2}|\theta\right]\approx
E_{x\sim{\mathcal{D}}}\left[\|\xi\|^{2}\right]$ and thus

$\|\theta\|=\sqrt{E_{x\sim{\mathcal{D}}}\left[\|\xi\|^{2}\right]}\sqrt{\lambda/2}=\|c\|\sqrt{\lambda/2}.$
(13)

$\blacksquare$

For a fixed learning rate we can expect the distance to the global minimum to
stabilize at a value proportional to $\sqrt{\lambda}$ and described by
Equation 13.

This model turns out of to be very similar to additive batch-noise model of Wu
et al. (2019) and Schaul et al. (2013), however instead of introducing it, we
arrive to it from a empirical assumption that each batch optimizes its own
loss function. Despite its simplicity it appears to capture well many of the
phenomena of large deep neural networks training.

An alternative derivation of the result above is based on Fokker-Planck
equation for describing the evolution of the state density (Sato & Nakagawa,
2014a; Li et al., 2017; Chaudhari & Soatto, 2018) and showing that the
stationary solution of this equation is a Gaussian distribution (for which the
samples are concentrated on a spherical shell). While in general this scaling
factor is opaque, there are several special cases allowing for further
simplification of the equation.

#### Weight averaging: effect on the stationary distribution.

Averaging method | Equivalent learning rate for last $k$ steps
---|---
Stochastic Weight Averaging (Izmailov et al., 2018), window size $k$, cycle $c=1$ | $\lambda(i)\frac{k+1-i}{k}$
Two point average $k$ step apart | ${\lambda(i)}/2$
Exponential Moving Average (decay $\delta$) |  $\lambda(i)(1-(1-\delta)^{k-i})$, where ($k\gg 1/\delta)$
Table 1: Equivalent learning rate schedules for different averaging methods.
Here $i$ denotes the current step of SGD trajectory and $\lambda(i)$ is a
learning rate pre-aggregation. Even though EMA is computed over the entire
trajectory, the contributions of points beyond $1/\delta$ decay exponentially,
thus the equivalent schedule only need to be applied to the last
$k\gg\frac{1}{\delta}$.

As we saw earlier, for a fixed learning rate the solution trajectory
stabilizes an ellipsoid whose size is proportional to $\sqrt{\lambda}$, while
traversing it indefinitely. If we let it continue sufficiently long we end up
with two solutions $\theta_{1}$ and $\theta_{2}$ that are samples on the same
sphere. Therefore the average $\theta=(\theta_{1}+\theta_{2})/2$ has norm
$\|\theta\|=\|\theta\|/\sqrt{2}$. Thus averaging two solutions with a higher
learning rate $\lambda$ gives us practically identical solution as if
following the trajectory from $\theta_{1}$ to $\theta_{2}$, but with the
learning rate $\lambda/2$. Geometrically, the interpolation between two points
on a random trajectory results in a point inside the sphere (Figure 4(b)),
which translates to a lower loss. Remarkably similar effects are also observed
for ImageNet as we show in Figure 5. In Appendix A we generalise this result
to an arbitrary averaging kernel.

We note that previous interpretations proposed in Kingma & Ba (2014) suggest
that weight averaging works by “smoothing” the model. Instead, we show that
there is a natural geometric interpretation that arises from the basic
properties of stochastic gradient descent. Similar intuition can also be
applied to stochastic weight averaging (SWA; Izmailov et al., 2018) and
exponential moving averaging (EMA; Kingma & Ba, 2014).

#### Weight averaging: equivalence to the reduced learning rate along the
trajectory.

As we saw above, aggregation is equivalent to a specific learning schedule in
the stationary regime. Here we consider another angle: we show that a similar
result also holds when we look at window sizes, where we can assume that
gradient on a fixed batch doesn’t change much along the trajectory. We
formalize this result in the lemma below.

###### Lemma 2

(for proof see Appendix A) If gradient of the loss with respect to a fixed
batch is approximately constant, then for SWA with cycle of length $c=1$, two-
point averaging and EMA, the equivalent learning schedule is described in
Table 1.

Note that under our simplifying assumption of matching gradients, we have
shown a much stronger result: not only losses match, but the actual
trajectories match. In the actual training, as can be seen in Figure 7 the
gradients, while staying fairly aligned, do exhibit some amount of divergence.
Further the divergence in $\theta$ space is also significant (though much less
then if batches were sampled independently). Despite that, we can see from
Figure 5 that _loss_ exhibit remarkable match between different averaging
schemas and learning rate schedules. We conjecture that the requirements of
this lemma can be relaxed in favor of loss match, where instead of matching
the full trajectory, the different methods reach different points on the loss
surface.

To recap, in our model we have shown that _both_ stationary regime and early
in the trajectory the weight aggregation has equivalent learning schedules. It
remains a subject of future work to bridge the theory to include the _entire_
trajectory, as appears to be the case for real deep neural networks.

Figure 5: Comparing learning rate schedule with aggregation. The dotted
vertical red lines show the set of independent trajectories “side-trips” with
appropriate learning rate schedule that start at corresponding point in the
main trajectory. For midpoint and average, the red circles at midpoint show
the accuracy at 1, 10, 100 and 1000 steps. For EMA the side-trips are 3000
steps. The alternating dash/dot show the running averages of the main
trajectory (solid line, bottom). Figure 6: Evolution of $L_{2}$ distances
divergence in weight space, the starting point $\theta_{0}$ was selected at
75K, with the starting learning rate at $\lambda=0.05$ and $\lambda/2$.
$\theta^{\prime}$ and $\theta^{\prime\prime}$ correspond to two different
trajectories using differently sampled mini-batches. Note how the loss of the
average point of the trajectory with learning rate $\lambda$ tracks the
trajectory with $\lambda/2$. Figure 7: Evolution of $L_{2}$ gradient
divergence, the starting point $\theta_{0}$ was selected at 75K for different
LR schedules. We compare fixed batch $b$ and training batch $b_{t}$. For fixed
batch the direction of the gradient changes very slowly over 1000 steps (blue
curve), while for different batches ( orange curve) all nearly orthogonal even
at the _same step_. The top most ( green curve) on the left chart is the
$\cos$ of the angle between two gradients at trajectory $\theta_{t}$ \- fixed
learning rate of $0.05$, and $\theta^{\prime}_{t}$ \- the linearly decaying
learning rate $\lambda=0.05\frac{1000-i}{1000}$. The middle graph shows the
change in gradient norm between different schedules. The rightmost graph,
(orange curve) shows the losses for training batch, while the green curve and
blue curve show the loss of a fixed batch $b$ for constant and linearly
decaying schedule respectively. Note that while green and blue curves maintain
good gradient alignment, despite green curve exhibits significant loss change.
Best viewed in color.

##  5 Experiments

In this section we describe our additional experiments. Following our setup
from Section 3, we experiment with ImageNet, but also include results from
Cifar10 and Cifar100 (Krizhevsky, 2009). For all datasets we use identical
setup with ResNet-34, and use SGD with momentum 0.9.

#### Different aggregation and learning rate schedule on ImageNet and CIFAR
datasets

In this experiment we show that learning rate schedules described in Table 1
match stochastic averaging on several large datasets. On Figure 5 we show our
results for the ImageNet. On Figure 5 we show validation accuracy, howeve
similar results also hold for training splits as well for the actual losses,
as we show in supplementary materials Figure 14. Additionally we include the
results on Cifar10 and Cifar100 in the supplementary materials Figure 10.

On Figure 6 we show how the solution for two-point averaging diverges from the
equivalent learning rate schedule. As can be seen, the training loss matches
our estimate fairly close, while arriving at two distinct solutions. Even
though the average point is significantly apart, in relative terms the two
solutions are much closer than two independently trained solution with
identical learning rate.

#### Gradient alignment and divergence of trajectories

As we recall from Lemma 2, we relied on the gradients for a fixed batch
essentially unchanging as we move along the trajectory. On Figure 7 we show
that gradients on the fixed batch $b$ indeed change very little. We consider
two _side-trips_ , using standard sequence of mini-batches, but we measure the
gradient on a _fixed_ mini-batch and compare its direction with the gradient
at the initial step. As can be seen, the gradient norm changed less than 5%,
while $\mathop{cos}$ stayed generally above 0.5. For high dimensional space
this means that the gradient on a fixed batch stays within a very narrow cone
_throughout_ the trajectory. Remarkably, the angle between two parallel points
of two trajectories (green curve) is even closer aligned than the gradient
between the end-point and the start point.

For reference we also include the angle between two _independent_ batches – as
one can see those gradients are typically almost perfectly orthogonal. All
this suggests that we can generally expect the requirements of lemma Figure 7
to hold even on large datasets.

#### Basins of attraction

Figure 8: Distance between intermediate weights $\theta_{t}$ and an
independently trained alternative solution $\theta^{\prime}$

It has been known that multiple trajectories sharing initial segment end up in
the attraction basin as measured by the absence of the loss barrier (Frankle
et al., 2019). Here we show that in addition to being part of the same basin,
the final point of one trajectory and an independent trajectory re-trained
using different sequence of mini-batches monotonically decreases until
saturates at some constant level. On Figure 8 we show how this distance
changes for one such trajectory. This provides further evidence that
independent trajectories land on a sphere around some fixed minimum.

###  5.1 Synthetic model experiments

#### Fast and slow convergence

We have previously shown that the effect of weight averaging can be simulated
approximately by using a properly chosen learning rate schedule. However, in
practice, it is often advantageous to use fixed learning rate throughout
training, especially if the new data is constantly arriving and the model
needs to be trained continuously. In this case, the weight averaging can
improve model performance, while still allowing us to train the model with a
large learning rate guaranteeing fast convergence.

The beneficial effect of weight averaging is particularly pronounced in
situations where training dynamics is characterized by a wide spectrum of time
scales, which is typical in practical scenarios (Ghorbani et al., 2019). This
rich spectrum of timescales manifests itself as the existence of elongated
“trenches” in the loss landscape and long tails characteristic for the
evolution of the loss during training.

Since weight averaging is generally sensitive to the time scale of the
underlying weight evolution, one can expect it to have different effects on
“slow” and “fast” degrees of freedom. As discussed in Appendix A in more
detail, weight averaging will have little effect on the slow dynamics, but is
expected to reduce the size of the stationary distribution for fast degrees of
freedom and reduce the corresponding contribution to the loss as if the
learning rate was in fact smaller. Notice that using smaller learning rate
without weight averaging would have a disadvantage of slowing down the
convergence in slow coordinates and thus using large learning rate with weight
averaging is expected to help us improve model performance without sacrificing
the convergence speed.

We illustrate this intuition by solving equation 5 for
$b_{x}\sim\mathcal{N}(0,\mathbbm{1})$ and a diagonal $\langle
A^{T}_{x}A_{x}\rangle$ with values $1$ and $0.015$ for the fast and slow
degrees of freedom correspondingly. The evolution of the loss function in this
system is shown in Figure 9 for two different learning rates: (a)
$\lambda_{0}=5\cdot 10^{-2}$ and (b) $\lambda_{1}=2\cdot 10^{-2}$. Experiments
with $\lambda=\lambda_{0}$ are conducted both with and without the exponential
moving weight averaging (with a decay rate of $450$ steps). Figure 9
illustrates that there are indeed two time scales in $L(t)$, fast and slow,
and that weight averaging can dramatically reduce the stationary loss $L$
without sacrificing the speed of convergence, unlike when training with the
smaller learning rate $\lambda=\lambda_{1}$. Also, as shown in Figure 11b, in
the appendix, weight averaging has a much larger effect on the fast degrees of
freedom, significantly reducing corresponding contributions to the overall
loss, while having only a mild effect on slow coordinates.

Figure 9:  Loss dynamics in a system governed by with “fast” ($\Omega_{ii}=1$)
and “slow” ($\Omega_{ii}=0.015$) degrees of freedom. Larger
($\lambda_{0}=5\cdot 10^{-2}$) and smaller ($\lambda_{1}=2\cdot 10^{-2}$)
learning rates v.s. EMA;

##  6 Open question and conclusions

We explored some remarkable properties of the learning rate in stochastic
gradient descent, and demonstrated the novel connection between iterate
averaging and learning rate schedules. We showed that this connection can be
observed both in simple theoretical models and in large-scale training on
multiple datasets. We hope that this work paves the way to further
understanding on the role of the learning rate in training. One direction that
we see is developing more general models that would allow to study other
phenomena commonly observed during training, such as overfitting and
trajectory-dependent learning rate.

## References

  * Ainsworth et al. (2022) Samuel K. Ainsworth, Jonathan Hayase, and Siddhartha Srinivasa.  Git re-basin: Merging models modulo permutation symmetries, 2022.  URL https://arxiv.org/abs/2209.04836.
  * Bae et al. (2022) Juhan Bae, Paul Vicol, Jeff Z HaoChen, and Roger Grosse.  Amortized proximal optimization, February 2022.  URL http://arxiv.org/abs/2203.00089.
  * Bengio (2012) Yoshua Bengio.  Practical recommendations for gradient-based training of deep architectures.  _arXiv_ , June 2012.  URL http://arxiv.org/abs/1206.5533.
  * Benzing et al. (2022) Frederik Benzing, Simon Schug, Robert Meier, Johannes von Oswald, Yassir Akram, Nicolas Zucchet, Laurence Aitchison, and Angelika Steger.  Random initialisations performing above chance and how to find them.  _arXiv_ , September 2022.  URL http://arxiv.org/abs/2209.07509.
  * Chaudhari & Soatto (2018) Pratik Chaudhari and Stefano Soatto.  Stochastic gradient descent performs variational inference, converges to limit cycles for deep networks.  In _6th International Conference on Learning Representations, ICLR 2018, Vancouver, BC, Canada, April 30 - May 3, 2018, Conference Track Proceedings_ , 2018.
  * Cohen et al. (2021) Jeremy M. Cohen, Simran Kaur, Yuanzhi Li, J. Zico Kolter, and Ameet Talwalkar.  Gradient descent on neural networks typically occurs at the edge of stability.  _CoRR_ , abs/2103.00065, 2021.  URL https://arxiv.org/abs/2103.00065.
  * Damian et al. (2022) Alex Damian, Eshaan Nichani, and Jason D Lee.  Self-Stabilization: The implicit bias of gradient descent at the edge of stability.  _arXiv_ , September 2022.  URL http://arxiv.org/abs/2209.15594.
  * Frankle et al. (2019) Jonathan Frankle, Gintare Karolina Dziugaite, Daniel M. Roy, and Michael Carbin.  Linear mode connectivity and the lottery ticket hypothesis.  _CoRR_ , abs/1912.05671, 2019.  URL http://arxiv.org/abs/1912.05671.
  * Ghorbani et al. (2019) Behrooz Ghorbani, Shankar Krishnan, and Ying Xiao.  An investigation into neural net optimization via hessian eigenvalue density.  In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), _Proceedings of the 36th International Conference on Machine Learning, ICML 2019, 9-15 June 2019, Long Beach, California, USA_ , volume 97 of _Proceedings of Machine Learning Research_ , pp. 2232–2241. PMLR, 2019\.
  * Goodfellow et al. (2016) Ian Goodfellow, Yoshua Bengio, Aaron Courville, and Yoshua Bengio.  _Deep learning_ , volume 1.  MIT Press, 2016.
  * He et al. (2019) Haowei He, Gao Huang, and Yang Yuan.  Asymmetric valleys: Beyond sharp and flat local minima.  _coRR_ , February 2019.  URL http://arxiv.org/abs/1902.00744.
  * He et al. (2015) Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.  Deep residual learning for image recognition.  _arXiv_ , December 2015.  URL http://arxiv.org/abs/1512.03385.
  * Izmailov et al. (2018) Pavel Izmailov, Dmitrii Podoprikhin, Timur Garipov, Dmitry Vetrov, and Andrew Gordon Wilson.  Averaging weights leads to wider optima and better generalization.  _arXiv_ , March 2018.  URL http://arxiv.org/abs/1803.05407.
  * Kaddour (2022) Jean Kaddour.  Stop wasting my time! saving days of imagenet and bert training with latest weight averaging.  _CoRR_ , 2022.  URL https://arxiv.org/abs/2209.14981.
  * Kingma & Ba (2014) Diederik Kingma and Jimmy Ba.  Adam: A method for stochastic optimization.  _arXiv_ , 2014.
  * Krizhevsky (2009) Alex Krizhevsky.  Learning multiple layers of features from tiny images.  Technical report, U. Toronto, 2009.  URL https://www.cs.toronto.edu/~kriz/learning-features-2009-TR.pdf.
  * Li et al. (2017) Qianxiao Li, Cheng Tai, and Weinan E.  Stochastic modified equations and adaptive stochastic gradient algorithms.  In Doina Precup and Yee Whye Teh (eds.), _Proceedings of the 34th International Conference on Machine Learning, ICML 2017, Sydney, NSW, Australia, 6-11 August 2017_ , volume 70 of _Proceedings of Machine Learning Research_ , pp. 2101–2110. PMLR, 2017.
  * Loshchilov & Hutter (2016) Ilya Loshchilov and Frank Hutter.  SGDR: Stochastic gradient descent with warm restarts.  _coRR_ , August 2016.  URL http://arxiv.org/abs/1608.03983.
  * Ma et al. (2022) Chao Ma, Daniel Kunin, Lei Wu, and Lexing Ying.  Beyond the quadratic approximation: the multiscale structure of neural network loss landscapes.  _coRR_ , April 2022.  URL http://arxiv.org/abs/2204.11326.
  * Mandt et al. (2017) Stephan Mandt, Matthew D Hoffman, and David M Blei.  Stochastic gradient descent as approximate bayesian inference.  _J. of ML Research_ , pp. 1–35, April 2017.  URL http://arxiv.org/abs/1704.04289.
  * Polyak & Juditsky (1992) B T Polyak and A B Juditsky.  Acceleration of stochastic approximation by averaging.  _SIAM J. Control Optim._ , 30(4):838–855, July 1992.  URL https://doi.org/10.1137/0330046.
  * Robbins & Monro (1951) Herbert Robbins and Sutton Monro.  A Stochastic Approximation Method.  _The Annals of Mathematical Statistics_ , 22(3):400 – 407, 1951.  doi: 10.1214/aoms/1177729586.  URL https://doi.org/10.1214/aoms/1177729586.
  * Ruppert (1988) David Ruppert.  Efficient estimations from a slowly convergent robbins-monro process.  Technical report, Cornell University, 02 1988.
  * Russakovsky et al. (2015) Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, Alexander C. Berg, and Li Fei-Fei.  Imagenet large scale visual recognition challenge.  _Int. J. Comput. Vision_ , 115(3):211–252, December 2015.  ISSN 0920-5691.  doi: 10.1007/s11263-015-0816-y.  URL http://dx.doi.org/10.1007/s11263-015-0816-y.
  * Sato & Nakagawa (2014a) Issei Sato and Hiroshi Nakagawa.  Approximation analysis of stochastic gradient langevin dynamics by using fokker-planck equation and ito process.  In _Proceedings of the 31th International Conference on Machine Learning, ICML 2014, Beijing, China, 21-26 June 2014_ , volume 32 of _JMLR Workshop and Conference Proceedings_ , pp. 982–990. JMLR.org, 2014a.
  * Sato & Nakagawa (2014b) Issei Sato and Hiroshi Nakagawa.  Approximation analysis of stochastic gradient langevin dynamics by using Fokker-Planck equation and ito process.  _Proceedings of Machine Learning Research_ , 32(2):982–990, 2014b.  URL https://proceedings.mlr.press/v32/satoa14.html.
  * Schaul et al. (2013) Tom Schaul, Sixin Zhang, and Yann LeCun.  No more pesky learning rates.  In _International conference on machine learning_ , pp. 343–351. PMLR, 2013.
  * Smith (2015) Leslie N Smith.  Cyclical learning rates for training neural networks.  _arXiv_ , June 2015.  URL http://arxiv.org/abs/1506.01186.
  * Wu et al. (2019) Jingfeng Wu, Wenqing Hu, Haoyi Xiong, Jun Huan, Vladimir Braverman, and Zhanxing Zhu.  On the noisy gradient descent that generalizes as SGD.  _arXiv_ , June 2019.  URL http://arxiv.org/abs/1906.07405.
  * Ziyin et al. (2021) Liu Ziyin, Kangqiao Liu, Takashi Mori, and Masahito Ueda.  Strength of minibatch noise in SGD.  _arXiv_ , February 2021.  URL https://openreview.net/pdf?id=uorVGbWV5sw.

Supplementary materials for
“Training trajectories, aggregation and the curious role of the learning
rate”.

##  Appendix A Weight Averaging

In this section, we consider the evolution of $\theta$ governed by

$\theta_{t+1}=\theta_{t}-\lambda\mathrm{A}\,\theta_{t}-\lambda\tilde{b}_{{\mathbb{B}}},$
(14)

a simplified version of the finite-batch version of equation 5:

$\theta_{t+1}=\theta_{t}-\frac{\lambda}{|{\mathbb{B}}|}\sum_{x\in{\mathbb{B}}}A^{T}_{x}A_{x}\theta_{t}-\frac{\lambda}{|{\mathbb{B}}|}\sum_{x\in{\mathbb{B}}}b_{x},$

where ${\mathbb{B}}$ is a batch of samples, $\mathrm{A}=\langle
A^{T}_{x}A_{x}\rangle$ is a matrix and
$\tilde{b}_{{\mathbb{B}}}\sim\mathcal{N}(0,\mathbbm{\mathrm{C}})$ is a
multivariate normal random variable with covariance
$\mathrm{C}=\mathrm{Q}\mathrm{Q}^{\top}$.

Here, we first show how equation 14 can be reduced even further by using a
coordinate transformation that whitens the batch noise. We then consider a
moving average of the training trajectory, derive a simple equation governing
it and characterize the resulting converged steady-state distribution.

###  A.1 Simplifying gradient descent equation

Stochastic gradient descent trajectories following equation 14 depend on both
$\mathrm{A}$ and $\mathrm{C}$ matrices. Performing a change of coordinates
$\theta_{t}=\mathrm{Q}u_{t}$, we can remove the dependence on one of the
matrices if $\mathrm{Q}$ is an invertible matrix such that
$\mathrm{Q}\mathrm{Q}^{\top}=\mathrm{C}$. Indeed, noticing that
$\tilde{b}_{\mathbb{B}}$ can be represented as $-\mathrm{Q}\xi_{\mathbb{B}}$
with $\xi_{\mathbb{B}}\sim\mathcal{N}(0,1)$, we obtain:

$\mathrm{Q}u_{t+1}=\mathrm{Q}u_{t}-\lambda\mathrm{A}\mathrm{Q}u_{t}+\lambda\mathrm{Q}\xi_{\mathbb{B}},$

and finally get a simpler description of the stochastic gradient descent
trajectory:

$u_{t+1}=(\mathrm{I}-\lambda\Omega)u_{t}+\lambda\xi_{t},$ (15)

where our new equation depends on only one matrix
$\Omega:=\mathrm{Q}^{-1}\mathrm{A}\mathrm{Q}$.

The solution of equation 15 can be obtained recursively and reads:

$\displaystyle
u_{t+1}=(\mathrm{I}-\lambda\Omega)^{t+1-\tau}u_{\tau}+\lambda\sum_{t^{\prime}=\tau}^{t}(\mathrm{I}-\lambda\Omega)^{t-t^{\prime}}\xi_{t^{\prime}}.$
(16)

###  A.2 Moving average of the SGD trajectory

In Section 4, we presented a simple intuitive explanation of the fact that the
two-point average evolves similarly to an ordinary SGD trajectory with an
effective learning rate of $\hat{\lambda}=\lambda/2$. Here we prove a more
general result for an arbitrary averaging procedure.

Consider a training trajectory $u_{t}$ solving equation 15 and let
$\hat{u}_{t}\equiv\sum_{k=0}^{\infty}\mu_{k}u_{t-k}$ be some average of the
training trajectory with the real averaging kernel
$\\{\mu_{k}\in\mathbb{R}|k\in 0,\dots,\infty\\}$. It is easy to see that
$\hat{u}_{t}$ satisfies the following equation:

$\displaystyle\hat{u}_{t+1}=(\mathrm{I}-\lambda\Omega)\hat{u}_{t}+\lambda\sum_{k=0}^{\infty}\mu_{k}\xi_{t-k}.$
(17)

Introducing $\nu_{t}\equiv\sum_{k=0}^{\infty}\mu_{k}\xi_{t-k}$, we see that
this random process is characterized by $\langle\nu_{t}\rangle=0$ and an
autocorrelation
$C_{t,\delta}\equiv\langle\nu_{t}\nu_{t+\delta}\rangle=\sum_{k=0}^{\infty}\mu_{k}\mu_{k+\delta}$
independent of time $t$, where the averaging is performed over different
realizations of the batch noise $\xi$.

The solution of equation 17 can be obtained by analogy with how solution 16
was obtained for equation 15:

$\displaystyle\hat{u}_{t}=\Gamma^{t}\hat{u}_{0}+\lambda\sum_{\tau=1}^{t}\Gamma^{t-\tau}\nu_{\tau},$

where we introduce $\Gamma\equiv\mathrm{I}-\lambda\Omega$. Choosing $t=T$ for
some sufficiently large $T$ and assuming that all eigenvalues $\kappa_{k}$ of
$\Gamma$ are characterized by $0<\kappa_{k}<1$, the first term ends up being
exponentially small $\sim\exp(-\kappa_{\rm max}T)$ and

$\displaystyle\hat{u}_{T}\approx\lambda\sum_{\tau=0}^{T-1}\Gamma^{\tau}\hat{\nu}_{\tau},$

where $\hat{\nu}_{\tau}\equiv\nu_{T-\tau}$.

Now let us look at the distribution $p(\hat{u}_{T})$ for a fixed time $T$ and
different realizations of $\nu$. It is easy to see that
$\langle\hat{u}_{T}\rangle=0$ since $\langle\hat{\nu}_{\tau}\rangle=0$ for
each $\tau$. Furthermore,

$\displaystyle\langle\hat{u}_{T}\hat{u}_{T}\rangle=\lambda^{2}\left\langle\sum_{\tau,\tau^{\prime}=0}^{T-1}\Gamma^{\tau+\tau^{\prime}}\hat{\nu}_{\tau}\hat{\nu}_{\tau^{\prime}}\right\rangle=\lambda^{2}\sum_{\tau,\tau^{\prime}=0}^{T-1}\Gamma^{\tau+\tau^{\prime}}\left\langle\hat{\nu}_{\tau}\hat{\nu}_{\tau^{\prime}}\right\rangle.$

Notice that here we average over different realizations of $\nu_{t}$ and
assume homogeneity of $\nu_{t}$ in time since $C_{t,t^{\prime}-t}$ is
independent of $t$ and only depends on the difference $t^{\prime}-t$. Assuming
that $T$ is sufficiently large for both $\kappa_{\rm max}T\gg 1$ and
$C_{t>T}\ll C_{0}$ to hold, we can approximate (changing the integration
bounds where $\Gamma^{\tau+\tau^{\prime}}C_{\delta}$ is vanishingly small):

$\displaystyle\langle\hat{u}_{T}\hat{u}_{T}\rangle\approx\lambda^{2}\left(\sum_{\tau=0}^{T-1}\Gamma^{2\tau}C_{0}+2\sum_{\delta=1}^{T-1}\sum_{\tau=0}^{T-1}\Gamma^{2\tau+\delta}C_{\delta}\right).$

Here the coefficient of $2$ emerges because $C_{\delta}=-C_{\delta}$ for
$\delta>0$. This can also be rewritten as (replacing $T$ with $T+1$):

$\displaystyle\langle\hat{u}_{T+1}\hat{u}_{T+1}\rangle\approx\lambda^{2}X_{T}\left(C_{0}+2\sum_{\delta=1}^{T}C_{\delta}\Gamma^{\delta}\right),$
(18)

where $X_{T}\equiv\sum_{\tau=0}^{T}\Gamma^{2\tau}$.

Now let us look at the limit of $T\to\infty$, assuming for simplicity that
$C_{\delta}$ is bounded as $\delta\to\infty$ and assuming that the eigenvalues
of $\Gamma=\mathrm{I}-\lambda\Omega$ are all smaller than $1$. The expression
for $X_{\infty}$ can be simplified by relating it to Taylor series for
$1/(1-x)$, specifically notice that:

$\displaystyle(\mathrm{I}-\Gamma^{2})X_{\infty}=(\mathrm{I}-\Gamma^{2})\sum_{\tau=0}^{\infty}\Gamma^{2\tau}=\sum_{\tau=0}^{\infty}\Gamma^{2\tau}-\sum_{\tau=1}^{\infty}\Gamma^{2\tau}=\mathrm{I}$

and therefore $X_{\infty}=(\mathrm{I}-\Gamma^{2})^{-1}$. We can then
approximate equation 18 as follows:

$\displaystyle
F\equiv\lim_{T\to\infty}\langle\hat{u}_{T+1}\hat{u}_{T+1}\rangle\approx\lambda^{2}(\mathrm{I}-\Gamma^{2})^{-1}\left(C_{0}+2\sum_{\delta=1}^{\infty}C_{\delta}\Gamma^{\delta}\right),$

or recalling that $\Gamma=\mathrm{I}-\lambda\Omega$:

$\displaystyle
F\approx\lambda(2\Omega-\lambda\Omega^{2})^{-1}\left(C_{0}+2\sum_{\delta=1}^{\infty}C_{\delta}\Gamma^{\delta}\right).$
(19)

This final expression for $F$ connects the covariance of the stationary
distribution in the weight space to the autocorrelation of the averaging
kernel $\mu_{t}$. One important conclusion is that since
$\mathrm{I}-\lambda\Omega$ is contracting (assuming $\Omega$ is positive-
definite), the long tail of $C_{\delta}$ can be effectively cancelled by
$(\mathrm{I}-\lambda\Omega)^{\delta}$. In other words, if the characteristic
averaging time exceeds $1/\kappa$, where $\kappa$ is the eigenvalue of
$\mathrm{I}-\lambda\Omega$, then the corresponding width of the stationary
distribution will be inhibited by weight averaging.

###  A.3 Two-point average

The final expression 19 can then be applied to an arbitrary averaging
procedure. For example, the following result generalizes the geometric
derivation for the two-point average to an arbitrary111$\lambda$ still needs
to be sufficiently small for system dynamics to not diverge learning rate
$\lambda$ and the distance $\Delta$ between steps.

Since we define the averaging procedure via $\mu_{0}=1/2$ and
$\mu_{\Delta}=1/2$ (and $\mu_{k}=0$ otherwise), we obtain $C_{0}=1/2$ and
$C_{\Delta}=1/4$. Substituting this expression into equation 19, we obtain:

$F_{\Delta}\equiv\lim_{T\to\infty}\langle\hat{u}_{T}\hat{u}_{T}\rangle\approx
S_{\lambda}\frac{\mathrm{I}+(\mathrm{I}-\lambda\Omega)^{\Delta}}{2},$

where $S_{\lambda}\equiv\lambda(2\Omega-\lambda\Omega^{2})^{-1}$. This
expression can also be interpreted as:

$\hat{\lambda}=\lambda\frac{1+(1-\lambda\kappa)^{\Delta}}{2},$ (20)

where $\hat{\lambda}$ is the effective learning rate for a degree of freedom
corresponding to an eigenvector of $\Omega$ with an eigenvalue of $\kappa$.
This expression generalizes our previous observation that
$\hat{\lambda}\approx\lambda$ for sufficiently small $\Delta\sim 1$ and
$\hat{\lambda}\approx\lambda/2$ for $\Delta\gg 1$ when
$(1-\lambda\kappa)^{\Delta}\approx 0$.

###  A.4 Multi-point average

Now let us study weight averaging over $n$ points separated by $\Delta$ steps
each, i.e.,
$\hat{u}_{t}=n^{-1}(u_{t}+u_{t-\Delta}+u_{t-2\Delta}+\dots+u_{t-(n-1)\Delta})$.
Since this corresponds to $\mu_{k\Delta}=1/n$ for $k=0,\dots,n-1$, we obtain:

$\displaystyle
C_{k\Delta}=\sum_{\tau=0}^{\infty}\mu_{k\Delta}\mu_{(k+\tau)\Delta}=\frac{n-k}{n^{2}}.$

Substituting this expression in equation 19, we obtain:

$\displaystyle
C_{0}+2\sum_{\delta=1}^{\infty}C_{\delta}\Gamma^{\delta}=\frac{1}{n}+\frac{2}{n}\sum_{k=1}^{n-1}\left(1-\frac{k}{n}\right)\Gamma^{k\Delta}=\frac{1}{n}+\frac{2}{n}\left(G_{1}-\frac{G_{2}}{n}\right),$
(21)

where

$\displaystyle G_{1}\equiv\sum_{k=1}^{n-1}\Gamma^{k\Delta},\qquad
G_{2}\equiv\sum_{k=1}^{n-1}k\Gamma^{k\Delta}.$

It is not difficult then to verify that:

$\displaystyle
G_{1}=(\mathrm{I}-\Gamma^{\Delta})^{-1}\Gamma^{\Delta}\left(\mathrm{I}-\Gamma^{(n-1)\Delta}\right),$
$\displaystyle
G_{2}=(\mathrm{I}-\Gamma^{\Delta})^{-2}\Gamma^{\Delta}\left(\mathrm{I}-n\Gamma^{(n-1)\Delta}+(n-1)\Gamma^{n\Delta}\right).$

These expressions for $G_{1}$ and $G_{2}$ together with equation 19 and
equation 21 fully define the covariance $F$ for arbitrary values of $\lambda$,
$\Delta$ and an arbitrary $\Omega$.

As before, in the limit of $\Delta\to\infty$ when $\Gamma^{\Delta}\to 0$, we
see that

$\displaystyle
F\approx\frac{S_{\lambda}}{n}\left(\mathrm{I}+\frac{2(n-1)}{n}\Gamma^{\Delta}\right)+o(\Gamma^{\Delta}).$

In other words, as one would expect, the effective learning rate is decreased
by approximately a factor of $n$.

###  A.5 Proof of Lemma 2

#### Proof:

We have
$\theta_{t_{2},\lambda(t)}=\theta_{t_{1}}+\sum_{i={t_{1}}}^{t_{2}}\lambda(i)\xi_{x_{i}}(\theta_{i})$,
where $\xi_{x_{i}}$ is a stochastic gradient of batch $x_{i}$ with respect to
$\theta_{i}$. thus

$\frac{\theta_{t_{2}}+\theta_{t_{1}}}{2}=\theta_{t_{1}}+\sum_{i={t_{1}}}^{t_{2}-1}\frac{\lambda(i)}{2}\xi_{x_{i}}\approx\theta_{t_{2},\lambda(t)/2},$

where we assumed that for fixed $x_{i}$, the gradient $\xi_{x_{i}}$ changes
very slowly and thus $\xi_{x_{i},\lambda(t)}\approx\xi_{x_{i},\lambda/2(t)}$.

Similarly for stochastic weight averaging we have

$\displaystyle\theta_{\text{swa}}$
$\displaystyle=\frac{1}{t_{2}-t_{1}}\sum_{t=t_{1}}^{t_{2}}\theta_{t}$ (22)
$\displaystyle=\frac{1}{t_{2}-t_{1}}\sum_{t=t_{1}}^{t_{2}}\left[\theta_{t_{1}}+\sum_{i=t_{1}}^{t-1}\lambda(i)\xi_{x_{i}}\right]$
(23)
$\displaystyle=\theta_{t_{1}}+\sum_{t=t_{1}}^{t_{2}-1}\lambda(t)\frac{t_{2}-t}{t_{2}-t_{1}}\xi_{x_{t}}\approx\theta_{t_{2},\frac{t_{2}-t}{t_{2}-t_{1}}\lambda(t)}$
(24)

Finally, for exponential moving average we have:

$\displaystyle\theta_{\text{ema},t_{2}}$ $\displaystyle=$
$\displaystyle\delta\sum_{t=0}^{t_{2}}\theta_{t}(1-\delta)^{t_{2}-t}$ (25)
$\displaystyle\approx$
$\displaystyle=\delta\sum_{t=t_{1}}^{t_{2}}\theta_{t}(1-\delta)^{t_{2}-t}$
(26) $\displaystyle=$
$\displaystyle\delta[\theta_{t_{1}}\sum_{t=t_{1}}^{t_{2}}(1-\delta)^{t_{2}-t}+$
(28)
$\displaystyle+\sum_{t=t_{1}}^{t_{2}-1}\lambda(t)\xi_{x_{t}}\sum_{j=t}^{t_{2}-1}(1-\delta)^{t_{2}-j-1}]$
$\displaystyle\approx$
$\displaystyle\theta_{t_{1}}+\sum_{t=t_{1}}^{t_{2}-1}\lambda(t)\left(1-(1-\delta)^{t_{2}-t}\right)\xi_{t}$
(29)

where the first transition holds because we assume $t_{2}-t_{1}\gg 1/\delta$,
and thus contribution of terms up-to $t_{1}$ is negligible. $\blacksquare$

##  Appendix B Additional Experiments

Comparison of a two-point average and an equivalent learning schedule is shown
in Figure 10.

(a)

(b)

Figure 10: Two-point averaging on validation data for CIFAR10 and CIFAR100

(a)

(b)

Figure 11:  Loss dynamics in a system governed by with “fast”
($\Omega_{ii}=1$) and “slow” ($\Omega_{ii}=0.015$) degrees of freedom. (a)
larger ($\lambda_{0}=5\cdot 10^{-2}$) and smaller ($\lambda_{1}=2\cdot
10^{-2}$) learning rates vs. Exponential Moving Average; (b) the effect of
exponential moving averaging of the model weights for fast and slow degrees of
freedom.

###  B.1 Does single-step behavior findings generalize to different batch
sizes, or at different trajectory

In this section we show that our observations from the main experimental hold
both at different stages of trajectory and for different batch sizes. On
Figure 12 we show how the loss changes in a single step when performing this
step at different location of SGD trajectory and using different batch size.
On Figure 13 we show the maximum loss that is achievable if we pick a step
size that would pick “optimal” step. In other words the points correspond to
minima of each line of Figure 12.

(a)

(b)

(c)

(d)

Figure 12: Loss change as a function of step size starting off different
points of SGD trajectory, and using different batches. Figure 13: Maximum loss
drop (log scale) for single step, when performed at different stages of SGD
trajectory. As one can see even for batch size that are in the typical range
(256-4096), training batch drops several orders of magnitude more than held
out batch, however as the batch size reaches larger values such as 32768 the
difference decreases.

###  B.2 Additional experiments on Imagenet

On Figure 14 we show that sidetrip with appropriate learning rate schedule
v.s. aggregation holds for actual loss and the accuracy, as well as for both
training and test data.

Figure 14: Accuracy and loss for training and validation splits on Imagenet.
The bottom row shows the trajectory for cosine learning rate schedule, which
achieves state of the art accuracy for Resnet-34 architecture. Each graph
shows momentary v.s. Stochastic Averaging (SWA) and side-trips (red) with
appropriate learning rate schedule.