On the generalization of learning algorithms that do not converge
Anonymous 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 to generalization performance.
Our findingsprovide evidence that networks that “train stably generalize better” even when thetraining continues indefinitely 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., 2019b, 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 complexity ofthe 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].
Some works consider theflat/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 of the hessianof the loss affect the local linear stability of optimization orbits.
However, given that we adoptthe 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 ( · , w ))
=
E 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, m, 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]
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 stabilityproperties of individual orbits, we make an assumption about the ergodic properties of ϕ S thatreasonably matches empirical evidence.
To describe the assumption, we give a short primer on therelevant concepts from ergodic theory 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]
) 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 ′ .
A higher value of β indicates a lower SAS algorithm.
The definition of SAS differsfrom Definition 1 as we do not assume convergence of the learning algorithm.
Nevertheless, whenevery orbit of ϕ S converges to a fixed point w ∗ S , for almost every S, then Definition (7) reduces tothe standard notion of algorithmic stability since ergodic averages of ⟨ ℓ z ⟩ S along almost every orbitconverge 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] in the total variation (TV) norm.
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 .
5 use VGG16 and ResNet18 modelsrespectively.
As expected, the greater the noise corruption p, the larger is the generalization error, asalso reported by previous studies Zhang et al.
[2021a], Loukas et al.
The phenomenon isshown in Figure 3 (left), where the test loss timeseries (an estimate of the generalization error) isplotted for these different cases.
Each line in Figure 3 is a sample mean over 10 random initializations.
In Figure 2, 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.
Figure2 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 2(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 is anestimate on the lower bound of β, and clearly increases with p, i.e, with increasing generalizationerrors (Figure 3(left)).
In Figure 3(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 3 (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) shownin Figure 3(center).
Curiously, since both test error and training error are larger when p = 0 .
5 incomparison with smaller p ′ s the generalization gap is smaller, as shown in Figure 3 (center) (see also Appendix D).
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., 2019b, 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.
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’scontributions and scope?
[Yes](b) Did you describe the limitations of your work?
[Yes](c)
Did you discuss any potential negative societal impacts of your work?
[Yes](d) Have you read the ethics review guidelines and ensured that your paper conforms tothem?
[Yes]2.
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[Yes](b) Did you include complete proofs of all theoretical results?
[Yes]3.
If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[Yes] In thesupplementary material(b) Did you specify all the training details (e.g., data splits, hyperparameters, how theywere chosen)?
[Yes] see section 5 and Appendix D(c)
Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[Yes] See section 5 and Appendix D(d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)?
[No]4.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators
[Table] [Table] [Table]
any
[Table]
either
[Table]
[Table]
[N
[Table]
[Table]
[Table]
using/
[N
[Table]
[Table]
[Table]
[Table]
[Table] [Table]
any potential
risks
[Table] [Table] [Table] [Table]
Did you
[Table]
the estimated hourly wage paid to participants and the total
spent
on participant compensation?
[N
[Table]
A Bifurcation analysis of smooth and non-convex optimization
In this section, we discuss some examples of non-convex optimization in one dimension performedwith gradient descent (GD).
We illustrate that, with increasing learning rate, the asymptotic behaviorof orbits may alter from being periodic to quasiperiodic to chaotic.
These qualitative changes arebrought about through period-doubling bifurcations , which are observed in many physical systems(e.g.
, [Quail et al., 2015, Aron and Schwartz, 1984, Zhao et al., 2004]). We consider smooth andnon-convex objective functions of the form ℓ s = g s ◦ g s ◦ g s , with g s being the canonical quadraticmap g s ( w ) : = 1 − sw (1 − w ) of the unit interval [0 , 1] .
Since g s is smooth, the compositions of g swith itself are smooth. The first bifurcation point appears just above η = 2 / | d 2 ℓ s /dw 2 | , which is thestability threshold for convex optimization [Nesterov, 2003]. In Figure 4, the left, center and right columns correspond to s = 1 , 3 and 4 respectively. In thefirst row, we plot the loss functions ℓ s , which are non-convex (with multiple global minima) ats = 3 and s = 4 .
The second row shows the sharpness – absolute value of the second derivative,a ( w ) =
| d 2 ℓ s ( w ) /dw | . At s = 1 , the unique global minimum at w = 0 . 5 is flatter than the globalminima at s = 3 and s = 4 .
The third row of Figure 4 is a bifurcation diagram, which shows the attractor on the y -axis as afunction of the learning rate.
The attractor is approximated by the asymptotic orbits of the dynamicsat multiple (100) different initial conditions chosen randomly on the unit interval.
Note that at smallvalues of η < 2 / ∥ a ∥ with ∥ a ∥ := sup w ∈ [0 , 1] a ( w ) orbits from different initial conditions converge tofixed points corresponding to the local/global minima at each value of s . Periodic orbits emerge atη > 2 / ∥ a ∥ , that are ultimately shown to become chaotic for larger learning rates. This can be notedfrom the last row of Figure 4, where Lyapunov exponents (see e.g., [Katok and Hasselblatt, 1997, Wilkinson, 2017]) computed at different initial conditions are plotted.
Given the gradient descentdynamics,
[Equation]
(13) where f ′ ( w ) = ( df/dw )( w ) , the Lyapunov exponent , λ s : [0 , 1] → R is defined as, [Equation] Roughly speaking, this function λ s ( w ) measures the asymptotic stability of infinitesimal linearperturbations along the orbit of w.
Here, we use the exponential notation, ϕ ts w = ϕ s ◦ ϕ t − 1 s w todenote compositions of ϕ s . One definition of chaotic orbits is that the Lyapunov exponent is positive. We see from the bottom row of Figure 4 that in the range of learning rates considered, chaos isobserved (positive Lyapunov exponents starting from uniformly random initial conditions) in thecase of sharper minima, for larger learning rates.
In ergodic systems, λ s ( w ) is independent of w. Indeed, for larger learning rates, we see that λ s appears to be independent of the initial conditions(the bottom row of Figure 4 shows λ s ( w ) for 100 different w ). On the other hand, smaller learningrates where convergence to fixed points or periodic behavior is observed, the Lyapunov exponentconverges to different negative values depending on the initial condition. This example supportsthe heuristic explanation for the primary assumption made in the main text about the asymptoticdynamics of learning algorithms.
In this section, we prove Theorem 1 from the main text. This result says that the statistical stability ofan algorithm implies generalization. Recall that here, the population risk is defined as
[Equation]
and the empirical risk is
[Equation] We closely follow the proof strategies of Bousquet and Elisseeff [2002] (see also [Bousquet et al. Define a function Φ S : = R S − ˆ R S , whose expected value is
[Equation]
Examining the second term, since z i ’s are chosen i.i.d.
according to D , E S ∼D n ⟨ ℓ z i ⟩ S is constantacross i and equal to E S ∼D n ⟨ ℓ z 1 ⟩ S , where z 1 ∈ S.
As in the main text, let S ′ denote any set thathas at most one element different from S, i.e., a stochastic perturbation of S. Using the stochasticperturbation S ′ , we can rewrite E S ∼D n ⟨ ℓ z 1 ⟩ S as E S ∼D n E z ∈D ⟨ ℓ z ⟩ S ′ .
Substituting this equivalentexpression in (14), and using the fact that ϕ S is SAS with stability coefficient β :
[Equation]
The proof is based on applying Mcdiarmid’s inequality (see e.g., Bousquet et al.
[2020]) to Φ S , whichgives a high-probability upper bound on Φ S − E (Φ S ) in terms of the deviation of Φ S from a Φ S ′ .
Inorder to obtain an upper bound for the latter quantity, note that
[Equation] Considering the difference of the population risks, [Equation] [Equation]
Next considering the difference of empirical risks,
[Equation]
[Equation]
where, recall that L := sup z ∈ R d × R sup w ∈ W | ℓ ( z, · ) | .
Putting (17) together with (18) into Inequality(16), we obtain | Φ S − Φ S ′
| ≤ 2 β + 1 n (2 L − β ) . Lemma 2. (Mcdiarmid’s inequality) Let Z 1 , · · · , Z n be random variables taking values in Z .
Afunction f : Z n → R is said to satisfy the bounded differences property if there exists a constantc > 0 such that
[Equation]
[Equation]
for all single coordinate changes.
For an f that satisfies the bounded differences with a constantc > 0 , given a δ > 0 , with probability at least 1 − δ,
[Equation]
Applying Mcdiarmid’s inequality to Φ S , and recalling (15)
, we obtain the generalization bound statedin Theorem 1 of the main text.
That is, with probability ≥ 1 − δ,
[Equation]
The above proof may also be repeated using the concentration inequality obtained in Theorem 4 of
Bousquet et al.
This leads to a tighter bound analogous to Corollary 8 of Bousquet et al.[2020].
C Predicting generalization with loss timeseries
In this section, we substantiate the connection drawn in section 4 of the main text between the rateof decay of correlations in the loss function and SAS.
The study of dynamics lifted to the spaceof observables [Koopman, 1931] has a substantial precedent in dynamical systems theory (e.g.,[Crimmins and Froyland, 2020, Keller and Liverani, 1999, Dellnitz et al., 2000]) and computationalmethods Arbabi and Mezic [2017], Budiši´c et al.
[2012], Williams et al.
[2014], Korda and Mezi´c[2018].
In particular, the idea of relating the correlation decay rate or the convergence rate of Fokker-Planck/Frobenius Perron operators with some notion of global sensitivity has also extensivelyappeared both in the statistical learning theory and SDE literature (see e.g. [Bartlett et al., 2021, Sirignano and Spiliopoulos, 2022]) and in the dynamical systems literature (e.g., [Kato, 2013, Chekroun et al., 2014]).
In this work, inspired by the existence of such relationships in variouscontexts, we are able to show that the autocorrelations in the loss function can serve as predictors forthe generalization gap.
This section provides the complete setting for the results in section 4 of the main text.
First we definethe transition operator for probability distributions on the weight space W and analogous operator onthe loss space.
Markov operator for weight space.
Recall that Ξ t is a batch of m indices chosen uniformly fromthe set [ n ] := { 1 , · · · , n } .
For GD, Ξ t is deterministic and equal to the set [ n ] .
Let (Ω , Σ , P ) be aprobability space.
Let K S : W × B ( W ) →
R + be the Markov kernel associated with the update ϕ S ,i.e., for a Borel subset A ∈ B ( W ) and a point w ∈ W ,
K S ( w
[Equation]
where P ( ϕ S ( w ) ∈
A | Ξ t ) is the probability of the event that ϕ S ( w )
∈
A when the m indices Ξ t arechosen.
Correspondingly, we may define the Markov operator P S , also called the Frobenius-Perronoperator Lasota and Mackey [1998], on the space of probability measures on W ,
[Equation]
From the above definition, it is clear that any ϕ S -invariant probability measure µ is an eigendistribution of P S with eigenvalue 1.
In our setting (see section 2 of the main text), there are potentiallymultiple eigendistributions corresponding to eigenvalue 1.
Each invariant measure µ S also definesdifferent transition probabilities on the weight space W :
[Equation]
Markov operator for the loss.
Now, instead of the phase space W , for each z, consider the image,
I z ≡ ℓ ( z, W ) ⊆ [0 , L ] of W under ℓ ( z, · ) , with L as defined in Theorem 1.
Let B ( I z ) denote the
Borel sigma algebra on I z .
Analogous to the kernel (21), we may now define a kernel K zµ S , for an( ξ, E ) ∈
I z × B ( I z ) , as
[Equation]
This, in turn, gives rise to a Markov operator analogous to the Frobenius-Perron operator (22) on thefull-dimensional weight space:
[Equation]
[Equation]
Note that this operator is well-defined when the level sets of ℓ z are µ S -measurable.
One sufficientcondition for this is when the foliation of W by these level sets is subordinate to a measurablepartition.
We may then consider disintegrations of µ S on this measurable partition and define thekernel (24) using conditional measures supported on elements of the partition.
Further, it is clear that
P zµ S satisfies the properties of a Markov operator (positive unity preserving contraction).
Thus, in order to prove Lemma 1, it remains to show that the operator defined by (25) is mixing.
Forthis, we make an additional assumption.
We assume that the Frobenius-Perron operator P S mixes tothe measure µ S starting from any measure of the form µ = ℓ ( z, · ) − 1 ν, for an absolutely continuousprobability measure ν on [0 , L ] . That is, [Equation] where ∥ · − ∗ ∥ TV indicates the total variation distance and ζ ∈ (0 , 1) is the rate of mixing.
On theother hand, since P tS µ → µ S weakly, ν zt := (cid:0) P zµ S (cid:1) t ℓ ( z, · ) ♯ µ → ν S weakly on B ( I z ) .
Intuitively,we expect that rate of mixing of the latter, say λ z ∈ (0 , 1) correlates with ζ since
[Equation]
Assuming that W is a Polish space, the above relationship (27) conveys that when a measure µ tconverges to µ S in the TV norm, expectations with respect to µ t of all continuous functions, ofwhich { ℓ ( z, · ) } z is a subset, also converge to expectations with respect to µ S .
Thus, intuitively weexpect that when sup z λ z is a lower bound for ζ ((26)).
Our setting is different from previous works([Chekroun et al., 2014] and references therein) in the dynamics literature that use observable-specific
Markov operators, in that we do not assume uniqueness of the ergodic, invariant measure on thefull-dimensional (weight) space.
Effect of stochastic perturbations In section 4 of the main text, we use the perturbation theory ofmixing Markov operators to relate the SAS coefficient to the mixing rate of P zµ S .
This perturbationbound (from [Mitrophanov, 2005]) is given in terms of the TV norm of the perturbation δ P zµ Sto the operator P z µ S when a stochastic perturbation S → S ′ is applied to the training set.
Herewe discuss the size of δ P zµ S , elaborating on the proof of Theorem 2 in the main text.
First notethat ∥
δP zµ S ∥ ≤ ∥P S − P S ′ ∥ , and hence it suffices to obtain an upper bound for ∥P S − P S ′ ∥ .
Theperturbation to Markov kernel in the weight space due to a stochastic perturbation is given by
[Equation]
[Equation]
where ∆ is the set of m indices from [ n ] that contain k, the index at which S and S ′ differ.
When Ξis a uniform random variable as we have assumed, the cardinality ∆ is (cid:0) n − 1 m − 1 (cid:1) / (cid:0) nm (cid:1) = m/n.
Hence,we obtain the following bound on the Markov kernel perturbation,
[Equation]
Using the definition 22,
[Equation]
which is O (1 /n ) for small batch sizes as claimed in the main text, and vacuous for GD.
On the otherhand, when ϕ S is deterministic, (28) becomes
[Equation]
30) which leads to the following upper bound on the perturbation size in the Wasserstein norm,
[Equation]
[Equation]
In order to apply this bound, we use a result from the perturbation theory of Markov chains from[Rudolf and Schweizer, 2018]
(Corollary 3.2 in the asymptotic limit) to replace the TV distancebased bound used in section 4.
Thus, we obtain again, an upper bound for β
=
O ( L D / (1 − λ ) n ) asclaimed in Theorem 2.
D Stability experiments on ResNets
We obtain similar results for SAS with the ResNet18 model as with the VGG16 model shown in
Figure 3 of the main text.
In Figure 5, we plot the difference in the cumulative average of the test loss at runs with the ResNet18 architecture.
The difference is taken between two SGD runs with the sameparameters as in section 5 of the main text and with training data that are stochastic perturbations ofeach other.
We consider 45 pairs of stochastically perturbed datasets (see section 2 of main text fordefinition of stochastic perturbation) each for each value of p .
The value of p = 0 , 0 . 1 , 0 . 17 , 0 . 25 , 50indicates the probability of error injected into the labels of the CIFAR10 dataset. The mean of theabsolute difference in the cumulative test loss is shown in dark colors while the standard error in meanin the corresponding lighter color. The time averages are calculated over 1200 epochs after a runup time of 200 epochs.
The results indicate that greater the noise probability p, greater the estimateof SAS (less statistical stability).
Hence, these results indicate that statistical algorithmic stability(see section 2) correlates with the generalization performance of the models at the difference labelcorruption (noise) levels.
Furthermore, considering together with the VGG16 results presented in themain paper, this relationship between SAS and generalization holds across the different architectureswe have employed in this paper.
The generalization gap plotted in Figure 3 (center) (of the main text) is the cumulative time average(ergodic average) of the absolute difference between the test and training errors.
This is approximatelyequal to (strictly, an upper bound for) | R S − ˆ R S | .
A loose upper bound for this quantity comes fromthe theoretical generalization bound (in Theorem 1).
In practice, this quantity is estimated to be smallrelative to the test error with corrupt datasets.
When the noise probability is 0% (original dataset), thegap estimate is about 80% of the test error because the training error is small.
But, when the noiseprobability is 50%, the gap estimate is about 20% of the test error since the training error is alsolarge.
We may not observe this reduction in the generalization gap if the errors (test and training)were defined using pointwise values.
That is, since, sup w | E z ∼ D ℓ z ( w ) − (1 /n ) (cid:80) ni
=1 ℓ z i ( w ) | ≥| R S − ˆ R S | , upon early stopping when training error is low, we may not observe this phenomenon.
wemay not observe this reduction in the generalization gap if the errors (test and training) were definedusing pointwise values.
That is, since, sup w | E z ∼ D ℓ z ( w ) − (1 /n ) (cid:80) ni
=1 ℓ z i ( w ) | ≥ | R S − ˆ R S | , uponearly stopping when training error is low, we may not observe this phenomenon.
Finally, we remarkthat our empirically estimated autocorrelation function serves as a proxy until a more sophisticatedmethod for the estimation of λ is developed.
As mentioned in the main text (section 4), this is achallenging problem that is beyond the scope of this work.
E Revisiting stability in the linear regime
Since its introduction by Jacot et al.
[2018], training in the Neural Tangent Kernel (NTK) regime hasbeen analyzed thoroughly in numerous works, wherein its convergence to kernel ridge regression hasbeen formally proved in two-layer, infinitely wide networks [Montanari and Zhong, 2020, Bartlettet al., 2021], infinitely wide fully connected networks [Arora et al., 2019a], convolutional networks
Arora et al.
[2019b] and so on.
Using its equivalence with kernel ridge regression, the generalizationproperties under the NT regime have also been well-studied (see Bartlett et al.
[2021] for a review).
Here, we revisit the NT regime with the purpose of demonstrating that the analysis in the presentpaper also holds when ϕ S achieves a fixed point.
The Neural Tangent Kernel (NTK) model is an approximation of a neural network whose parameters remain close to initialization during training.
Given that training in the NT regime iswell-approximated by the dynamics of linear regression (see Theorem 3.1 and 3.2 of Arora et al.[2019b]), in order to apply our dynamics-based generalization analyses to the NTK regime, we needonly consider linear regression dynamics.
That is, let ˜ w t be an orbit of a perturbed dynamics (thelinear regression dynamics, see Lemma 1 of [Arora et al., 2019b]) close to an orbit w t =
ϕ S
w t − 1 )for all time, so that, for some small ϵ > 0 ,
[Equation]
As mentioned in the main text (section 3), SAS reduces to the standard notion of algorithmic stability(see
e.g., [Mohri et al., 2018] chapter 14 for a survey) when the dynamics converges to a fixed point.
Now suppose that the dynamics ˜ w ∗ is algorithmically stable.
That is, for all stochastic perturbations S ′ of S, and for all z, [Equation] Assuming Lipschitz loss with Lipschitz constant C z Lip and denoting by C Lip = sup
z C z Lip
[Equation]
That is, the NTK orbit is stable with the stability coefficient β + 2 C Lip ϵ.
Thus, in order to prove thealgorithmic stability of the NTK orbit, it is enough to consider the stability of the linear regressionorbit, which is a linear dynamical system as we describe below.
The second main idea that we develop is the prediction of the stability coefficient by the rate of decayof correlations.
In this case, we show that the speed of convergence to the fixed point determines thegeneralization properties via algorithmic stability.
We derive stability-based generalization bounds,an alternative to Rademacher complexity-based bounds in Arora et al.
[2019a].
Comparing withexisting generalization results in the well-understood NT regime is an ideal test bed for the alternativedynamical perspective of this present work.
Linear dynamics Let w r ∈ W be the parameter at which a NN h NN ( · , w r ) is the zero function from 853 R d to R , i.e., h NN ( x, w r ) = 0 , for all x. Now consider training the weights w ∈ W of the NN, withthe initialization w 0 = w r .
With GD on the squared loss L S ( w ) =
{+(1 / 2) (cid:80) n i =
+}
( y i − h NN
( x i , w )) 2 ,and learning rate η > 0 , the dynamics of the weights are as follows,
[Equation]
Here, X =
[ x 1 , · · · , x n ] T ∈
R d × n and the notation h NN ( X, w ) ∈
R n represents[ h NN ( x 1 , w ) , · · · , h NN ( x n , w )] T .
Note that the above dynamics ϕ S ( w ) is a nonlinear function of w.
Now we consider the NTK setting described in Bartlett et al.
[2021] so that we replace h NN with itslinearization about w r , ∇ h NN ( x, w r )( w − w r ) .
With this linearization about w r , the above dynamicsϕ S becomes linear in w ,
[Equation] [Equation] where Y S = [ y 1 , · · · , y n ] T ∈ R n .
Recalling that W ⊂ R d w , Φ S is an n × d w matrix with the i throw being ∇ h NN ( x i , w r ) .
In the NTK regime, the dynamics ϕ S is well-approximated by ˜ ϕ S (see
Theorem 5.1 of [Bartlett et al., 2021] for conditions under which the approximation holds).
That is,the linear dynamics close to the NTK dynamics, referred to as the linear regression dynamics above,is the following,
[Equation]
37) where K S :=
Φ TS Φ S , A S :
= ( I − ηK S ) ∈ R d w × d w and b S :
= η (cid:0) K S w r + Φ TS Y S (cid:1)
∈ R d w .
Thedynamics converges to a fixed point ˜ w ∗ S =
A S ˜ w ∗ S + b S , as long as ∥ A S ∥ < 1 .
Evolution on function space Note that the invariant distribution of the above dynamics is singular:the delta distribution centered at ˜ w ∗ S .
In order to repeat the analysis in section 4 for this specialcase, we need to obtain the relationship between the rate of decay of correlations with respect to thisinvariant measure and the stability of the fixed point.
First isolating the rate of decay of correlations,this rate is equal to the second largest eigenvalue of the associated Frobenius-Perron on L 2 ( W ) orequivalently, the Koopman operator on L 2 ( W ) (since these two linear operators are adjoint to eachother, they share isolated spectra).
Since the dynamical system is linear, it is easy to verify thatthe eigenvalues of A S are also Koopman eigenvalues.
We can also check that the eigenfunctionscorresponding to eigenvalue θ i are of the form v Ti w + 1 / ( θ i − 1) v Ti b S , where v i are left eigenvectorsof A S .
Note that A S is a symmetric matrix whose largest absolute eigenvalue is equal to 1 − ηθ min , whereθ min > 0 is the smallest eigenvalue of the NTK Φ S Φ TS .
Stability of the fixed point In order to relate the rate of convergence, λ :=
1 − ηθ min with SAS in thiscase, we now describe SAS in this regime.
As we noted previously, SAS reduces to the algorithmicstability of the fixed point for the dynamics (37).
In order to deduce the stability of the fixed point tostochastic perturbations in the input, we need to recognize that the fixed point is an exact interpolant.
We can check that ˜ ϕ S ( ˜ w ∗ S ) =
˜ w ∗ S iff Y S =
˜ ϕ S ( ˜ w ∗ S − w r ) .
That is, the function ∇ h NN ( x, w r )( w ∗ S −w r ) exactly interpolates at the data points, among the class of linearized functions.
In other words, the fixed point is ˜ w ∗ S =
w r + a S , where a S is the minimum norm interpolation solution given bya S =
Φ TS (Φ S Φ TS ) − 1 Y S .
Thus, having a closed form expression for w ∗ S , we can obtain an upperbound on the stability of the algorithm ϕ S .
Since ˆ K S := Φ S Φ TS and its inverse are symmetric, ∥ ( ˆ K S ) − 1 ∥ is also the maximum eigenvalue of( ˆ K S ) − 1 .
A stochastic perturbation S ′ of S introduces a rank-one change denoted δK to ( ˆ K S ) − 1 .
From Weyl’s inequality, ∥ ( ˆ K S ) − 1 − ( ˆ K S ′ ) − 1 ∥ ≤ ∥ δK ∥ .
In the case of Lipschitz loss, an upperbound on β therefore depends on the maximum eigenvalue of ˆ K − 1 S , which is equal to 1 /θ min .
Thus,we see that a smaller θ min implies a smaller rate of convergence (slower convergence) as well as alarger upper bound on β (lesser algorithmic stability).
In this way, the linear regime also supportsthe analysis in section 4, which discusses a more general scenario of convergence of weights indistribution.
Remark 2.
As an aside that applies to the entire paper, we clarify that by “statistics” we refer tostatistics over the parameter space W .
The distribution over the weight space is specified in eachcontext.
Recall that the randomness in the sense of randomness over the weights arises due to thestochastic nature of SGD as well as the randomness over initial conditions.
Remark 3.
While we have considered linearization about a point, we may repeat the above analysisby considering another linear network by modifying the definition of the empirical matrix Φ S Φ T S to
K S =
E w ∈D w ∇ h NN
( X, W ) ∇ h T NN ( X, w ) .
This is the well-studied limit of the empirical kernel asthe number of neurons N tends to infinity.