On the generalization of learning algorithms that do not converge
Author s)
Abstract
Common generalization analyses of deep learning assume that the training converges to a fixed point.
But, in practice, the weights of deep neural networksoptimized with stochastic gradient descent often oscillate indefinitely.
To reducethis discrepancy between theory and practice, we analyze the generalization ofalgorithms when the weights only converge in distribution.
Our main contributionis to propose a notion of statistical algorithmic stability (SAS) that extends classicalalgorithmic stability to non-convergent algorithms and to study its connection togeneralization.
This ergodic-theoretic approach to learning theory leads to newinsights when compared to the traditional optimization and learning theory perspectives.
We prove that the stability of the time-asymptotic behavior of a learningalgorithm can be used to predict its generalization and empirically demonstratehow loss dynamics can provide clues of overfitting.
Our findings provide evidencethat networks that “train stably generalize better” even when the training continuesindefinitely and the weights do not converge.
Introduction
It is common practice that when the training loss of a neural networks converges close to somevalue, the learning algorithm—typically some variant of Stochastic Gradient Descent (SGD)—isterminated.
Perhaps surprisingly, recent works indicate that the network function at termination timeis typically not a fixed point of the learning algorithm: if run longer, the gradient norm does notvanish and the learning algorithm outputs functions with significantly different parameters [Cohenet al., 2021, Zhang et al., 2021b, Lobacheva et al., 2021].
This observation stands in contrast tocommon generalization analyses taking into account the learning algorithm that assume convergenceto a fixed point.
It also raises the question of when (and why) non-convergent learning algorithmsshould be expected to generalize.
Standard approaches for obtaining generalization bounds find ways to bound the complexity measureof the function class expressible by a neural network in a data-independent [Vapnik, 1999, Bartlettet al., 2017, Neyshabur et al., 2018, Golowich et al., 2018, Bartlett et al., 2019, Arora et al., 2018]or data-dependent manner [von Luxburg and Bousquet, 2004, Xu and Mannor, 2012, Sokoli´c et al.,2017].
The list of works is too long to exhaustively cover here.
Nevertheless, even with recenttechnical improvements, these approaches do not explicitly account for the relationship betweengeneralization and the dynamics of the learning algorithm.
A second approach to generalization analysis yields algorithm-dependent bounds by connecting thegeneralization performance of the trained network to the stability of the training to data perturbations[Bousquet and Elisseeff, 2002, Feldman and Vondrak, 2018, Kuzborskij and Lampert, 2018, Bousquetet al., 2020, Zhang et al., 2021c, Rakhlin, 2006].
However, these bounds do not exploit dynamicalinformation about the learning algorithm and only apply to algorithms that converge to a fixed point.
One line of inquiry into algorithm-dependent generalization has hence focused on settings wherethe algorithm provably converges to a fixed point, such as two-layer overparameterized networks,deep linear networks, certain matrix factorization etc.
[Frei et al., 2019, Allen-Zhu et al., 2019,
Soudry et al., 2018, Arora et al., 2019, Gunasekar et al., 2018, Haeffele and Vidal, 2015].
In amore general setting, Hardt et al.
[2016] prove algorithmic stability-based generalization bounds thatworsen with increasing training time whereas Loukas et al.
[2021] connect stable training dynamicsnear convergence under Dropout with good generalization.
For algorithms that do not converge(close) to a fixed point, the above analyses result in vacuous bounds.
Contributions.
We start with the supposition that robustness of the learning algorithm’s fixed pointsto small changes in the training set is not the mechanism underlying the good generalization propertiesof deep networks.
This is because, as is commonly known and many recent works indirectly remarkupon [Cohen et al., 2021, Ahn et al., 2022, Zhang et al., 2021b, Garipov et al., 2018], training can belinearly unstable : even if assuming the presence of isolated local minima that the learning algorithmconvergences to, the weights encountered in training are not robust to small perturbations.
Smallperturbations can accumulate over time leading to completely different orbits in weight space.
Yet,despite unstable dynamics, the generalization properties of two completely different fixed points (e.g.,obtained from two different random initializations or due to a different order of SGD updates) aretypically comparable.
Our first contribution entails proposing a generalization of stability that applies to non-convergingalgorithms.
To achieve this, we depart from the typical optimization perspective of analyzing thegeneralization properties of minima [Keskar et al., 2016, Ge et al., 2015, Sagun et al., 2016].
Instead,we study the generalization performance of the learning algorithm on average (see section 2.2).
Adopting a statistical viewpoint, in Section 3 we define a new notion of statistical algorithmicstability that measures the stability of this average performance to perturbations of the training data,as opposed to stability of individual orbits in the weight space.
Section 3 then proves upper boundsfor the generalization error based on our new notion of statistical stability.
We then move on to consider how statistical stability connects to the behavior of the loss function.
To this end, in Sections 2.2 and 4 we provide conditions such that the spectral gap of a Markovoperator associated with the dynamics of the loss function is predictive of statistical algorithmicstability and hence of generalization.
Empirically, we estimate this spectral gap based on the rateof convergence of ergodic averages of the loss function.
Our numerical results demonstrate that ourestimate correlates with generalization for image classification under corruption.
1.1 Related work
Our work builds on previous analyses of the stability of learning algorithms that converge [Bousquetand Elisseeff, 2002, Feldman and Vondrak, 2018, Kuzborskij and Lampert, 2018, Bousquet et al.
,2020, Zhang et al., 2021c].
Below, we also briefly discuss connections (beyond the literature onstability) with previous analyses on the dynamics of learning algorithms.
Convergence and generalization of SGD.
Exponential convergence rates to global minima areknown for SGD in the smooth and convex setting, with decaying learning rates and small constantlearning rates (e.g., [Moulines and Bach, 2011, Needell et al., 2014]).
In the training of deep neuralnetworks, which is the focus of this paper, the loss function is non-convex and is typically notwell-approximated locally by convex functions [Ma et al., 2018].
However, in the overparameterizedregime, a Polyak-Lojasiewicz-type inequality is automatically satisfied when the loss is smooth andconvex, and can be shown to hold in the non-convex setting as well [Liu et al., 2022, Ma et al.
, 2018],allowing convergence to a global minimum.
In certain non-convex problems, convergence to localminima and even global minima is proved, assuming a strict saddle property [Ge et al., 2015] or localconvexity [Kleinberg et al., 2018].
Raginsky et al.
[2017] show a favorable convergence for a slightlydifferent version of SGD by approximating the dynamics with a suitable continuous time Langevindiffusion.
A number of previous works [Safran and Shamir, 2016, Freeman and Bruna, 2016, Li and
Yuan, 2017, Nguyen et al., 2018] have also provided arguments that SGD converges to a solution thatgeneralizes given suitable initialization and significant overparameterization.
Another line of workconsiders the mean field limit of SGD dynamics to show generalization [Mei et al., 2019, Gabriéet al., 2018, Chen et al., 2020].
The ergodic theoretic perspective we adopt here deviates from theoptimization perspectives but relates to the mean field perspective in that we implicitly consider (seesection 4) the evolution of probability distributions on weight space.
Flat minima.
In the generalization literature, flat minima correspond to large connected regions withlow empirical error in weight space.
Flat minima have been argued to connect with the complexityof the resulting hypothesis and, hence, can imply good generalization Hochreiter and Schmidhuber[1997].
It has also been shown that SGD converges more frequently to flat minima when the batchsize is small or the learning rate and the number of iterations are suitably adjusted [Keskar et al.
,2016, Hoffer et al., 2017, Jastrz˛ebski et al., 2017, Smith and Le, 2018, Zhang et al., 2018].
Someworks consider the flat/sharp dichotomy an oversimplification
[Dinh et al., 2017, Sagun et al., 2017, He et al., 2019] using e.g. reparameterization arguments Sagun et al.
The eigenvalues ofthe hessian of the loss affect the local linear stability of optimization orbits.
However, given that weadopt the time-asymptotic/statistical picture, we do not explicitly make assumptions on local stability(flatness/sharpness) or the global topology of the loss landscape [Montanari and Zhong, 2020, Venturiet al., 2019].
SGD as a Bayesian sampler.
The idea of looking at distributions of learners permeates PAC-Bayesiananalyses of generalization [McAllester, 1999, Dziugaite and Roy, 2017, Zhou et al., 2019, Pitas, 2020].
In contrast to Bayesian posteriors of the parameters, we study parameter distributions generated bythe learning algorithms.
Previous empirical work has suggested that SGD operates almost like a
Bayesian sampler [Mingard et al., 2021]
but the connection remains not fully understood.
Local descent algorithms as dynamical systems: a statistical viewpoint
This section lays out the basic definitions and assumptions of our work.
We begin in Section 2.1 byintroducing the learning problem and discussing the learning algorithms from a dynamical systemsperspective.
Section 2.2 then puts forth the notions of statistical convergence that give rise to ourmain results.
2.1 Learning as a dynamical system
We consider the supervised learning setup: a learner is given a training set S
=
{ z 1 , . . .
, z n }consisting of n pairs z i ≡ ( x i , y i ) of inputs x
i ∈ R d and their corresponding labels y i ∈ Y , drawni.i.d.
from a distribution D .
A class of parametric models or hypotheses h : R d × W → R on aparameter space W gives possible input-to-output relationships, h ( w, · ) : R d → Y .
The learner isgiven a risk or loss function ℓ : ( R d × Y ) × W
→ R + that describes the error of a given hypothesis.
A common choice of the loss function is the mean-squared, hinge, or cross-entropy loss.
The learnerattempts to minimize the population risk R S ( h ( ·
[Equation] w [Equation]
z
[Equation]
ℓ ( z, w ) , by minimizing the empiricalrisk
[Equation]
This minimization is achieved using a local descent algorithm given by iterative updates ϕ S :
W → Wof the form
[Equation]
where η t is the learning rate or step size at time t.
In the gradient descent (GD) algorithm, ˆ ∇ L S ( w ) = ∇ L S ( w ) is the gradient w.r.t.
w , the iterates ofϕ S represent a deterministic dynamical system.
The critical points (including saddle points, localand global minima) of L S are fixed points of ϕ S (i.e., points w such that ϕ S w = w ).
For stochasticgradient descent (SGD), ˆ ∇ L S ( w ) is a random variable given by the sample mean of gradients∇ ℓ ( · , w ) over a batch of samples in S .
In this case, ϕ S is a random dynamical system and criticalpoints of L S are not necessarily fixed points of ϕ S .
To unify the definition of ϕ S for both GD and SGD, we introduce the random variable Ξ that indicatesthe choice of batch.
Suppose the batch size is fixed: m < n for SGD and m = n for GD.
Denoteby [ n ] the set { 1 , 2 , · · · , n } and let Ξ t be a collection of m elements from [ n ] chosen uniformly atrandom.
Then, we may write ϕ S as
[Equation]
40 20 0 20 40       40 loss 0.
[Figure]
[Figure]
[Figure] [Figure]
In our investigation, we opt for simplicity and consider a fixed learning rate η t = η .
More generally,our analysis can be extended to learning rates that eventually converge to η .
A rapidly vanishinglearning rate can obscure the learning dynamics by enforcing convergence to arbitrary points irrespective of the points’ loss.
In contrast, learning rates that do not decay to zero can induce interestingtransitions between neighborhoods of critical points and the iterates w t are generally more exploratoryof the loss landscape.
2.2 Statistical convergence of the learning algorithm
In deep learning, the choice of w T is often to some extent arbitrary as the loss landscape is generallynon-convex and stochastic learning algorithms are the norm.
Hence, rather than focusing on a singlehypothesis, it is attractive to consider the average generalization error induced by the stochasticdynamical system ϕ S .
We adopt this statistical perspective, rather than focus on a single orbit ofϕ S .
Thus, the setting below applies to any asymptotic behavior of orbits of the (typically nonlinear)dynamics ϕ S including fixed points, periodic, quasiperiodic and chaotic orbits.
To analyze the generalization performance without making assumptions on the global stability properties of individualorbits, we make an assumption about the ergodic properties of ϕ S that reasonably matches empiricalevidence.
To describe the assumption, we give a short primer on the relevant concepts from ergodictheory of dynamical systems.
Invariant measures.
Denote the Borel sigma algebra on W by B ( W ) .
A probability measure µ S iscalled invariant for ϕ S if µ S ( A ) = µ S ◦ ϕ − 1 S ( A ) for all Borel sets A ∈ B ( W ) .
Intuitively, for anysubset of the weight space, the probability that the dynamical system occupies it does not change.
See Liverani [2004] for an introduction to invariant measures.
Can we expect the dynamics of learning algorithms to give rise to invariant measures?
If ϕ S is adeterministic continuous function on the set W and W is compact, the Krylov–Bogolioubov theorem(see e.g., Theorem 4.1.1 of Katok and Hasselblatt [1997]) gives the existence of at least one invariantmeasure for ϕ S .
On the other hand, when ϕ S is an SGD update, Dieuleveut et al.
[2020] showthe existence of invariant distributions for convex loss functions, whereas Fort and Pagès [1999]and Mattingly et al.
[2002] show the existence of invariant distributions in more general settingsunder mild assumptions on the gradient noise (in the estimate ˆ ∇ L S ) .
Several works have analyzedthe asymptotic properties of such invariant distributions, e.g., Chee and Toulis [2018] analyze theoscillation of the iterates w t about a mean that is a critical point, as a phase separate from the transientphase of convergence to this invariant distribution.
In general, there are infinitely many invariant measures.
A pertinent class of invariant measures inpractice is that of ergodic, invariant measures, where we can only understand the statistical propertiesof a system through evolution of individual orbits.
Roughly speaking, for an ergodic, invariant measure, sets that are invariant under the dynamics ϕ S (such as a set of periodic points, quasiperiodicor chaotic attractors) are trivial: they have measure 0 or 1.
Ergodicity.
When an invariant measure µ S is ergodic for ϕ S , time averages converge to ensembleaverages according to µ S .
That is, for µ S -almost every initial state w 0 , and for all continuous scalarfunctions h :
[Equation]
and the dependence on initial conditions w 0 is forgotten.
Unfortunately, we cannot generally expect the dynamics of the learning algorithm to have a uniqueergodic invariant measure on W .
Instead, ϕ S -invariant sets could be smaller subsets of W withdifferent ergodic invariant measures supported on them.
As a result, starting from two different pointschosen Lebesgue almost everywhere (or sampled from any probability density on W ), we do notexpect time averages to converge to the same limit.
Indeed, as shown in Figure 1, the weights visitedby ϕ S depend on the initial conditions and can vary significantly across different runs.
In other words,the limit of infinite time averages described in Equation (4) does depend on w 0 .
Dynamics of the hypothesis.
To circumvent the above issues, rather than focusing on the weightsdirectly, we shift perspective and consider the dynamics of the learned hypothesis h ( · , w t ) .
Specifically, we posit that the infinite-time averages of the learned hypothesis are ergodic: that is, the timeaverages of h ( z, w t ) converge to the same function irrespective of the initial function h
( z, w 0 ) .
Asa result, functionals that depend only on h and not explicitly on the weights are also ergodic.
Ourmain assumption is that the time averages of loss functions are ergodic, since they depend only on thenetwork functions:
Assumption 1.
Given any S ∼ D n , there exists a map z → ⟨ ℓ z ⟩ S , such that for Lebesgue-a.e. w 0and every z ∈ R d × R , the following holds:
[Equation]
where { w t = ϕ S ( w t − 1 ) } is an orbit of ϕ S .
As an intuitive justification for Assumption 1, symmetries may result in several sets of parameters(weights and biases) that represent the same network function h. Hence, assuming that there existsa unique probability on the space of neural network functions that is ergodic is less restrictive andcloser to reality than assuming a unique ergodic measure on the parameter space.
Figure 1 provides evidence that Assumption 1 can hold in practice.
It shows two observablescalculated along 4 different orbits of a VGG16 model trained using SGD with momentum (batchesof size 128, learning rate of 0.01, momentum of 0.9).
On the left, we can see a low-dimensionalprojection of the second layer weights obtained by projecting the weights onto 40 dimensions using
PCA and then employing t-SNE [Van der Maaten and Hinton, 2008], whereas the observable on theright is the test accuracy, which is a functional of the hypothesis function.
Here, the test accuracy canbe seen to converge to a distribution independent of the initial conditions, whilst the weight orbitsvary significantly across different initializations.
This experiment corroborates our hypothesis thattime averages of functionals on the loss space converge to distributions independent of the initialcondition (i.e., unique ergodic invariant measures), even if time averages of a generic observable onthe weight space do not.
Remark 1.
Increasing learning rate may induce ergodicity.
Assumption 1 may hold for all continuousfunctions for sufficiently large constant learning rates, as long as orbits do not diverge.
Intuitively,large learning rates can cause more frequent transitions from the basin of attraction of one localminimum to another.
On the other hand, for small learning rates, we may be able to detect thepresence of multiple attractors, since different initial conditions lead to different long-time averages.
These ideas are illustrated with a toy example in Appendix A.
Statistical stability implies generalization
3.1 Statistical algorithmic stability
Classically, the derivations of algorithmic stability-based generalization (see e.g., Hardt et al.
Chapter 14 of Mohri et al.
[2018], Bousquet et al.
[2020]) utilize an input perturbation of one elementin S by replacing it with a different element from the input distribution D :
Defnition 1 (Algorithmic stability, adapted from [Bousquet and Elisseeff, 2002]) .
Consider theweights w ∗ S and w ∗ S ′ obtained by running the learning algorithm on two training sets S , S ′ sampledfrom D that differ by exactly one sample.
We say that the learning algorithm is algorithmically stable(AS) with a stability coefficient β ≥
if
[Equation]
We refer to this type of perturbation as a stochastic perturbation.
Algorithmic stability does notdirectly apply to learning algorithms that do not converge to a fixed point.
Hence, next, we extendalgorithmic stability to loss statistics .
The resulting notion of statistical algorithmic stability (SAS)extends algorithmic stability to algorithms whose loss statistics converge as prescribed by Assumption1, even if the weights do not converge.
Defnition 2 (Statistical algorithmic stability) .
We say that a learning algorithm with loss statisticsdenoted by ⟨ ℓ z ⟩ S is statistically algorithmically stable (SAS) with a stability coefficient β ≥
if
[Equation]
7) for all S , S ′ sampled from D that differ by exactly one point.
In the above definition, ⟨ ℓ z ⟩ S ′ refers to the ergodic average of ℓ ( z, · ) observed along almost everyorbit of ϕ S ′ .
The definition of SAS differs from Definition 1 as we do not assume convergence of thelearning algorithm.
Nevertheless, when every orbit of ϕ S converges to a fixed point w ∗ S , for almostevery S, then Definition (7) reduces to the standard notion of algorithmic stability since ergodicaverages of ⟨ ℓ z ⟩ S along almost every orbit converge to ℓ ( z, w ∗ S ) .
Crucially, and in line with the observations in Figure 1, we quantify stability based on the statistics ofthe loss function after it has statistically converged to some ergodic ϕ S -invariant measure µ S and notthe ensemble mean or any one point in the weight space W .
In other words, the definition of SASdoes not use or provide information about the algorithmic stability of E w ∈ µ S w for any µ S .
3.2 Learning theoretic implications
It is well-known that Definition 1 leads to generalization bounds (see [Hardt et al., 2016, Bousquetet al., 2020] and references therein).
Next, we show that the more general statistical algorithmicstability from Definition 2 also can imply a notion of generalization.
To obtain generalization bounds, we first redefine the empirical and population risk to use the ergodicaverage loss statistics rather than fixed points:
[Equation]
Although the definitions above appear deceivingly similar with the standard ones introduced in
Section 2.1, the two notions of risk are defined on different spaces, with the standard ones beingmaps from the hypothesis space H to scalar values and (8) being maps from the space of learningalgorithms .
Using these definitions, the following generalization bound holds:
Theorem 1.
Given a β -statistically algorithmically stable algorithm ϕ S (see Definition 2), for anyδ ∈ (0 , 1) , with probability greater than 1 − δ ,
[Equation]
where L := sup z sup w ∈ W | ℓ ( z, w ) | is an upper bound on the loss.
The proof can be found in
Appendix B. Like [Bousquet and Elisseeff, 2002]’s bound
, the stabilitycoefficient needs to grow at most like β ∼ O (1 / √ n ) for the generalization bound to be non-vacuous.
Note that our proof is also naturally coupled with [Bousquet and Elisseeff, 2002, Feldman and
Vondrak, 2018]’s improved technique, which provides tighter bounds when β ∼ O (1 / √ n ) .
(see
Appendix B).
Dynamical interpretation of SAS
Having defined the notion of statistical algorithmic stability, the natural next question is when analgorithm ϕ S is SAS.
Moreover, with computable timeseries data, can we determine SAS?
The latterquestion is, of course, of immense practical importance because S-SAS is sufficient for generalization(via Theorem 1).
We now delineate one route, taking an operator-theoretic perspective to provide amechanistic understanding of what makes an algorithm SAS.
The theoretical explanation we developyields a rough heuristic – albeit not a definitive predictor of generalization – that relates a giventraining dynamics to the SAS coefficient β .
Since ergodic, invariant measures on the weight spaceare not unique, it may seem daunting at first to obtain an estimate of the SAS coefficient.
However,assuming that we have a unique ergodic, invariant measure on the loss space leads us naturally to acomputable proxy for β, as we show below:
Lemma 1.
(A Markov loss process) Let µ S be an ergodic, invariant measure for ϕ S .
Assume that aloss function ℓ ( z, · ) :
W → I z ⊆ [0 , L ] is such that its level sets form a smooth foliation of the supportof µ S .
Let ν
zS
=
ℓ
( z, · )
♯ µ S be the pushforward measure on the loss space, i.e., ν zS : B ( I z )
→ R + ,where B ( I z ) is the Borel sigma algebra on I z .
Then, there exists a mixing Markov operator P zµ S onthe space of probability measures on B ( I z ) , with an invariant measure ν zS .
This lemma parallels Chekroun et al.
[2014]’s Theorem A, but in a quite different setting, withoutusing existence of a unique physical measure.
The proof, deferred to Appendix C, involves usingthe assumptions to define a transition probability operator on the loss space using the transitionprobability between level sets of ℓ ( z, · ) on the weight space.
The main consequence of this lemma isthat the Markov process on the loss space associated with the operator P zµ S gives us a mechanisticunderstanding of SAS, as we discuss below.
First, using Assumption 1, it is clear that the pushforward measure ν S defined in Lemma 1 isindependent of µ S .
In other words, for any ergodic, invariant measure µ S , we have a unique ergodic,invariant measure ν zS = ℓ ( z, · )
♯ µ S .
Since this measure is unique, by construction in Lemma 1, itis also the stationary distribution associated with the loss process { ℓ ( z, w t ) } t .
Note that the lossprocess is, in general, non-Markovian.
But, we are able to leverage the perturbation theory of Markovprocesses to study the effect of stochastic perturbations on ν zS , thanks to Lemma 1.
Theorem 2.
If the mixing operator P zµ S constructed in Lemma 1 is uniformly ergodic, and the
Lipschitz constant of ∇ ℓ ( · , w ) on R d is bounded above by L D for all w ∈ W , then ϕ S is SAS with stability coefficient
[Equation]
where λ = sup
z λ
z is the supremum over z of the mixing rates of P zµ S and α
=
O
Proof.
(Sketch)
Following Chekroun et al.
[2014], the proof relies on the perturbation bounds of
Mitrophanov [2005].
If P zµ S is uniformly ergodic [Mitrophanov, 2005], there exists a λ z < 1 and C > 0 such that for all ξ ∈ I z ,
[Equation]
Further, if a stochastic perturbation S ′ introduces a change δ P zµ S to P zµ S , the perturbation boundfrom [Mitrophanov, 2005] gives
[Equation]
(11) where α :=
⌊ log 1 /λ z C ⌋ and ∥ δ P zµ S ∥ is the operator norm of δ P zµ S induced by TV distance.
Now,since |⟨ ℓ z ⟩ S − ⟨ ℓ z ⟩ S ′ | =
(cid:12)(cid:12)(cid:12) E ν
zS ξ − E ν
zS ′ ξ (cid:12)(cid:12)(cid:12) ≤
L ∥
ν zS − ν zS ∥ TV , we obtain
[Equation]
Since this holds for all z, the right hand side of (12) is an upper bound for the stability coefficient βin
Definition (2), by taking the supremum over z .
It is easy to see that if Ξ t is uniform, the change inthe Markov operator on the weight space can be bounded above by m/n in the TV-induced norm.
This bound means that the stability upper bound (12) is O (1 /n ) for small batch sizes.
For GD, theupper bound on (cid:13)(cid:13) δ P zµ S (cid:13)(cid:13) is vacuous.
However, when the Lipschitz constant of ∇ ℓ ( · , w ) is uniformly(in w ) small on R d (discussed further in Appendix C), the size of δ P zµ S can still be controlled suchthat
β =
O (1 /n ) .
Theorem 2 implies that an algorithm that exhibits faster convergence to the stationary measure on theloss space generalizes better .
The bound in Theorem 2 is sensitive to changes in λ, when the latter isclose to 1.
This implies that smaller λ (faster convergence of the loss to the ergodic invariant measure)yield smaller upper bounds on β (more statistically stable algorithm) and better generalization.
Strictly speaking, to determine λ z , we must consider finite-dimensional approximations of theoperator P zµ S (as done in the setting of chaotic systems in e.g., [Crimmins and Froyland, 2020,
Chekroun et al., 2014]) and compute its spectral decomposition.
However, in our setting, it is difficultin practice to generate samples in the loss space that obey the Markov process defined in Lemma 1.
Instead, we argue for using the readily available non-Markovian loss process to estimate convergenceto the equilibrium distribution.
In particular, we hypothesize below that the auto-correlation functionof the test loss can separate algorithms based on their SAS.
Dropping the superscripts z , we define the auto-correlation C ℓ ( τ ) in a loss function ℓ by(cid:12)(cid:12) (1 /T )
(cid:80) t< = T ℓ ( w t ) ℓ ( w t + τ ) − ⟨ ℓ ⟩ 2 (cid:12)(cid:12) / ⟨ ℓ ⟩ 2 , in the limit T → ∞ .
That is, C ℓ ( τ ) gives the correlation between the random variables ℓ ( w t ) / ⟨ ℓ ⟩ and ℓ ( w t + τ ) / ⟨ ℓ ⟩ , where the randomness comes fromΞ t and the initialization of weights in SGD and only from the latter in GD.
Once the distribution onthe space of ℓ has converged to ν S , i.e., for sufficiently large τ, C ℓ ( τ ) gives the correlation betweentwo dependent samples from ν S .
This dependence is not lost since the process is non-Markovian.
Now, for small times τ , even the Markov process from Lemma 1 on loss space has not mixedcompletely.
This also makes C ℓ ( τ ) larger because the loss process mixes at least as slowly as the
Markov process.
In short, C ℓ encodes information both about speed of convergence to equilibriumdistribution and the time-dependence that arises due to the dynamics (non-Markovian effect).
Thus, ifthe non-Markovian effects in the loss are comparable between two different algorithms, higher valuesof C ℓ indicate larger λ, and hence less statistical stability.
Numerical results
We numerically validate the main ideas of section 3 and 4 on VGG16 and ResNet18 models trainedon a CIFAR10 dataset (see Appendix D for further numerical results).
For all our experiments, ϕ S is an SGD update with momentum 0.9, fixed learning rate 0.01 and batch size of 128.
In all figures,time indicates number of epochs.
We generate different choices of the training set S p by corrupting
CIFAR10’s labels with probability p, with S 0 being the original CIFAR10 dataset.
Figures 2 and3 show results corresponding to p = 0 , 0 . 1 , 0 .
25 and 0 .
As expected, the greater the noisecorruption p, the larger is the generalization error, as also reported by previous studies Zhang et al.[2021a], Loukas et al.
The phenomenon is shown in Figure 2 (left), where the test losstimeseries (an estimate of the generalization error) is plotted for these different cases.
Each line in Figure 2 is a sample mean over 10 random initializations.
In Figure 3, we depict also (lower bounds of) the SAS coefficient β
(see Definition 2).
The cumulativetime average along orbits of ϕ S p of length 1200 epochs are used as estimators for statistics.
Figure3 shows the difference in the test loss cumulative averages between random orbits of ϕ S p and ϕ S ′ p ,with S
p being a stochastic perturbation of S p .
The errorbars indicate the standard error over 45 pairsof orbits.
In Figure 3(right), we compute an estimate of the stability coefficient as the difference ofthe (estimated) test loss statistics at different values of p using 45 pairs of orbits for each p.
This isan estimate on the lower bound of β, and clearly increases with p, illustrating that cases with worsegeneralization errors (Figure 2(left)) have larger β.
In Figure 2(right), we plot C ℓ ( τ ) (section 4) where ℓ is taken as the test loss.
The test loss statistic⟨ ℓ ⟩ is estimated as a time average over 1200 epochs.
At each p, we show the sample average of thenormalized autocorrelations over 10 independent runs.
Since we have normalized by ⟨ ℓ ⟩ , the y-axistruly extracts correlations over a long time range (up to 700 epochs) in a function of L 1 norm 1 at allp .
Thus, Figure 2 (right) illustrates that the correlations at p = 0 are smaller when compared to higherp ′ s not because the loss itself has smaller magnitude, but because the convergence to equilibriumdistribution on the loss space is faster.
We see that the magnitude of the test loss autocorrelationscorrelates well with the generalization gap (absolute difference in test and training losses) shown in
Figure 2(center).
Overall, the loss process does indeed codify the phenomenological explanation for
SAS that is expressed in 12.
Discussion and conclusion
Predicting generalization is an active area of research and has implications for more reliable use ofmachine learning.
In this work, we introduce statistical algorithmic stability (SAS) and show how itimplies generalization bounds.
Here, we add a few important remarks.
Statistical stability only requires robustness of statistics on loss space.
A central challenge in theanalyses of non-convergent algorithms is the fact that there may be multiple, very different invariantmeasures µ S .
For this reason, our stability criterion focuses on statistics of the loss function.
Thisway, an algorithm can be stable even if the measures µ S and µ S ′ are not close in the total variationnorm.
In fact, an algorithm can be stable even if µ S and µ S ′ are mutually singular, if the loss functionshave similar statistics.
Since our SAS analysis hinges on a reasonable yet nonrestrictive assumptionon ergodic properties of the algorithm, we believe it is broadly applicable.
Algorithms that converge are special cases of the above analysis.
In Appendix E, we show that theproposed relationship (section 4) between statistical stability and convergence rates to stationarymeasures can also be observed in the Neural Tangent Kernel (NTK) regime
[Jacot et al., 2018].
Inthis case, the training dynamics, ϕ S , can be approximated by a linear function of the weights, whichconverges to a fixed point.
Thus, this provides an alternative, ergodic theoretic, interpretation ofgeneralization in the NTK regime [Arora et al., 2019, Montanari and Zhong, 2020, Bartlett et al.,2021].
Broader view.
While our focus is on algorithmic stability-based generalization, this work gathersmore evidence to support the broader view [Wojtowytsch, 2021, Zhang et al., 2021b] that exploitingtheoretically and empirically available dynamical information about the training algorithm is a fruitfulcomplement to understanding generalization from the optimization landscape and learning theoryperspectives.