On the generalization of learning algorithms that do not converge
Abstract
Generalization analyses of deep learning typically assume that the training converges to a fixed point.
But, recent results indicate that in practice, the weights of deep neural networks optimized with stochastic gradient descent often oscillate indefinitely.
To reduce this discrepancy between theory and practice, this paper focuses on the generalization of neural networks whose training dynamics do not necessarily converge to fixed points.
Our main contribution is to propose a notion of statistical algorithmic stability (SAS) that extends classical algorithmic stability to non-convergent algorithms and to study its connection to generalization.
This ergodic-theoretic approach leads to new insights when compared to the traditional optimization and learning theory perspectives.
We prove that the stability of the time-asymptotic behavior of a learning algorithm relates to its generalization and empirically demonstrate how loss dynamics can provide clues to generalization performance.
Our findings provide evidence that networks that “train stably generalize better” even when the training 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 some value, the learning algorithm—typically some variant of Stochastic Gradient Descent (SGD)—is terminated.
Perhaps surprisingly, recent works indicate that the network function at termination time is typically not a fixed point of the learning algorithm: if run longer, the gradient norm does not vanish and the learning algorithm outputs functions with significantly different parameters [Cohen et al., 2021, Zhang et al., 2022, Lobacheva et al., 2021].
This observation stands in contrast to common generalization analyses that assume convergence of the training algorithm to a fixed point.
It also raises the question of when (and why) non-convergent learning algorithms should be expected to generalize.
Standard approaches for obtaining generalization bounds find ways to bound the complexity measure of the function class expressible by a neural network in a data-independent e.g., [Vapnik, 1999, Bartlett et al., 2017, Neyshabur et al., 2018, Golowich et al., 2018, Bartlett et al., 2019, Arora et al., 2018] or data-dependent manner, e.g., [von Luxburg and Bousquet, 2004, Xu and Mannor, 2012, Sokoli´c et al., 2017].
Nevertheless, even with recent technical improvements, by and large these approaches do not explicitly account for the relationship between generalization and the dynamics of the learning algorithm.
A different approach to generalization analysis yields algorithm-dependent bounds by connecting the generalization 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, Bousquet et al., 2020, Zhang et al., 2021b, Rakhlin, 2006]. A key claim of these works has been that networks that are trained fast generalize better.
This claim, though intuitively meaningful, hinges on the premise of a convergent algorithm, which deviates from the observed behavior of deep neural network training.
Further work has been done for settings where the algorithm provably converges to a fixed point, such as two-layer overparameterized networks, deep linear networks, certain matrix factorizations, 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 a more general setting, Hardt et al.
[2016] prove algorithmic stability-based generalization bounds that worsen with increasing training time, whereas Loukas et al.
[2021] connect stable training dynamics near 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 points to small changes in the training set is not the mechanism underlying the good generalization properties of deep networks.
This is because, as is commonly known and many recent works indirectly remark upon [Cohen et al., 2021, Ahn et al., 2022, Garipov et al., 2018, Zhang et al., 2022], training can be linearly unstable : even if assuming the presence of isolated local minima that the learning algorithm converges to, the weights encountered in training are not robust to small perturbations.
Small perturbations 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) are typically comparable.
Our first contribution entails proposing a generalization of stability that applies to non-converging algorithms.
To achieve this, we depart from the typical optimization perspective of analyzing the generalization 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 algorithmic stability 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 bounds for the average 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 Section 4 we provide conditions such that the spectral gap of a Markov operator associated with the dynamics of the loss function is predictive of statistical algorithmic stability and hence of generalization.
Empirically, we estimate this spectral gap based on the rate of convergence of ergodic averages of the loss function.
We show via experiments that our estimate 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 [Bousquet and Elisseeff, 2002, Feldman and Vondrak, 2018, Kuzborskij and Lampert, 2018, Bousquet et al., 2020, Zhang et al., 2021b].
Below, we also briefly discuss connections (beyond the literature on stability) with previous analyses on the dynamics of learning algorithms.
Convergence and generalization of SGD.
Exponential convergence rates to global minima are known for SGD in the smooth and convex setting, with decaying learning rates and small constant learning rates (e.g., [Moulines and Bach, 2011, Needell et al., 2014]).
In the training of deep neural networks, which is the focus of this paper, the loss function is non-convex and is typically not well-approximated locally by convex functions [Ma et al., 2018].
However, in the overparameterized regime, a Polyak-Lojasiewicz-type inequality that is automatically satisfied when the loss is smooth and convex 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 local minima and even global minima has been proved, assuming a strict saddle property [Ge et al., 2015] or local convexity [Kleinberg et al., 2018].
Raginsky et al.
[2017] show a favorable convergence for a slightly different version of SGD by approximating the dynamics with a suitable continuous time Langevin diffusion.
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 that generalizes given suitable initialization and significant overparameterization.
A few works have studied the nonlinear dynamics of training [Kong and Tao, 2020] and provided insights into generalization in the overparameterized regime [Dogra and Redman, 2020, Saxe et al., 2013, Advani et al., 2020].
Another line of work considers 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 the optimization perspectives but relates to the mean field perspective in that we implicitly consider (see section 4) the evolution of probability distributions on weight space.
Flat minima.
In the generalization literature, flat minima correspond to large connected regions with low empirical error in weight space.
Flat minima have been argued to be related to the complexity of 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 batch size 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].
Some works consider the flat/sharp dichotomy an oversimplification [Dinh et al., 2017, Sagun et al., 2017,
He et al., 2019], e.g., using a reparameterization argument [Sagun et al., 2017].
The eigenvalues of the hessian of the loss affect the local linear stability of optimization orbits.
However, given that we adopt the time-asymptotic/statistical picture, we do not explicitly make assumptions on local stability (flatness/sharpness) or the topology of the loss landscape [Montanari and Zhong, 2020, Venturi et al., 2019].
SGD as a Bayesian sampler.
The idea of looking at distributions of learners permeates PAC-Bayesian analyses 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 by the learning algorithms.
Previous empirical work has suggested that SGD operates almost like a
Bayesian sampler [Mingard et al., 2021]
but the connection remains to be 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. by introducing the learning problem and discussing learning algorithms from a dynamical systems perspective.
Section 2.2 then puts forth the notions of statistical convergence that give rise to our main 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 ⊆ R , drawn i.i.d.
from a distribution D .
A class of parametric models or hypotheses h : R d × W → R on a parameter space W gives possible input-to-output relationships, h ( w, · ) : R d → R .
The learner is given a risk or loss function ℓ : ( R d × R ) × W
→ R + that describes the error of a given hypothesis.
Common choices for loss functions are the mean-squared, hinge and cross-entropy loss.
The learner attempts to minimize the population risk R S ( h ( · , w ))
=
E z ∼D ℓ ( z, w ) , by minimizing the empirical risk
[Equation]
This minimization is achieved using a local descent algorithm given by iterative updates ϕ S :
W → W of the form
[Equation]
where η t is the learning rate or step size at time t. [Figure] 40 loss [Figure]
[Figure] [Figure] [Figure] In the gradient descent (GD) algorithm, ˆ ∇ L S ( w ) = ∇ L S ( w ) is the gradient w.r.t.
w , and the iterates of ϕ S represent a deterministic dynamical system.
The critical points (including saddle points, local and global minima) of L S are fixed points of ϕ S (i.e., points w such that ϕ S w = w ).
For stochastic gradient 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 critical points 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 indicates the choice of batch.
Suppose the batch size m is fixed: m < n for SGD and m = n for GD.
Denote by [ n ] the set { 1 , 2 , · · · , n } and let Ξ t be a collection of m elements from [ n ] chosen uniformly at random.
Then, we may write ϕ S as
[Equation]
with Ξ t being the set [ n ] at all t for GD.
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 asymptotically converge to η .
A rapidly vanishing learning rate can obscure the learning dynamics by enforcing convergence to arbitrary points irrespective of the points’ losses.
In contrast, learning rates that do not decay to zero can induce interesting transitions between neighborhoods of critical points and the iterates w t are generally more exploratory of the loss landscape.
2.2 Statistical convergence of the learning algorithm
In deep learning, any specific choice of w T , at some time T, is arbitrary as the loss landscape is generally non-convex and stochastic learning algorithms are the norm.
Hence, rather than focusing on a single hypothesis, it is attractive to consider the average generalization error induced by the stochastic dynamical 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 individual orbits, we make an assumption about the ergodic properties of ϕ S that reasonably matches empirical evidence.
To describe the assumption, we give a short primer on the relevant concepts from ergodic theory of dynamical systems.
Invariant measures.
Denote the Borel sigma algebra on W by B ( W ) .
A probability measure µ S is called invariant for ϕ S if µ S ( A ) = µ S ◦ ϕ − 1 S ( A ) for all Borel sets A ∈ B ( W ) .
Intuitively, for any subset of the weight space, the probability that the dynamical system occupies it does not change.
See Liverani [2004] for an introduction to invariant measures.
If ϕ S is a deterministic continuous function on the set W and W is compact, classical ergodic theorems (see e.g., Theorem 4.1.1 of Katok and Hasselblatt [1997]) give the existence of at least one invariant, ergodic measure for ϕ S . Ergodic measures are a pertinent class of invariant measures in practice, that allow us to understand the statistical properties of a system through evolution of individual orbits. For an ergodic measure, sets that are invariant under the dynamics ϕ S (such as a set of periodic points, quasiperiodic or chaotic attractors) are trivial: they have measure 0 or 1.
When ϕ S is an SGD update, Dieuleveut et al.
[2020] show the 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 settings under mild assumptions on the gradient noise (in the estimate ˆ ∇ L S ) .
Several works have analyzed the asymptotic properties of such invariant distributions, e.g., Chee and Toulis [2018] analyze the oscillation of the iterates w t about a mean that is a critical point, as a phase separate from the transient phase of convergence to this invariant distribution.
Kong and Tao [2020] study the effect of large learning rates (see also [Wang et al., 2021] for matrix factorization problems) and multiscale loss functions to show the existence of Gibbs invariant measures in GD.
Ergodicity.
When an invariant measure µ S is ergodic for ϕ S , time averages converge to ensemble averages according to µ S .
That is, for µ S -almost every initial state w 0 , and for all continuous scalar functions 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 unique ergodic invariant measure on W .
Instead, ϕ S -invariant sets could be smaller subsets of W with different ergodic invariant measures supported on them.
As a result, starting from two different points chosen Lebesgue almost everywhere (or sampled from any probability density on W ), we do not expect time averages to converge to the same limit.
Indeed, as shown in Figure 1, the weights visited by ϕ S depend on the initial conditions and can vary significantly across different runs.
In other words, the limit of infinite time averages described in (4) does depend on w 0 .
Dynamics of the hypothesis.
To circumvent the above issues, rather than focusing on the weights, we consider the dynamics of the learned hypothesis h ( · , w t ) .
Specifically, suppose that the infinite-time averages of the learned hypothesis are ergodic: that is, the time averages of h ( z, w t ) converge to the same function irrespective of the initial function h ( z, w 0 ) .
Then, functionals that depend only on h and not explicitly on the weights are also ergodic.
This is one such scenario that leads to our main assumption: the time averages of loss functions are ergodic.
Assumption 1.
Given any S ∼ D n , there exists a map z → ⟨ ℓ z ⟩ S , such that for Lebesgue-a.e. wand 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 exists a unique probability on the space of neural network functions that is ergodic is less restrictive and closer to reality than assuming a unique ergodic measure on the parameter space. Furthermore, ergodicity of the network functions is a stronger condition than our Assumption 1 and serves to illustrate one possible sufficient condition for our assumption to hold.
Figure 1 provides evidence that Assumption 1 can hold in practice.
It shows two observables calculated along 4 different orbits of a VGG16 model trained using SGD with momentum (batches of size 128, learning rate of 0.01, momentum of 0.9).
On the left, we can see a low-dimensional projection 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 the right is the test accuracy, which is a functional of the hypothesis function.
Here, the test accuracy can be seen to converge to a distribution independent of the initial conditions, whilst the weight orbits vary significantly across different initializations.
This experiment corroborates our hypothesis that time averages of functionals on the loss space converge to distributions independent of the initial condition, even if time averages of a generic observable on the weight space do not.
See Remark 1 on the idea that a larger learning rate induces ergodicity and Appendix A, which illustrates the idea with a toy example.
Statistical stability implies generalization
In Section 3.1, we generalize algorithmic stability beyond algorithms that converge to fixed points.
Section 3.2 then proceeds to examine the implications of our definition to generalization error.
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 element in S by replacing it with a different element from the input distribution D :
Defnition 1 (Algorithmic stability, adapted from [Bousquet and Elisseeff, 2002]) .
Consider the weights w ∗ S and w ∗ S ′ obtained by running the learning algorithm on two training sets S , S ′ sampled from 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 not directly apply to learning algorithms that do not converge to a fixed point.
Hence, next, we extend algorithmic stability to loss statistics .
The resulting notion of statistical algorithmic stability (SAS) extends algorithmic stability to algorithms whose loss statistics converge as prescribed by Assumption 1, even if the weights do not converge.
Defnition 2 (Statistical algorithmic stability) .
We say that a learning algorithm with loss statistics denoted by ⟨ ℓ z ⟩ S is statistically algorithmically stable (SAS) with a stability coefficient β ≥
if
[Equation]
where S , S ′ differ in exactly one element.
In the above definition, ⟨ ℓ z ⟩ S ′ refers to the ergodic average of ℓ ( z, · ) observed along almost every orbit of ϕ S ′ .
A higher value of β indicates a lower SAS algorithm.
The definition of SAS differs from Definition 1 as we do not assume convergence of the learning algorithm.
Nevertheless, when every orbit of ϕ S converges to a fixed point w ∗ S , for almost every S, then Definition (7) reduces to the standard notion of algorithmic stability since ergodic averages 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 of the loss function and not the loss function at an ensemble mean or at any one point in the weight space W .
In other words, the definition of SAS does not use or provide information about the algorithmic stability of E w ∈ µ S w for any invariant measure µ S .
3.2 Learning theoretic implications
It is well-known that Definition 1 leads to generalization bounds (see [Hardt et al., 2016, Bousquet et al., 2020] and references therein).
Next, we show that the more general statistical algorithmic stability 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 loss statistics rather than loss values at fixed points:
[Equation]
Although the definitions above appear deceivingly similar to the standard ones introduced in Section 2.1, the two notions of risk are defined on different spaces, with the standard ones being maps from the hypothesis space H to scalar values and (8) being maps from the space of learning algorithms . That is, the latter risks only depend on algorithmic parameters and are not functions of H .
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 and is based on applying Mcdiarmid’s inequality to the generalization gap.
Like Bousquet and Elisseeff [2002]’s bound, the SAS coefficient β has to decay at least as fast as ∼ O (1 / √ n ) for the generalization bound to be non-vacuous.
Note that our proof is also naturally coupled with [Bousquet et al., 2020, Feldman and Vondrak, 2018]
’s improved technique, which provides tighter bounds when β ∼ O (1 / √ n ) .
(see Appendix B).
Dynamical systems interpretation of SAS
How can we distinguish between two algorithms with different SAS?
Can the algorithm dynamics provide clues to its SAS?
These questions are of practical importance because a more (statistically) stable algorithm generalizes better (via Theorem 1).
Since the SAS coefficient β is defined as a supremum over infinite pairs of training sets and inputs, a numerical estimation of it using its definition only gives a rough lower bound and obtaining this lower bound is also computationally expensive (see Figure 3 and section 5). Moreover, a numerical lower bound on β does not give us a mechanistic understanding of statistical stability, which we seek here. Here we develop an operator-theoretic explanation for what makes an algorithm statistically stable.
We also derive a rough heuristic – albeit not a definitive predictor of generalization – that relates a given training dynamics to its SAS coefficient β . 4.1 Bounding the SAS coefficient Before we present our bound, we discuss why we must deviate from the trajectory-based approach that is commonly employed to study algorithmic stability. We recall that Hardt et al.
[2016] base their proofs of classical algorithmic stability (Theorems 3.7 and 3.8 in Hardt et al. [2016]) on the accumulating differences between two orbits corresponding to S and S ′ , whenever the different element is chosen in the mini-batch. The greater the difference in the weights, the greater the difference in the (upper bound on the) loss functions at those weights, in the case of Lipschitz loss. Thus, Hardt et al.
[2016] argue that training faster (or stopping earlier) will lead to stabler algorithms.
Since SAS is about robustness of loss statistics or infinitely long time averages, the analysis à la Hardt et al.
[2016] generically leads to vacuous bounds for the SAS coefficient.
Our analysis is based on the realization that an algorithm can be SAS even if two orbits corresponding to S and S ′ diverge from each other (within W ) , but lead to similar loss statistics.
Thus, to bound the
SAS coefficient, a perturbation analysis of transition operators (global information), rather than a finite-time perturbation analysis of an orbit (local information), is needed.
Since SAS only demands robustness on loss space, we consider Markov transition operators on the loss space ( ⊆
[0 , L ] ), as opposed to on the weight space ( W ). In general, at any z and w 0 , the loss process, { ℓ ( z, w t ) } t is not Markovian. But, a family of Markov operators can be associated with the family of loss functions.
Lemma 1.
(Markov operators) Let µ S be an ergodic, invariant measure for ϕ S .
Assume that a loss function ℓ ( z, · ) : W → I z ⊆ [0 , L ] is such that the pushforward,
ν zS : B
( I z ) → R + , of µ S on the loss space is well-defined.
Here, B ( I z ) is the Borel sigma algebra on I z and ν zS = ℓ ( z, · ) ♯ µ S , where ♯ refers to the pushforward operation.
Then, there exists a uniformly ergodic Markov operator P zµ S on the space of probability measures on I z , with its invariant measure being ν zS .
This lemma (see Appendix C for a proof) parallels Chekroun et al.
[2014]’s Theorem A, but in a quite different setting, without using existence of a unique physical measure on the phase space ( W ).
Combined with Assumption 1, perturbation results on the Markov operators, P zµ S , defined by Lemma 1 lead us to SAS. Since P zµ S is uniformly ergodic, there exist λ z ∈ (0 , 1) and C > 0 such that for all ξ ∈ I z , in the Wasserstein norm ( ∥ · ∥ W ), [Equation]
Let λ := sup z λ z be the rate of mixing associated with the family of operators P zµ S . We leverage the perturbation theory of Markov operators to study the effect of stochastic perturbations on ν zS , and subsequently, SAS.
To see the connection with SAS, we recall that, by Assumption 1, for any choice of µ S , we have that ⟨ ℓ z ⟩ S =
E ξ ∼ ν zS ξ.
Thus, given any pair of µ S and µ S ′ , we have (see Appendix C for details) [Equation] Thus, a uniform upper bound on ∥ ν
zS − ν zS ′ ∥ W gives an upper bound for the SAS coefficient, β. Such a bound is obtained from the straightforward use of an appropriate perturbation result on uniformly ergodic Markov chains.
Theorem 2.
Given a uniformly ergodic Markov operator P zµ S constructed in Lemma 1, and that the Lipschitz constant of ∇ ℓ ( · , w ) on R d is bounded above by L D for all w ∈ W , ϕ S is SAS with stability coefficient
[Equation]
where λ = sup
z λ
z is the supremum over z of the mixing rates of P zµ S .
Proof.
(Sketch)
The proof relies on the perturbation bounds of Rudolf and Schweizer [2018], Corollary 3.2.
Note that when (10) is satisfied, and if a stochastic perturbation S ′ introduces a change δ P zµ S to P zµ S , the perturbation bound from [Rudolf and Schweizer, 2018] gives
[Equation]
[Equation]
where ∥ δ P zµ S ∥ is the operator norm of δ P zµ S induced by Wasserstein norm.
From (11),
[Equation]
Since this holds for all z, the right hand side of (13) is an upper bound for the stability coefficient β in Definition (2), by taking the supremum over z .
An upper bound on (cid:13)(cid:13)
δ P zµ S (cid:13)(cid:13) can be derived when the Lipschitz constant of ∇ ℓ ( · , w ) is uniformly (in w ) bounded on R d , as shown in Appendix C).
Theorem 2 implies that an algorithm that exhibits faster convergence to the stationary measure on the loss space generalizes better .
The bound in Theorem 2 implies that smaller λ (faster convergence of the loss to the ergodic invariant measure) yields smaller upper bounds on β (more statistically stable algorithm) and better generalization.
Strictly speaking, to determine λ z , we must consider finite-dimensional approximations of the operator 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 difficult in 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 convergence to the equilibrium distribution.
In particular, we hypothesize below that the auto-correlation function of the test loss can separate algorithms based on their SAS.
Dropping the superscripts z , we define the auto-correlation C ℓ ( τ ) in a loss function ℓ (see also Remark 3 in Appendix C) by
[Equation]
That is, C ℓ ( τ ) gives the correlation between the random variables ℓ ( w t ) / ⟨ ℓ ⟩ and ℓ ( w t + τ ) / ⟨ ℓ ⟩ , where the randomness comes from the batch selection Ξ t and the initialization of weights in SGD, and only from the latter in GD.
Suppose that the loss process is Markov, i e., ℓ ( w t +1 ) conditioned on ℓ ( w t ) is independent of { ℓ ( w t ′ ) :
t ′ ≤ t − 1 } . Setting t = 0 after a sufficiently long runup time, the function C ℓ ( τ ) in (14) is the auto-correlation function of a stationary Markov chain, which typically decays exponentially with τ.
This decay rate is lower (slower correlation decay/longer correlation times) when λ is closer to 1. Qualitatively, comparing two different stationary Markov chains, the one with higher values of C ℓ has longer correlation times and hence, larger λ.
In general, however, the loss process { ℓ ( w t ) : t ≥ 0 } is non-Markovian. Thus, we do not expect
C ℓ ( τ ) to decay with τ.
This is because the correlation function C ℓ encodes both the Markovian and non-Markovian components component of the loss dynamics (see Remark 4 in Appendix C; [Zwanzig, 2001, Kondrashov et al., 2015]). Hence, we expect the correlation times in the loss process to not be equal to that of a Markov process generated according to P µ S .
For the purpose of qualitatively comparing the SAS coefficients of two algorithms, however, the loss process can be considered a heuristic for a Markov process generated according to P µ S . Then, higher values of C ℓ computed from the loss process 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 trained on the 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 versions of the training set S p by corrupting
CIFAR10’s labels with probability p, with S 0 being the original CIFAR10 dataset.
Figures 2 andshow results corresponding to p = 0 , 0 . 1 , 0 .
25 and 0 .
Each line in Figure 3 is a sample mean over 10 random initializations.
In Figure 2, we numerically estimate a lower bound on the SAS coefficient β using its definition (see Definition 2).
On the left, we plot the difference in test loss time-averages between between random orbits of ϕ S p and ϕ S ′ p , with S ′ p being a stochastic perturbation of S p .
The mean over 45 pairs of orbits is shown as dots and the error bars indicate the standard error in mean.
The cumulative time average along orbits of length 1200 epochs are used as estimators for statistics.
With these estimators for statistics, in Figure 2(center), we compute an estimate of the stability coefficient as the difference of the (estimated) test loss statistics at different values of p.
This is hence an estimate on the lower bound of β, and clearly increases with p. This illustrates that cases with worse generalization errors (Figure 2 (right)) have larger lower bounds on β. In Figure 3 (left), we show the test loss timeseries obtained with the ResNet18 model. Again, greater the noise corruption p, the larger is the generalization error (estimated by test error), consistent with previous studies Zhang et al.
[2021a], Loukas et al.
On the other hand, the generalization gap – difference between training and test error – is bigger for intermediate levels of noise and decreases when p = 0 . 5 , as shown in Figure 3 (center). The gap does not always increase with p because, when the labels are close to random, the network cannot fit the training data and thus the training error is also large. 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 the auto-correlations over 10 independent runs.
We see that the magnitude of the test loss autocorrelations preserves the same order (across p ′ s) as the generalization gap (absolute difference in test and training losses) shown in Figure 3(center).
Hence, we empirically observe that the loss process codifies the phenomenological explanation for SAS that is expressed in Theorem 2.
Discussion and conclusion
Predicting generalization is an active area of research and has implications for more reliable use of machine learning.
In this work, we introduce statistical algorithmic stability (SAS) and show how it implies generalization bounds.
Here, we add a few important remarks.
Statistical stability only requires robustness of statistics on loss space.
A central challenge in the analyses of non-convergent algorithms is the fact that there may be multiple, very different invariant measures µ S .
For this reason, our stability criterion focuses on statistics of the loss function.
This way, an algorithm can be stable even if the measures µ S and µ S ′ are not close in the total variation norm.
In fact, an algorithm can be stable even if µ S and µ S ′ are mutually singular, if the loss functions have similar statistics.
Since our SAS analysis hinges on a reasonable yet nonrestrictive assumption on 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 the proposed relationship (section 4) between statistical stability and convergence rates to stationary measures can also be observed in the Neural Tangent Kernel (NTK) regime
[Jacot et al., 2018].
In this case, the training dynamics, ϕ S , can be approximated by a linear function of the weights, which converges to a fixed point.
Thus, this provides an alternative, ergodic theoretic, interpretation of generalization 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 gathers more evidence to support the broader view [Wojtowytsch, 2021, Zhang et al., 2022] that exploiting theoretically and empirically available dynamical information about the training algorithm is a fruitful complement to understanding generalization from the optimization landscape and learning theory perspectives. Funding: This work was partially funded by the NSF AI Institute TILOS, NSF award 2134108, and ONR grant N00014-20-1-2023 (MURI ML-SCOPE).
Acknolwedgments: N.C. would like to thank Nandhini Chandramoorthy and Derek Lim for helping with the experimental setup and Sven Wang, Benjamin Zhang, Matt Li and Youssef Marzouk for valuable discussions.
The authors also thank the reviewers for their constructive suggestions.
Checklist
For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions 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 to them?
[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 the supplementary material (b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were 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., type of 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?
[N/A] (b) Did you mention the license of the assets?
[N/A] (c) Did you include any new assets either in the supplemental material or as a URL?
[N/A] (d) Did you discuss whether
and how consent was obtained from people whose data you’re using/curating?
[N/A] (e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content?
[Yes] 5.
If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?
[N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
[N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
[N/A]
A Bifurcation analysis of smooth and non-convex optimization
In this section, we discuss some examples of non-convex optimization in one dimension performed with gradient descent (GD).
We illustrate that, with increasing learning rate, the asymptotic behavior of orbits may alter from being periodic to quasiperiodic to chaotic.
These qualitative changes are brought 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 and non-convex objective functions of the form ℓ s = g s ◦ g s ◦ g s , with g s being the canonical quadratic map g s ( w ) :
= 1 − sw (1 − w ) of the unit interval [0 , 1] .
Since g s is smooth, the compositions of g s with itself are smooth.
The first bifurcation point appears just above η = 2 / | d 2 ℓ s /dw 2 | , which is the stability threshold for convex optimization [Nesterov, 2003].
In Figure 4, the left, center and right columns correspond to s = 1 , 3 and 4 respectively.
In the first row, we plot the loss functions ℓ s , which are non-convex (with multiple global minima) at s = 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 global minima 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 a function of the learning rate.
The attractor is approximated by the asymptotic orbits of the dynamics at multiple (100) different initial conditions chosen randomly on the unit interval.
Note that at small values of η < 2 / ∥ a ∥ with ∥ a ∥ := sup w ∈ [0 , 1] a ( w ) orbits from different initial conditions converge to fixed 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 noted from 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 descent dynamics,
[Equation]
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 linear perturbations along the orbit of w.
Here, we use the exponential notation, ϕ ts w = ϕ
s ◦ ϕ t − 1 s w to denote compositions of ϕ s .
A positive Lyapunov exponent indicates dynamical instability, e.g., chaotic orbits.
We see from the bottom row of Figure 4 that in the range of learning rates considered, chaos is observed (positive Lyapunov exponents starting from uniformly random initial conditions) in the case 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 learning rates where convergence to fixed points or periodic behavior is observed, the Lyapunov exponent converges to different negative values depending on the initial condition.
This example supports the heuristic explanation for the primary assumption made in the main text about the asymptotic dynamics of learning algorithms.
The assumption of lack of ergodicity ensures that the analysis is applicable to a range of learning rates, even those smaller than 2 / ∥ a ∥ .
Remark 1.
Increasing the learning rate may induce ergodicity.
Assumption 1 may hold for all continuous functions 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 local minimum to another. On the other hand, for small learning rates, we may be able to detect the presence of multiple attractors, since different initial conditions lead to different long-time averages. The above toy example illustrates these ideas. B Proof of Theorem
In this section, we prove Theorem 1 from the main text.
This result says that the statistical stability of an 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., 2020]).
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 constant across 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 that has at most one element different from S, i.e., a stochastic perturbation of S. Using the stochastic perturbation S ′ , we can rewrite E S ∼D n ⟨ ℓ z 1 ⟩ S as E S ∼D n E z ∈D ⟨ ℓ z ⟩ S ′ .
Substituting this equivalent expression in (16), 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 , which gives a high-probability upper bound on Φ S − E (Φ S ) in terms of the deviation of Φ S from a Φ S ′ .
In order to obtain an upper bound for the latter quantity, note that
[Equation]
[Equation]
Considering the difference of the population risks,
[Equation]
Next considering the difference of empirical risks,
[Equation]
(21) where, recall that L := sup z ∈ R d × R sup w ∈ W | ℓ ( z, · ) | .
Putting (19) together with (20) into Inequality (18), 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 .
A function f : Z n → R is said to satisfy the bounded differences property if there exists a constant c > 0 such that
[Equation]
for all single coordinate changes.
For an f that satisfies the bounded differences with a constant c > 0 , given a δ > 0 , with probability at least 1 − δ,
[Equation]
Applying Mcdiarmid’s inequality to Φ S , and recalling (17)
, we obtain the generalization bound stated in 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.
C Predicting generalization with loss timeseries
In this section, we substantiate the connection drawn in section 4 of the main text between the rate of decay of correlations in the loss function and SAS.
The study of dynamics lifted to the space of 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 computational methods 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 extensively appeared 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 various contexts, we are able to show that the autocorrelations in the loss function can serve as predictors for the generalization gap.
This section provides the complete setting for the results in section 4 of the main text and completes the proof of Theorem 2.
First we define the transition operator for probability distributions on the weight space W and analogous operator on the loss space.
Markov operator for weight space.
Recall that Ξ t is a batch of m indices chosen uniformly from the set [ n ] := { 1 , · · · , n } .
For GD, Ξ t is deterministic and equal to the set [ n ] .
Let (Ω , Σ , P ) be a probability 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 , [Equation]
[Equation]
where P ( ϕ S ( w ) ∈
A | Ξ t ) is the probability of the event that ϕ S ( w )
∈
A when the m indices Ξ t are chosen.
Correspondingly, we may define the Markov operator P S , also called the Frobenius-Perron operator 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 potentially multiple eigendistributions corresponding to eigenvalue 1.
Each invariant measure µ S also defines different 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 (23), we may now define a kernel K zµ S , for an ( ξ, E ) ∈
I z × B ( I z ) , as
[Equation]
[Equation] This, in turn, gives rise to a Markov operator analogous to the Frobenius-Perron operator (24) on the full-dimensional weight space:
[Equation]
Note that this operator is well-defined when the level sets of ℓ z are µ S -measurable.
One sufficient condition for this is when the foliation of W by these level sets is subordinate to a measurable partition.
We may then consider disintegrations of µ S on this measurable partition and define the kernel (26) 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 (27) is mixing.
For this, we make an additional assumption.
We assume that the Frobenius-Perron operator P S mixes to the measure µ S starting from any measure of the form µ = ℓ ( z, · ) − 1 ν, for an absolutely continuous probability measure ν on [0 , L ] .
That is,
[Equation]
where ∥ · − ∗ ∥ TV indicates the total variation distance and ζ ∈ (0 , 1) is the rate of mixing.
On the other 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 (29) conveys that when a measure µ t converges to µ S in the TV norm, expectations with respect to µ t of all continuous functions, of which { ℓ ( z, · ) } z is a subset, also converge to expectations with respect to µ S .
Thus, intuitively we expect that when sup z λ z is a lower bound for ζ ((28)).
In the main text, below Lemma 1, we state the uniform ergodicity of P zµ S in terms of the Wasserstein norm.
This holds from (29) since convergence in TV distance implies convergence in Wasserstein.
This concludes the proof of Lemma 1.
Finally, note that 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 the full-dimensional (weight) space.
Effect of stochastic perturbations In section 4 of the main text, we use the perturbation theory of mixing Markov operators to relate the SAS coefficient to the mixing rate of P zµ S .
This perturbation bound (from [Rudolf and Schweizer, 2018]) is given in terms of the perturbation δ P zµ S to the operator P zµ S when a stochastic perturbation S → S ′ is applied to the training set.
Here we discuss the size of δ P zµ S , completing the proof of Theorem 2 in the main text.
First note that ∥ δP zµ S ∥ ≤ ∥P S − P S ′ ∥ , and hence it suffices to obtain an upper bound for ∥P S − P S ′ ∥ .
The perturbation to Markov kernel in the weight space due to a stochastic perturbation is given by
[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.
This leads to the following upper bound on the perturbation size in the Wasserstein norm, ∥P
[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) as explained in the proof sketch of Theorem 2 in section 4.
Thus, we obtain again, an upper bound for β
=
O ( L D / (1 − λ ) n ) as claimed in Theorem 2. Remark 2.
In this section as well as section 4, we use inequalities between different norms on the space of finite signed measures. This is a Banach space with the total variation norm, isomorphic to L 1 (Λ) for some background measure Λ . The inequalities (for instance, (11) ) follow from the dual characterization of the norms.
Remark 3.
Although the definition of the autocorrelation in section 4 appears to be a function of w 0 ,
C ℓ ( τ ) does not depend on w 0 due to Assumption 1.
Note that the rate of convergence to equilibrium, the coefficient λ in the definition of uniform ergodicity, also determines the correlation decay rate when ℓ is initialized out of equilibrium (That is, ℓ ( w 0 ) does not sample ν S ).
Remark 4.
One way to understand the non-Markovian loss process is through the Mori-Zwanzig formalism [Zwanzig, 2001] that originated in statistical mechanics and has found extensive applications in deriving reduced-order models for complex physics (e.g. see [Lin and Lu, 2021, Kondrashov et al., 2015]).
In this formalism, we can consider the exact evolution of a finite set of observables, Ψ( w ) =
[ ψ 1 ( w ) , · · · , ψ p ( w )] such that ℓ ∈ span { Ψ } (here, ℓ := ℓ z for an arbitrary z ).
At time t, Ψ t := Ψ( w t ) can be written as a sum of three terms: 1) a Markov term that depends on the values
Ψ t − 1 , 2) a non-Markovian memory term that is a function of { Ψ t − k : t ≥ k > 1 } , and 3) a noise term that is a function of w 0 .
The first two terms depend on the dynamics ϕ S and are hence different for different values of S.
The autocorrelation function C ℓ has contributions from both the Markovian and non-Markovian terms.
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 same parameters as in section 5 of the main text and with training data that are stochastic perturbations of each other.
We consider 45 pairs of stochastically perturbed datasets (see section 2 of main text for definition of stochastic perturbation) each for each value of p .
The value of p = 0 , 0 . 1 , 0 . 17 , 0 . 25 ,indicates the probability of error injected into the labels of the CIFAR10 dataset.
The mean of the absolute difference in the cumulative test loss is shown in dark colors while the standard error in mean in the corresponding lighter color.
The time averages are calculated over 1200 epochs after a run up time of 200 epochs.
The results indicate that greater the noise probability p, greater the estimate of 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 label corruption (noise) levels.
Furthermore, considering together with the VGG16 results presented in the main paper, this relationship between SAS and generalization holds across the different architectures we 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 approximately equal to (strictly, an upper bound for) | R S − ˆ R S | .
A loose upper bound for this quantity comes from the theoretical generalization bound (in Theorem 1).
In practice, this quantity is estimated to be small relative to the test error with corrupt datasets.
When the noise probability is 0% (original dataset), the gap estimate is about 80% of the test error because the training error is small.
But, when the noise probability is 50%, the gap estimate is about 20% of the test error since the training error is also large.
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.
Finally, we remark that our empirically estimated autocorrelation function serves as a proxy until a more sophisticated method for the estimation of λ is developed.
As mentioned in the main text (section 4), this is a challenging 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 has been analyzed thoroughly in numerous works, wherein its convergence to kernel ridge regression has been formally proved in two-layer, infinitely wide networks [Montanari and Zhong, 2020, Bartlett et 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 generalization properties 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 present paper 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 is well-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 need only consider linear regression dynamics.
That is, let ˜ w t be an orbit of a perturbed dynamics (the linear 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 the algorithmic stability of the NTK orbit, it is enough to consider the stability of the linear regression orbit, 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 decay of correlations.
In this case, we show that the speed of convergence to the fixed point determines the generalization properties via algorithmic stability.
We derive stability-based generalization bounds, an alternative to Rademacher complexity-based bounds in Arora et al.
[2019a].
Comparing with existing generalization results in the well-understood NT regime is an ideal test bed for the alternative dynamical 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 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, with the initialization w 0 = w r .
With GD on the squared loss L S ( w ) =
(1 / 2) (cid:80) ni =
( 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 its linearization 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]
36) 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 th row 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 .
The dynamics 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 special case, we need to obtain the relationship between the rate of decay of correlations with respect to this invariant 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 ) or equivalently, the Koopman operator on L 2 ( W ) (since these two linear operators are adjoint to each other, they share isolated spectra).
Since the dynamical system is linear, it is easy to verify that the eigenvalues of A S are also Koopman eigenvalues.
We can also check that the eigenfunctions corresponding to eigenvalue θ i are of the form v Ti w + 1 / ( θ i − 1) v Ti b S , where v i are left eigenvectors of 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 this case, we now describe SAS in this regime.
As we noted previously, SAS reduces to the algorithmic stability of the fixed point for the dynamics (37).
In order to deduce the stability of the fixed point to stochastic 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 by a S =
Φ TS (Φ S Φ TS ) − 1 Y S .
Thus, having a closed form expression for w ∗ S , we can obtain an upper bound 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 upper bound 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 a larger upper bound on β (lesser algorithmic stability).
Hence, the linear regime also supports the analysis in section 4, which discusses a more general scenario of convergence of weights in distribution.
Remark 5.
As an aside that applies to the entire paper, we clarify that by “statistics” we refer to statistics over the parameter space W .
The distribution over the weight space is specified in each context.
Recall that the randomness in the sense of randomness over the weights arises due to the stochastic nature of SGD as well as the randomness over initial conditions.
Remark 6.
While we have considered linearization about a point, we may repeat the above analysis by considering another linear network via modifying the definition of the empirical matrix Φ S Φ TS to
[Equation] X
[Equation] NN
[Equation] X, w ) .
This is the well-studied limit of the empirical kernel as the number of neurons N tends to infinity.