Stability and Generalization of Learning Algorithms that Converge to Global Optima

Introduction

The recent success of training complex models at state-of-the-art accuracy in many common machine learning tasks has sparked significant interest and research in algorithmic machine learning. In practice, not only can these complex deep neural models yield zero training loss, they can also generalize surprisingly well . Although there has been significant recent work in analyzing the training loss performance of several learning algorithms, our theoretical understanding of their generalization properties falls far below what has been observed empirically.

A useful proxy for analyzing the generalization performance of learning algorithms is that of stability. A training algorithm is stable if small changes in the training set result in small differences in the output predictions of the trained model. In their foundational work, Bousquet and Elisseeff establish that stability begets generalization.

While there has been stability analysis for empirical risk minimizers , there are far fewer results for commonly used iterative learning algorithms. In a recent novel work, Hardt et al. establish stability bounds for SGD, and discuss algorithmic heuristics that provably increase the stability of SGD models. Unfortunately, generalizing their techniques to establish stability bounds for other first-order methods can be a strenuous task. Showing non-trivial stability for more involved algorithms like SVRG (even in the convex case), or SGD in more nuanced non-convex setups is far from straightforward. While provides a clean and elegant analysis that shows stability of SGD for non-convex loss functions, the result requires very small step-sizes. The step-size is small enough that one may require exponentially many steps for provable convergence to an approximate critical point, under standard smoothness assumptions (see subsection A.10). Generally, there seems to be a trade-off between convergence and stability of algorithms. In this work we show that under certain geometric assumptions on the loss function around global minima, we can actually leverage the convergence properties of an algorithm to prove that it is stable.

The goal of this work is to provide black-box and easy-to-use stability results for a variety of learning algorithms in non-convex settings. We show that this is in some cases possible by decoupling the stability of critical points and their proximity to models trained by iterative algorithms.

We establish that models trained by algorithms that converge to local minima are stable under the Polyak-Łojasiewicz (PL) and the quadratic growth (QG) conditions . Informally, these conditions assert that the suboptimality of a model is upper bounded by the norm of its gradient and lower bounded by its distance to the closest global minimizer. As we see in the following, these conditions are sufficient for stability and are general enough to yield useful bounds for a variety of settings.

Our results require weaker conditions compared to the state-of-the art, while recovering several prior stability bounds. For example in the authors require convexity, or strong convexity. Gonen and Shalev-Shwartz prove the stability of ERMs for nonconvex, but locally strongly convex loss functions obeying strict saddle inequalities . By contrast, we develop comparable stability results for a large class of functions, where no convexity, local convexity, or saddle point conditions are imposed. We note that although establishes the stability of SGD for smooth non-convex objectives, the stepsize selection can be prohibitively small for convergence. In our bounds, we make no assumptions on the hyper-parameters of the algorithms.

We use our black-box results to directly compare the generalization performance of popular first-order methods in general learning setups. While direct proofs of stability seem to require a substantial amount of algorithm-specific analysis, our results are derived from known convergence rates of popular algorithms. For strong convexity—a special case of the PL condition—we recover order-wise the stability bounds of Hardt et al. , but for a large family of optimization algorithms (e.g., SGD, GD, SVRG, etc). We show that many of these algorithms offer order-wise similar stability as saddle-point avoiding algorithms in nonconvex problems where all local minima are global . We finally show that while SGD and GD have analogous stability in the convex setting, this breaks down in the non-convex setting. We give an explicit example of a simple 1-layer neural network on which SGD is stable but GD is not. Such an example was theorized in (i.e., Figure 10 in the aforementioned paper); here we formalize the authors’ intuition. Our results offer yet another indication that SGD trained models can be more generalizable than full-batch GD ones.

Finally, we give examples of some machine learning scenarios where the PL condition mentioned above holds true. Adapting techniques from , we show that deep networks with linear activation functions are PL almost everywhere in the parameter space. Our theory allows us to derive results similar to those in about local/global minimizers in linear neural networks.

Prior work

The idea of stability analysis has been around for more than 30 years since the work of Devroye and Wagner . Bousquet and Elisseef defined several notions of algorithmic stability and used them to derive bounds on generalization error. Further work has focused on stability of randomized algorithms and the interplay between uniform convergence and generalization . Mukherjee et al. show that stability implies consistency of empirical risk minimization. Shalev-Shwartz et al. show that stability can also imply learnability is some problems.

Our work is heavily influenced by that of Hardt et al. that establish stability bounds for stochastic gradient descent (SGD) in the convex, strongly convex, and non-convex case. The work by Lin et al. shows that stability of SGD can be controlled by forms of regularization. In , the authors give stability bounds for SGD that are data dependent. Since they do not rely on worst-case arguments, they lead to smaller generalization error bounds than that in , but require assumptions on the underlying data. The work by Liu et al. gives a related notion of uniform hypothesis stability and show that it implies guarantees on the generalization error.

Stability is closely related to the notion of differential privacy introduced in . Roughly speaking, differential privacy ensures that the probability of observing any outcome from a statistical query changes if you modify any single dataset element. Dwork et al. later showed that differentially private algorithms generalize well . These generalization bounds were later improved by Nissim and Stemmer . Such generalization bounds are similar to those guaranteed by stability but often require different tools to handle directly.

Preliminaries

The generalization performance of the model can then be measured by the generalization gap:

For our purposes, $w$ will be the output of some (potentially randomized) learning algorithm $\mathcal{A}$, trained on some data set $S$. We will denote this output by $\mathcal{A}(S)$.

Let us now define a related training set $S^{\prime}=\{z_{1},\ldots,z_{i-1},z_{i}^{\prime},z_{i+1},\ldots,z_{n}\}$, where $z_{i}^{\prime}\sim\mathcal{D}$. We then have the following notion of uniform stability that was first introduced in .

An algorithm $\mathcal{A}$ is uniformly $\epsilon$-stable, if for all data sets $S,S^{\prime}$ differing in at most one example, we have

The expectation is taken with respect to the (potential) randomness of the algorithm $\mathcal{A}$. Bousquet and Elisseeff establish that uniform stability implies small generalization gap .

Suppose $\mathcal{A}$ is uniformly $\epsilon$-stable. Then,

In practice, uniform stability may be too restrictive, since the bound above must hold for all $z$, irrespective of its marginal distribution. The following notion of stability, while weaker, is still enough to control the generalization gap. Given a data set $S=\{z_{1},\ldots,z_{n}\}$ and $i\in\{1,\ldots,n\}$, we define $S^{i}$ as $S\backslash z_{i}$.

Note that this is a weaker notion than uniform stability, but one can still use it to establish non-trivial generalization bounds:

In the following, we derive stability bounds like the above, for models trained on empirical risk functions satisfying the PL and QG conditions. To do so, we will assume that the functions in question are $L$-Lipschitz.

If $f$ is assumed to be differentiable, this is equivalent to saying that for all $x$, $\|\nabla f(x)\|_{2}\leq L$.

In a recent work, Karimi et al. in used the Polyak-Łojasiewicz condition to prove simplified nearly optimal convergence rates for several first-order methods. Notably, there are some non-convex functions that satisfy the PL condition. The condition is defined below.

Fix a set $\mathcal{X}$ and let $f^{*}$ denote the minimum value of $f$ on $\mathcal{X}$. We will say that a function $f$ satisfies the Polyak-Łojasiewicz (PL) condition on $\mathcal{X}$, if there exists $\mu>0$ such that for all $x\in\mathcal{X}$ we have

Note that for PL functions, every critical point is a global minimizer. While strong convexity implies PL, the reverse is not true. Moreover, PL functions are in general nonconvex (e.g., invex functions). We also consider a strictly larger family of functions that satisfy the quadratic growth condition.

We will say that a function $f$ satisfies the quadratic growth (QG) condition on a set $\mathcal{X}$, if there exists $\mu>0$ such that for all $x\in\mathcal{X}$ we have

where $x_{p}$ denotes the euclidean projection of $x$ onto the set of global minimizers of $f$ in $\mathcal{X}$ (i.e., $x_{p}$ is the closest point to $x$ in $\mathcal{X}$ satisfying $f(x_{p})=f^{*}$).

Both of these conditions have been considered in previous studies. The PL condition was first introduced by Polyak in , who showed that under this assumption, gradient descent converges linearly. The QG condition has been considered under various guises and can imply important properties about the geometry of critical points. For example, showed that local minima of nonlinear programs satisfying the QG condition are actually isolated stationary points. These kinds of geometric implications will allow us to derive stability results for large classes of algorithms.

Stability of Approximate Global Minima

In this section, we establish the stability of large classes of learning algorithms under the PL and QG conditions presented above. Our stability results are “black-box” in the sense that our bounds are decomposed as a sum of two terms: a term concerning the convergence of the algorithm to a global minimizer, and a term relevant to the geometry of the loss function around the global minima. Both terms are used to establish good generalization and provide some insights into the way that learning algorithms perform.

For a given data set $S$, suppose we use an algorithm $\mathcal{A}$ to train some model $w$. We let $w_{S}$ denote the output of our algorithm on $S$. The empirical training error on a data set $S$ is denoted $f_{S}(w)$ and is given by

We assume that each of these losses is $L$-Lipschitz with respect to the parameters of the model. We are interested in conditions on $f_{S}$ that allow us to make guarantees on the stability of $\mathcal{A}$. As it turns out, the PL and QG condition will allow us to prove such results. Although it seems unclear if these conditions are reasonable, in our last section we show that they arise in a large number of machine learning settings, including in certain deep neural networks.

To analyze the performance machine learning algorithms, it often suffices to understand the algorithm’s behavior with respect to critical points. This requires knowledge of the convergence of the algorithm, and an understanding of the geometric properties of the loss function around critical points. As it turns out, the PL and QG conditions allow us to understand the geometry underlying the minima of our function. Let $\mathcal{X}_{\min}$ denote the set of global minima of $f_{S}$.

Assume that for all $S$ and $w\in\mathcal{X}$, $f_{S}$ is PL with parameter $\mu$. We assume that applying $\mathcal{A}$ to $f_{S}$ produces output $w_{S}$ that is converging to some global minimizer $w_{S}^{*}$. Then $\mathcal{A}$ has pointwise hypothesis stability with parameter $\epsilon_{stab}$ satisfying the following conditions.

Case 1: If for all $S$, $\|w_{S}-w_{S}^{*}\|\leq O(\epsilon_{\mathcal{A}})$ then

Case 2: If for all $S$, $|f_{S}(w_{S})-f_{S}(w_{S}^{*})|\leq O(\epsilon^{\prime}_{\mathcal{A}})$ then

Case 3: If for all $S$, $\|\nabla f_{S}(w_{S})\|\leq O(\epsilon^{\prime\prime}_{\mathcal{A}})$, then

Suppose our loss functions are PL and our algorithm $\mathcal{A}$ is an oracle that returns a global optimizer $w_{S}^{*}$. Then the terms $\epsilon_{\mathcal{A}},\epsilon^{\prime}_{\mathcal{A}},\epsilon^{\prime\prime}_{\mathcal{A}}$ above are all identical to , leading to the following corollary.

Let $f_{S}$ satisfy the PL inequality with parameter $\mu$ and let $\mathcal{A}(S)=\text{arg}\min_{w\in\mathcal{X}}f_{S}(w)$. Then, $\mathcal{A}$ has pointwise hypothesis stability with

Bousquet and Ellisseef considered the stability of empirical risk minimizers where the loss function satisfied strong convexity. Their work implies that for $\lambda$-strongly convex functions, the empirical risk minimizer has we stability satisfying $\epsilon_{stab}\leq\frac{L^{2}}{\lambda n}$. Since $\lambda$-strongly convex implies $\lambda$-PL, Corollary 3.4 generalizes their result, with only a constant factor loss.

A similar result to Theorem 3.1 can be derived for empirical risk functions satisfy the QG condition and are realizable, e.g., where zero training loss is achievable.

Case 1: If for all $S$, $\|w_{S}-w_{S}^{*}\|\leq O(\epsilon_{\mathcal{A}})$ then

Case 2: If for all $S$, $|f_{S}(w_{S})-f_{S}(w_{S}^{*})|\leq O(\epsilon^{\prime}_{\mathcal{A}})$ then

Observe that unlike the case of PL empirical losses, QG empirical losses only allow for a $O(\frac{1}{\sqrt{n}})$ convergence rate of stability. Moreover, similarly to our result for PL loss functions, the result of Theorem 3.3 holds even if we only have information about the convergence of $\mathcal{A}$ in expectation.

Finally, we can obtain the following Corollary for empirical risk minimizers.

Let $f_{S}$ satisfy the QG inequality with parameter $\mu$ and let $\mathcal{A}(S)=\text{arg}\min_{w\in\mathcal{X}}f_{S}(w)$. Then, $\mathcal{A}$ has pointwise hypothesis stability with

2 Uniform Stability for PL/QG Loss Functions

Under a more restrictive setup, we can obtain similar bounds for uniform hypothesis stability, which is a stronger stability notion compare to its pointwise hypothesis variant. The usefulness of uniform stability compared to pointwise stability, is that it can lead to generalization bounds that concentrate exponentially faster with respect to the sample size $n$.

As before, given a data set $S$, we let denote $w_{S}$ be the model that $\mathcal{A}$ outputs. Let $\pi_{S}(w)$ denote the closest optimal point of $f_{S}$ to $w$. We will denote $\pi_{S}(w_{S})$ by $w_{S}^{*}$. Let $S,S^{\prime}$ be data sets differing in at most one entry. We will make the following technical assumption:

The empirical risk minimizers for $f_{S}$ and $f_{S^{\prime}}$, i.e., $w_{S}^{*},w_{S^{\prime}}^{*}$ satisfy $\pi_{S}(w_{S^{\prime}}^{*})=w_{S}^{*}$, where $\pi_{S}(w)$ is the projection of $w$ on the set of empirical risk minimizers of $f_{S}$. Note that this is satisfied if for every data set $S$, there is a unique minimizer $w^{*}_{S}$.

We would like to note that the above assumption is extremely strict, and in general does not apply to empirical losses with infinitely many global minima. To tackle the existence of infinitely many global minima, one could imagine designing $\mathcal{A}(S)$ to output a structured empirical risk minimizer, e.g., one such that if $\mathcal{A}$ is applied on $S^{\prime}$, its projection on the optima of $f_{S}$ would always yield back $A(S)$. This could be possible, if $A(S)$ corresponded to minimizing instead a regularized, or structure constrained cost function whose set of optimizers only contained a small subset of the global minima of $f_{S}$. Unfortunately, coming up with such a structured empirical risk minimizer for general nonconvex losses seems far from straightforward, and serves as an interesting open problem.

Assume that for all $S$, $f_{S}$ satisfies the PL condition with constant $\mu$, and suppose that Assumption 1 holds. Then $\mathcal{A}$ has uniform stability with parameter $\epsilon_{stab}$ satisfying the following conditions.

Case 1: If for all $S$, $\|w_{S}-w_{S}^{*}\|\leq O(\epsilon_{\mathcal{A}})$ then

Case 2: If for all $S$, $|f_{S}(w_{S})-f_{S}(w_{S}^{*})|\leq O(\epsilon^{\prime}_{\mathcal{A}})$ then

Case 3: If for all $S$, $\|\nabla f_{S}(w_{S})\|\leq O(\epsilon^{\prime\prime}_{\mathcal{A}}s)$, then

Since strong convexity is a special case of PL, this theorem implies that if we run enough iterations of a convergent algorithm $\mathcal{A}$ on a $\lambda$-strongly convex loss function, then we would expect uniform stability on the order of

In particular, this theorem recovers the stability estimates for ERMs and SGD applied to strongly convex functions proved in and , respectively.

In order to make this result more generally applicable, we would like to extend the theorem to a larger class of functions than just globally PL functions. If we assume boundedness of the loss function, then we can derive a similar result for globally QG functions. This leads us to the following theorem:

Case 1: If for all $S$, $\|w_{S}-w_{S}^{*}\|\leq O(\epsilon_{\mathcal{A}})$ then for all $z$ we have:

Case 2: If for all $S$, $|f_{S}(w_{S})-f_{S}^{*}|\leq O(\epsilon^{\prime}_{\mathcal{A}})$ then for all $z$ we have:

By analogous reasoning to that in Remark 3.1, both Theorem 3.5 and Theorem 3.6 hold if you only have information about the output of $\mathcal{A}$ in expectation.

PL loss functions in practice

As the bounds above show, the PL and QG conditions are sufficient for algorithmic stability and therefore imply good generalization. In this section, we show that the PL condition actually arises in some interesting machine learning setups, including least squares minimization, strongly convex functions composed with piecewise linear functions, and neural networks with linear activation functions. A first step towards a characterization of PL loss functions was proved by Karimi et al. , which established that the composition of a strongly-convex function and a linear function results in a loss that satisfies the PL condition.

Let $g$ be strongly-convex with parameter $\lambda$, $\sigma$ a leaky ReLU activation function with slopes $c1$ and $c_{2}$, and $X$ a matrix with minimum singular value $\sigma_{\min}(X)$. Let $c=\min\{|c_{1}|,|c_{2}|\}$. Then $f(w)=g(\sigma(Xw))$ is PL almost everywhere with parameter $\mu=\lambda\sigma_{\min}(X)^{2}c^{2}$.

In particular, 1-layer neural networks with a squared error loss and leaky ReLU activations satisfy the PL condition. More generally, this holds for any piecewise-linear activation function with slopes $\{c_{i}\}_{i=1}^{k}$. As long as each slope is non-zero and $X$ is full rank, the result above shows that the PL condition is satisfied.

2 Linear Neural Networks

The results above only concern one layer neural networks. Given the prevalence of deep networks, we would like to say something about the associated loss function. As it turns out, we can prove that a PL inequality holds in large regions of the parameter space for deep linear networks.

Suppose that the $W_{i}$ satisfy $\sigma_{\min}(W_{i})\geq\tau>0$ for all $i$. Then,

Combining Lemmas 4.2 and 4.3, we derive the following interesting corollary about when critical points are global minimizers. This result is not directly related to the work above, but gives an easy way to understand the landscape of critical points of deep linear networks.

Thematically similar results have been derived previously for 1 layer networks in and for deep neural networks in . In , Kawaguchi derives a similar result to ours for deep linear neural networks. Kawaguchi shows that every critical point is either a global minima or a saddle point. Our result, by contrast, implies that all full-rank critical points are global minima.

Lemmas 4.2 and 4.3 can also be combined to show that linear networks satisfy the PL condition in large regions of parameter space, as the follwing theorem says.

Stability of Some First-order Methods

We wish to apply our bounds from the previous section to popular convergent gradient-based methods. We consider SGD, GD, RCD, and SVRG. When we have $L$-Lipschitz, $\mu$-PL loss functions $f_{S}$ and $n$ training examples, Theorem 3.5 states that any learning algorithm $\mathcal{A}$ has uniform stability $\epsilon_{stab}$ satisfying

Here, $\epsilon_{\mathcal{A}}$ refers to how quickly $\mathcal{A}$ converges to the optimal value of the loss function. Specifically, this holds if the algorithm produces a model $w_{S}$ satisfying $|f_{S}(w_{s})-f_{S}^{*}|\leq O(\epsilon_{\mathcal{A}})$. For example, if we want to guarantee that our algorithm has the same stability as SGD in the strongly convex case, then we need to determine how many iterations $T$ we need to perform such that

The convergence rates of SGD, GD, RCD, and SVRG have been studied extensively in the literature . The results are given below. When necessary to state the result, we assume a constant step-size of $\gamma$. Figure 1 below summarizes the values of $\epsilon_{\mathcal{A}}$ for $T$ iterations of SGD, GD, RCD, and SVRG applied to $\lambda$-strongly convex loss functions and $\mu$-PL loss functions. Note that if Eq. (1) holds, then Corollary 3.4 implies that our algorithm is uniformly stable with parameter $\epsilon_{stab}=O(L^{2}/\mu n)$. Moreover, this is the same stability as that of SGD for strongly-convex functions , and saddle point avoiding algorithms on strict-saddle loss functions .

We use the above convergence rates of these algorithms in the $\lambda$-strongly convex and $\mu$-PL settings to determine how many iterations are required such that we get stability that is $O(L^{2}/\mu n)$. The results are summarized in Figure 2 below.

Note that in the $\mu$-PL case, although it is a nonconvex setup the above algorithms all exhibit the same stability for these values of $T$. This is not the case in general: several studies have observed that small-batch SGD offers superior generalization performance compared to large-batch SGD, or full-batch GD, when training deep neural networks .

Unfortunately our bounds above, are not nuanced enough to capture the difference in generalization performance between mini-batch and large-batch SGD. Below, we will make this observation formal. Although SGD and GD can be equally stable for nonconvex problems satisfying the PL condition, there exist nonconvex problems where full-batch GD is not stable and SGD is stable.

The Instability of Gradient Descent

In , Hardt et al. proved bounds on the uniform stability of SGD. They also noted that GD does not appear to be provably as stable for the nonconvex case and sketched a situation in which this difference would appear. Due to the similarity of SGD and GD, one may expect similar uniform stability. While this is true in every convex setting as we show in the Appendix, subsection A.6 (without requiring strong convexity), this breaks down in the non-convex setting. Below we construct an explicit example where GD is not uniformly stable, but SGD is. This example formalizes the intuition given in .

Intuitively, this is a generalized quadratic model where the predicted label $\hat{y}$ for a given $x$ is given by

The above predictive model can be described by the following 1-layer neural network using a quadratic and a linear activation function, denoted $z^{2}$ and $z$.

where $z_{i}=(-1,1)$ for $1\leq i\leq\frac{n-1}{2}$, $z_{i}=(-1/2,1)$ for $\frac{n-1}{2}<i\leq n-1$. By construction, we have

Then $f_{S}(w)$, $f_{S^{\prime}}(w)$ will approximately have the shape of the right-most function in Figure 4 above. However, recall that $\frac{d}{dw}g(w)=0$ at $w=\hat{w}$. Therefore, there is some $\delta$, with $0<\delta<\epsilon$, such that for all $w\in(\hat{w}-\delta,\hat{w}+\delta)$,

After enough iterations of gradient descent on $f_{S},f_{S^{\prime}}$, we will obtain models $w_{S}$ and $w_{S^{\prime}}$ that are close to the distinct local minima in the right-most graph in Figure 4. This will hold as long as $\gamma$ is not extremely large, in which case the steps of gradient descent could simply jump from one local minima to another. To ensure this does not happen, we restrict to $\gamma\leq 1$.

For all $n$, for all step-sizes $0<\gamma\leq 1$, there is a $K$ such that for all $k\geq K$, there are data sets $S,S^{\prime}$ of size $n$ differing in one entry and a non-zero measure set of initial starting points such that if we perform $k$ iterations of gradient descent with step-size $\gamma$ on $S$ and $S^{\prime}$ to get outputs $\mathcal{A}(S),\mathcal{A}(S^{\prime})$ then there is a $z^{*}$ such that

Theorem 6.1 establishes that there exist simple non-convex settings, for which the uniform stability of gradient descent does not decrease with $n$. In light of the work in , where the authors show that for very conservative step-sizes, SGD is stable on non-convex loss function, we might wonder whether SGD is stable in this setting with moderate step-sizes. We know that gradient descent is not stable, by Theorem 6.1, for $\gamma=1$. For simplicity of analysis, we focus on the case where $\gamma=1$.

Suppose we run SGD on the above $f_{S},f_{S^{\prime}}$ with step-size $1$ and initialize near $\hat{w}$ (we will be more concrete later about where we initialize). With probability $\frac{n-1}{n}$, the first iteration of SGD will use the same example for both $S$ and $S^{\prime}$, either $z=(-1,1)$ or $z=(-1/2,1)$. Computing derivatives at $\hat{w}$ shows

In both cases, the slope is at least $0.4$. Therefore, there is some $\eta$ such that for all $w\in[\hat{w}-\eta,\hat{w}+\eta]$,

Therefore, continuing to run SGD in this setting, even if we now decrease the step-size, will eventually lead us to the same basins of $f_{S}(w),f_{S^{\prime}}(w)$. Let $w_{S},w_{S^{\prime}}$ denote the outputs of SGD in this setting after enough steps so that we get convergence to within $\frac{1}{n}$ of a local minima. If our first sample $z$ was $(-1,1)$, we will end up in the right basin, while if our first sample $z$ was $(-1/2,1)$, we will end up in the left basin. In particular, for $\epsilon$ small, $z_{\pm}$ are close enough that the minima of $f_{S},f_{S^{\prime}}$ are within $\frac{1}{n}$ of each other. Note that the minima $w_{1},w_{2}$ that $w_{S},w_{S^{\prime}}$ are converging to are different. However, because they are in the same basin we know that for $\epsilon$ small, $z_{\pm}$ are close enough that $\|w_{S}-w_{S^{\prime}}\|\leq O(\frac{1}{n})$. Therefore, we have

Suppose that we initialize SGD in $[\hat{w}-\eta,\hat{w}+\eta]$ with a step-size of $\gamma=1$. Let $\mathcal{A}(S),\mathcal{A}(S^{\prime})$ denote the output of SGD after $k$ iteration for sufficiently large $k$. For $\|z\|\leq 2$,

This is in stark contrast to gradient descent, which is unstable in this setting. While work in suggests that this stability of SGD even in non-convex settings is a more general phenomenon, proving that this holds remains an open problem.

Conclusion

The success of machine learning algorithms in practice is often dictated by their ability to generalize. While recent work has developed great insight in to the training error of machine learning algorithms, much less is understood about their generalization error. Most work up to this point has either focused on specific algorithms or has made assumptions on the loss function (such as strong convexity) that may not be true in practice. By analyzing stability as coming from the convergence of an algorithm to global minima and the geometry surrounding them, we are able to derive much broader results. We develop easy-to-use stability results that encompass a general class of non-convex settings, and some of those can appear in interesting learning setups. Our bounds establish the stability for SGD, GD, SVRG, and RCD that is quantitatively comparable to more specialized results made in the past. Although our bounds are not nuanced enough to explain the generalization of mini-batch SGD compared to large-batch SGD, or full-batch GD, we hope that the generality of our nonconvex bounds serves as a step towards developing stability analyses that may help in understanding the generalization performance of practical machine learning algorithms.

There are still many exciting open problems concerning the stability and generalization of machine learning and optimization algorithms. We give a few below.

Stability for non-convex loss functions: While our results establish the stability of learning algorithms in some non-convex scenarios, it is unclear how to extend them directly to more general non-convex loss functions. Due to the wide variety of similar but not identical algorithms in machine learning, it would be particularly interesting to derive black-box results on the stability of learning algorithms for general non-convex loss functions, even when it concerns convergence to approximate local minima. In general, we expect these to be a function of the geometry of the loss function and the convergence of the algorithm, but it is quite possible that there are other factors that more directly control stability.

Generalization error of local minima: In general, we cannot expect a loss function to have only global minima. The question remains, among all local minima of a loss function, which one has the smallest generalization gap? The local minima with the smallest generalization error may not be a global minimizer. Even if we restrict to global minimizers, these may have different generalization errors. Is there a simple geometric characterization of their generalization error? This question may be connected to how sharp the loss function looks nearby the point. Steeper loss functions imply that small perturbations can greatly change the error, which suggests smaller generalization error. Theoretical results concerning this fact could be useful in stability analysis and the design of algorithms.

Generalization of SGD vs GD: We showed above that there are settings in which gradient descent is not uniformly stable, but SGD is. Empirically, SGD leads to small generalization error in neural networks . Theoretically, it is unclear how widespread this phenomenon is. Does SGD actually lead to more generalizable models than gradient descent? If so, why? How does this compare to other variants of SGD? Any theoretical results concerning the difference in generalization between SGD and other related algorithms could be extremely useful in the design of training algorithms and for understanding the empirical success of neural networks trained by SGD.

The geometry of critical points in real neural networks: We show above that linear neural networks obey relatively nice conditions on their local minima. In particular, as long as we restrict to weights that are full rank, all critical points are actually global minima. However, linear neural networks are not very exciting (e.g., they just correspond to linear models). Developing general theorems concerning the geometry underlying the critical points of neural networks is a challenging but very interesting open problem.

References

Appendix A Omitted Proofs

We have the following equivalent definition of PL due to Karimi et al.

The PL condition is equivalent to the condition that there is some constant $\mu>0$ such that for all $x$, $\|\nabla f(x)\|\geq\mu\|x_{p}-x\|$.

In this same paper, Karimi et al. show that PL functions also satisfy QG.

The PL condition implies the QG condition.

A.2 Proof of Theorem 3.1

Fix a training set $S$ and $i\in\{1,\ldots,n\}$. We will show pointwise hypothesis stability for all $S,i$ instead of for them in expectation. Let $w_{1}$ denote the output of $\mathcal{A}$ on $S$, and let $w_{2}$ denote the output of $\mathcal{A}$ on $S^{i}$. Let $w_{1}^{*}$ denote the critical point of $f_{S}$ to which $w_{1}$ is approaching, and $w_{2}^{*}$ denote the critical point of $f_{S^{i}}$ that $w_{2}$ is approaching. We then have,

We first wish to bound the first and third terms of (3). The bound depends on the case in Theorem 3.1.

Case 2: As stated in Lemma A.2, the PL condition implies the QG condition. Therefore,

By assumption on case 2, $|f_{S}(w_{1})-f_{S}(w_{1}^{*})|\leq O(\epsilon^{\prime}_{\mathcal{A}})$. This implies

Case 3: By Lemma A.1, the PL condition on $w_{1},w_{1}^{*}$ implies that

Using the fact that $f_{S}$ is $L$-Lipschitz and the fact that $\|\nabla f_{S}(w_{j})\|\leq O(\epsilon^{\prime\prime}_{\mathcal{A}})$ by assumption on Case 3, we find

By the PL condition, we can find a local minima $u$ of $f_{S}$ such that

Similarly, we can find a local minima $v$ of $f_{S^{i}}$ such that

Note that since $\nabla f_{S^{i}}(w_{2}^{*})=0$, we get:

Similarly, since $\nabla f_{S}(w_{1}^{*})=0$, we get:

Since all local minima of a PL function are global minima, we obtain

A.3 Proof of Theorem 3.3

Fix a training set $S$ and $i\in\{1,\ldots,n\}$. We will show pointwise hypothesis stability for all $S,i$ instead of for them in expectation. Let $w_{1}$ denote the output of $\mathcal{A}$ on $S$, and let $w_{2}$ denote the output of $\mathcal{A}$ on $S^{i}$. Let $w_{1}^{*}$ denote the critical point of $f_{S}$ to which $w_{1}$ is approaching, and $w_{2}^{*}$ denote the critical point of $f_{S^{i}}$ that $w_{2}$ is approaching. We then have,

We first wish to bound the first and third terms of (8). The bound depends on the case in Theorem 3.3.

By assumption on case 2, $|f_{S}(w_{1})-f_{S}(w_{1}^{*})|\leq O(\epsilon^{\prime}_{\mathcal{A}})$. This implies

Note that $|f_{S}(w_{1}^{*})-f_{S}(v)|=0$ by our realizability assumption. By assumption on $w_{1}^{*},w_{2}^{*}$, we know that $f_{S}(w_{1}^{*})\leq f_{S}(w_{2}^{*})$ and $f_{S^{i}}(w_{2}^{*})\leq f_{S^{i}}(w_{1}^{*})$. Some simple analysis shows

A.4 Proof of Theorem 3.5

Let $S_{1}=\{z_{1},\ldots,z_{n}\},S_{2}=\{z_{1},\ldots,z_{n-1},z_{n}^{\prime}\}$ be data sets of size $n$ differing only in one entry. Let $w_{i}$ denote the output of $\mathcal{A}$ on data set $S_{i}$ and let $w_{i}^{*}$ denote $w_{S_{i}}^{*}$. Let $f_{i}(w)=f_{S_{i}}(w)$.

Note that by A1, we know that $w_{1}^{*}$ is the closest optimal point of $f_{1}$ to $w_{2}^{*}$. By the PL condition,

This bounds the second term of 12. The first and third terms must be bounded differently depending on the case.

Case 1: By assumption, for $i=1,2$, $\|w_{i}-w_{i}^{*}\|=O(\epsilon_{\mathcal{A}})$, proving the result.

Case 2: As stated in Lemma A.2, PL implies QG. Therefore for $i=1,2$,

Case 3: As mentioned above in Lemma A.1, the PL condition implies that for $i=1,2$,

Since $\|\nabla f_{i}(w_{i})\|\leq O(\epsilon^{\prime\prime}_{\mathcal{A}})$, we get the desired result.∎

A.5 Proof of Theorem 3.6

Note that by A1, we know that $w_{1}^{*}$ is the closest optimal point of $f_{1}$ to $w_{2}^{*}$. By Lemma A.3,

This bounds the second term of 12. The first and third terms must be bounded differently depending on the case.

Case 1: By assumption, for $i=1,2$, $\|w_{i}-w_{i}^{*}\|=O(\epsilon_{\mathcal{A}})$, proving the result.

Case 2: By the QG property, we get that for $i=1,2$,

A.6 Stability of Gradient Descent for Convex Loss Functions

To prove the stability of gradient descent, we will assume that the underlying loss function is smooth.

In , Hardt et al. show the following theorem.

Performing similar analysis for gradient descent, we obtain the following theorem.

To prove this theorem, we use similar techniques to those in . We first consider the convex case.

So, if $\frac{\gamma_{t}}{n}\leq\frac{2}{\beta}$ for all $t$, we get:

We now move to the $\lambda$-strongly convex case. For simplicity of analysis, we assume that we use a constant step size $\gamma$ such that $\gamma\leq 1/\beta$.

The proof remains the same, except when using co-coercivity. Under this assumption, some plug and play in an analogous fashion will show:

Note that if $\gamma\leq\frac{1}{\beta}$ then the second term is nonnegative and one can show that this implies:

This implies uniform stability with parameter $\frac{2L^{2}}{\lambda n}$.∎

A.7 Proof of Theorem 4.1

For almost all $w$, we can write $\sigma(Xw)$ as $\text{diag}(b)Xw$ for a vector $b$ where $b_{i}(Xw)_{i}=\sigma((Xw)_{i})$ (this only exludes points $w$ such that $(Xw)_{i}$ is on a cusp of the piecewise-linear function). Then in an open neighborhood of such an $w$, we find:

For a given $w$, let $w_{p}$ be the closest global minima of $f$ (i.e., the closest point such that $f^{*}=f(w_{p})$). By strong convexity of $g$, we find:

Note that the minimum singular value of $\text{diag}(b)$ is the square root of the minimum eigenvalue of $\text{diag}(b)^{2}$. Since $\text{diag}(b)^{2}$ has entries $c_{i}^{2}$ on the diagonal, we know that the minimum singular value is at least $c=\min_{i}\{|c_{i}|\}$. Therefore we get:

A.8 Proof of Lemma 4.2

Using basic properties of the Frobenius norm and the definition of the pseudo-inverse, we have

This last step follows by basic properties of the pseudo-inverse. By the triangle inequality,

Note that $YX^{+}X-Y$ is the component of $Y$ that is orthogonal to the row-space of $X$, while $YX^{+}X$ is the projection of $Y$ on to this row space. Therefore, $YX^{+}X-Y$ is orthogonal to $WX-YX^{+}X$ with respect to the trace inner product. Therefore, the inequality above is actually an equality, that is

A.9 Proof of Lemma 4.3

Our proof uses similar techniques to that of Hardt and Ma . We wish to compute the gradient of $f$ with respect to a matrix $W_{j}$. One can show the following:

By assumption, $\sigma_{\min}(W_{j})\geq\tau$. Therefore:

Taking the gradient with respect to all $W_{i}$ we get:

In this subsection we show that a provable rate for SGD on smooth functions with learning rate proportional to $O(1/t)$ might require a large number of iterations. While show stability of SGD in non-convex settings with such a step-size, their stability bounds grow close to linearly with the number of iterations. When exponentially many steps are taken, this no longer implies useful generalization bounds on SGD.

When a function $f(x)=\sum_{i=1}^{n}f_{i}(x)$ is $\beta$ smooth on its domain, then the following holds:

Fix some $t$ and let $x=x_{t+1}=x_{t}-\gamma_{t}\nabla f_{s_{t}}(x_{t})$ and let $y=x_{t}$. Here, $x_{0}$ is set to some initial vector value, $s_{t}$ is a uniform iid sample from $\{1,\ldots,n]$, and $\gamma_{t}=c/t$ for some constant $c>0$. Then, due to the $\beta$-smoothness of $f$ we have the following:

Taking expectation with respect to all random samples $s_{t}$, yields

Summing the above inequality for all $t$ terms from 0 to $T$ we get

where $C_{1}$ is a universal constant that depends only on $c$. Observe that even if $C_{2}=0$, using the above simple bounding technique (a simplified version of the nonconvex convergence bounds of ), requires $O(e^{-\epsilon})$ steps to reach error $\epsilon$, i.e., an exponentially large number of steps.

We note that the above bound does not imply that there does not exist a smooth function for which $1/t$ stepsizes suffice for polynomial-time convergence (in fact there are several convex problems for wich $1/t$ suffices for fast convergence). However, the above implies that when we are only assuming smoothness on a nonconvex function, it may be the case that there exist nonconvex problems where $1/t$ implies exponentially slow convergence.

This implies that if the optimal model is at distance $\Omega(Md)$ from $x_{0}$, we would require at least $O(e^{M\cdot d})$ iterations to reach it in expectation.