Train faster, generalize better: Stability of stochastic gradient descent

Introduction

The most widely used optimization method in machine learning practice is stochastic gradient method (SGM). Stochastic gradient methods aim to minimize the empirical risk of a model by repeatedly computing the gradient of a loss function on a single training example, or a batch of few examples, and updating the model parameters accordingly. SGM is scalable, robust, and performs well across many different domains ranging from smooth and strongly convex problems to complex non-convex objectives.

In a nutshell, our results establish that:

Any model trained with stochastic gradient method in a reasonable amount of time attains small generalization error.

As training time is inevitably limited in practice, our results help to explain the strong generalization performance of stochastic gradient methods observed in practice. More concretely, we bound the generalization error of a model in terms of the number of iterations that stochastic gradient method took in order to train the model. Our main analysis tool is to employ the notion of algorithmic stability due to Bousquet and Elisseeff . We demonstrate that the stochastic gradient method is stable provided that the objective is relatively smooth and the number of steps taken is sufficiently small.

It is common in practice to perform a linear number of steps in the size of the sample and to access each data point multiple times. Our results show in a broad range of settings that, provided the number of iterations is linear in the number of data points, the generalization error is bounded by a vanishing function of the sample size. The results hold true even for complex models with large number of parameters and no explicit regularization term in the objective. Namely, fast training time by itself is sufficient to prevent overfitting.

Our bounds are algorithm specific: Since the number of iterations we allow can be larger than the sample size, an arbitrary algorithm could easily achieve small training error by memorizing all training data with no generalization ability whatsoever. In contrast, if the stochastic gradient method manages to fit the training data in a reasonable number of iterations, it is guaranteed to generalize.

Conceptually, we show that minimizing training time is not only beneficial for obvious computational advantages, but also has the important byproduct of decreasing generalization error. Consequently, it may make sense for practitioners to focus on minimizing training time, for instance, by designing model architectures for which stochastic gradient method converges fastest to a desired error level.

Our focus is on generating generalization bounds for models learned with stochastic gradient descent. Recall that the generalization bound is the expected difference between the error a model incurs on a training set versus the error incurred on a new data point, sampled from the same distribution that generated the training data. Throughout, we assume we are training models using $n$ sampled data points.

Our results build on a fundamental connection between the generalization error of an algorithm and its stability properties. Roughly speaking, an algorithm is stable if the training error it achieves varies only slightly if we change any single training data point. The precise notion of stability we use is known as uniform stability due to . It states that a randomized algorithm $A$ is uniformly stable if for all data sets differing in only one element, the learned models produce nearly the same predictions. We review this method in Section 2, and provide a new adaptation of this theory to iterative algorithms.

In Section 3, we show that stochastic gradient is uniformly stable, and our techniques mimic its convergence proofs. For convex loss functions, we prove that the stability measure decreases as a function of the sum of the step sizes. For strongly convex loss functions, we show that stochastic gradient is stable, even if we train for an arbitrarily long time. We can combine our bounds on the generalization error of stochastic gradient method with optimization bounds quantifying the convergence of the empirical loss achieved by SGM. In Section 5, we show that models trained for multiple epochs match classic bounds for stochastic gradient .

2 Related work

There is a venerable line of work on stability and generalization dating back more than thirty years . The landmark work by Bousquet and Elisseeff introduced the notion of uniform stability that we rely on. They showed that several important classification techniques are uniformly stable. In particular, under certain regularity assumptions, it was shown that the optimizer of a regularized empirical loss minimization problem is uniformly stable. Previous work generally applies only to the exact minimizer of specific optimization problems. It is not immediately evident on how to compute a generalization bound for an approximate minimizer such as one found by using stochastic gradient. Subsequent work studied stability bounds for randomized algorithms but focused on random perturbations of the cost function, such as those induced by bootstrapping or bagging . This manuscript differs from this foundational work in that it derives stability bounds about the learning procedure, analyzing algorithmic properties that induce stability.

Stochastic gradient descent, of course, is closely related to our inquiry. Classic results by Nemirovski and Yudin show that the stochastic gradient method produces is nearly optimal for empirical risk minimization of convex loss functions . These results have been extended by many machine learning researchers, yielding tighter bounds and probabilistic guarantees . However, there is an important limitation of all of this prior art. The derived generalization bounds only hold for single passes over the data. That is, in order for the bounds to be valid, each training example must be used no more than once in a stochastic gradient update. In practice, of course, one tends to run multiple epochs of the stochastic gradient method. Our results resolve this issue by combining stability with optimization error. We use the foundational results to estimate the error on the empirical risk and then use stability to derive a deviation from the true risk. This enables us to study the risk incurred by multiple epochs and provide simple analyses of regularization methods for convex stochastic gradient. We compare our results to this related work in Section 5. We note that Rosasco and Villa obtain risk bounds for least squares minimization with an incremental gradient method in terms of the number of epochs . These bounds are akin to our study in Section 5, although our results are incomparable due to various different assumptions.

Finally, we note that in the non-convex case, the stochastic gradient method is remarkably successful for training large neural networks . However, our theoretical understanding of this method is limited. Several authors have shown that the stochastic gradient method finds a stationary point of nonconvex cost functions . Beyond asymptotic convergence to stationary points, little is known about finding models with low training or generalization error in the nonconvex case. There have recently been several important studies investigating optimal training of neural nets. For example Livni et al. show that networks with polynomial activations can be learned in a greedy fashion . Janzamin et al. show that two layer neural networks can be learned using tensor methods. Arora et al. show that two-layer sparse coding dictionaries can be learned via stochastic gradient. Our work complements these developments: rather than providing new insights into mechanisms that yield low training error, we provide insights into mechanisms that yield low generalization error. If one can achieve low training error quickly on a nonconvex problem with stochastic gradient, our results guarantee that the resulting model generalizes well.

Stability of randomized iterative algorithms

Consider the following general setting of supervised learning. There is an unknown distribution ${\cal D}$ over examples from some space $Z.$ We receive a sample $S=(z_{1},\dots,z_{n})$ of $n$ examples drawn i.i.d. from ${\cal D}.$ Our goal is to find a model $w$ with small population risk, defined as:

Here, where $f$ is a loss function and $f(w;z)$ designates the loss of the model described by $w$ encountered on example $z$.

Since we cannot measure the objective $R[w]$ directly, we instead use a sample-averaged proxy, the empirical risk, defined as

The generalization error of a model $w$ is the difference

When $w=A(S)$ is chosen as a function of the data by a potentially randomized algorithm $A$ it makes sense to consider the expected generalization error

where the expectation is over the randomness of $A$ and the sample $S.$

In order to bound the generalization error of an algorithm, we employ the following notion of uniform stability in which we allow randomized algorithms as well.

A randomized algorithm $A$ is $\epsilon$-uniformly stable if for all data sets $S,S^{\prime}\in Z^{n}$ such that $S$ and $S^{\prime}$ differ in at most one example, we have

We recall the important theorem that uniform stability implies generalization in expectation. Since our notion of stability differs slightly from existing ones with respect to the randomness of the algorithm, we include a proof for the sake of completeness. The proof is based on an argument in Lemma 7 of and very similar to Lemma 11 in .

[Generalization in expectation] Let $A$ be $\epsilon$-uniformly stable. Then,

Denote by $S=(z_{1},\dots,z_{n})$ and $S^{\prime}=(z_{1}^{\prime},\dots,z_{n}^{\prime})$ two independent random samples and let $S^{(i)}=(z_{1},\dots,z_{i-1},z_{i}^{\prime},z_{i+1},\dots,z_{n})$ be the sample that is identical to $S$ except in the $i$’th example where we replace $z_{i}$ with $z_{i}^{\prime}$. With this notation, we get that

Furthermore, taking the supremum over any two data sets $S,S^{\prime}$ differing in only one sample, we can bound the difference as

by our assumption on the uniform stability of $A.$ The claim follows. ∎

Theorem 2.2 proves that if an algorithm is uniformly stable, then its generalization error is small. We now turn to some properties of iterative algorithms that control their uniform stability.

We consider general update rules of the form $G\colon\Omega\to\Omega$ which map a point $w\in\Omega$ in the parameter space to another point $G(w).$ The most common update is the gradient update rule

The canonical update rule we will consider in this manuscript is an incremental gradient update, where $G(w)=w-\alpha\nabla f(w)$ for some convex function $f$. We will return to a detailed discussion of this specific update in the sequel, but the reader should keep this particular example in mind throughout the remainder of this section.

The following two definitions provide the foundation of our analysis of how two different sequences of update rules diverge when iterated from the same starting point. These definitions will ultimately be useful when analyzing the stability of stochastic gradient descent.

With these two properties, we can establish the following lemma of how a sequence of updates to a model diverge when the training set is perturbed.

Fix an arbitrary sequence of updates $G_{1},\dots,G_{T}$ and another sequence $G_{1}^{\prime},\dots,G_{T}^{\prime}.$ Let $w_{0}=w_{0}^{\prime}$ be a starting point in $\Omega$ and define $\delta_{t}=\|w_{t}^{\prime}-w_{t}\|$ where $w_{t},w^{\prime}_{t}$ are defined recursively through

The first bound on $\delta_{t}$ follow directly from the assumption that $G_{t}=G_{t}^{\prime}$ and the definition of expansiveness. For the second bound, recall from Definition 2.4 that if $G_{t}$ and $G_{t}^{\prime}$ are $\sigma$-bounded, then by the triangle inequality,

which gives half of the second bound. We can alternatively bound $\delta_{t+1}$ as

Stability of Stochastic Gradient Method

Given $n$ labeled examples $S=(z_{1},\dots,z_{n})$ where $z_{i}\in Z,$ consider a decomposable objective function

where $f(w;z_{i})$ denotes the loss of $w$ on the example $z_{i}.$ The stochastic gradient update for this problem with learning rate $\alpha_{t}>0$ is given by

Stochastic gradient method (SGM) is the algorithm resulting from performing stochastic gradient updates $T$ times where the indices $i_{t}$ are randomly chosen. There are two popular schemes for choosing the examples’ indices. One is to pick $i_{t}$ uniformly at random in $\{1,\dots,n\}$ at each step. The other is to choose a random permutation over $\{1,\dots,n\}$ and cycle through the examples repeatedly in the order determined by the permutation. Our results hold for both variants.

In parallel with the previous section the stochastic gradient method is akin to applying the gradient update rule defined as follows.

There are two cases to consider. In the first case, SGM selects the index of an example at step $t$ on which is identical in $S$ and $S^{\prime}$. Unfortunately, it could still be the case that $\delta_{t}$ grows, since $w_{t}$ and $w_{t}^{\prime}$ differ and so the gradients at these two points may still differ. Below, we will show how to control $\delta_{t}$ in terms of the convexity and smoothness properties of the stochastic gradients.

The second case to consider is when SGM selects the one example to update in which $S$ and $S^{\prime}$ differ. Note that this happens only with probability $1/n$ if examples are selected randomly. In this case, we simply bound the increase in $\delta_{t}$ by the norm of the two gradient $\nabla f(w_{t-1};z)$ and $\nabla f(w_{t-1}^{\prime};z^{\prime}).$ The sum of the norms is bounded by $2\alpha_{t}L$ and we obtain $\delta_{t}\leq\delta_{t}+2\alpha_{t}L.$ Combining the two cases, we can then solve a simple recurrence relation to obtain a bound on $\delta_{T}.$

This simple approach suffices to obtain the desired result in the convex case, but there are additional difficulties in the non-convex case. Here, we need to use an intriguing stability property of stochastic gradient method. Specifically, the first time step $t_{0}$ at which SGM even encounters the example in which $S$ and $S^{\prime}$ differ is a random variable in $\{1,\dots,n\}$ which tends to be relatively large. Specifically, for any $m\in\{1,\dots,n\},$ the probability that $t_{0}\leq m$ is upper bounded by $m/n.$ This allows us to argue that SGM has a long “burn-in period” where $\delta_{t}$ does not grow at all. Once $\delta_{t}$ begins to grow, the step size has already decayed allowing us to obtain a non-trivial bound.

We now turn to making this argument precise.

2 Expansion properties of stochastic gradients

Let us now record some of the core properties of the stochastic gradient update. The gradient update rule is bounded provided that the function $f$ satisfies the following common Lipschitz condition.

We say that $f$ is $L$-Lipschitz if for all points $u$ in the domain of $f$ we have $\|\nabla f(x)\|\leq L.$ This implies that

Assume that $f$ is $L$-Lipschitz. Then, the gradient update $G_{f,\alpha}$ is $(\alpha L)$-bounded.

By our Lipschitz assumption, $\|w-G_{f,\alpha}(w)\|=\|\alpha\nabla f(w)\|\leq\alpha L\,$ . ∎

We now turn to expansiveness. As we will see shortly, different expansion properties are achieved for non-convex, convex, and strongly convex functions.

The following standard notion of smoothness leads to a bound on how expansive the gradient update is.

In general, smoothness will imply that the gradient updates cannot be overly expansive. When the function is also convex and the step size is sufficiently small the gradient update becomes non-expansive. When the function is additionally strongly convex, the gradient update becomes contractive in the sense that $\eta$ will be less than one and $u$ and $v$ will actually shrink closer to one another. The majority of the following results can be found in several textbooks and monographs. Notable references are Polyak and Nesterov . We include proofs in the appendix for completeness.

Assume that $f$ is $\beta$-smooth. Then, the following properties hold.

$G_{f,\alpha}$ is $(1+\alpha\beta)$-expansive.

Assume in addition that $f$ is convex. Then, for any $\alpha\leq 2/\beta,$ the gradient update $G_{f,\alpha}$ is $1$-expansive.

Assume in addition that $f$ is $\gamma$-strongly convex. Then, for $\alpha\leq\frac{2}{\beta+\gamma}$, $G_{f,\alpha}$ is $\left(1-\frac{\alpha\beta\gamma}{\beta+\gamma}\right)$-expansive.

Henceforth we will no longer mention which random selection rule we use as the proofs are almost identical for both rules.

3 Convex optimization

We begin with a simple stability bound for convex loss minimization via stochastic gradient method.

Assume that the loss function $f(\cdot\,;z)$ is $\beta$-smooth, convex and $L$-Lipschitz for every $z.$ Suppose that we run SGM with step sizes $\alpha_{t}\leq 2/\beta$ for $T$ steps. Then, SGM satisfies uniform stability with

Let $S$ and $S^{\prime}$ be two samples of size $n$ differing in only a single example. Consider the gradient updates $G_{1},\dots,G_{T}$ and $G_{1}^{\prime},\dots,G_{T}^{\prime}$ induced by running SGM on sample $S$ and $S^{\prime},$ respectively. Let $w_{T}$ and $w_{T}^{\prime}$ denote the corresponding outputs of SGM.

We now fix an example $z\in Z$ and apply the Lipschitz condition on $f(\cdot\,;z)$ to get

where $\delta_{T}=\|w_{T}-w_{T}^{\prime}\|.$ Observe that at step $t,$ with probability $1-1/n,$ the example selected by SGM is the same in both $S$ and $S^{\prime}.$ In this case we have that $G_{t}=G_{t}^{\prime}$ and we can use the $1$-expansivity of the update rule $G_{t}$ which follows from Lemma 3.7.0 using the fact that the objective function is convex and that $\alpha_{t}\leq 2/\beta$. With probability $1/n$ the selected example is different in which case we use that both $G_{t}$ and $G_{t}^{\prime}$ are $\alpha_{t}L$-bounded as a consequence of Lemma 3.3. Hence, we can apply Lemma 2.5 and linearity of expectation to conclude that for every $t,$

Plugging this back into equation (3.3), we obtain

Since this bounds holds for all $S,S^{\prime}$ and $z,$ we obtain the desired bound on the uniform stability. ∎

4 Strongly Convex Optimization

In the strongly convex case we can bound stability with no dependence on the number of steps at all. Assume that the function $f(w;z)$ is strongly convex with respect to $w$ for all $z$. Let $\Omega$ be a compact, convex set over which we wish to optimize. Assume further that we can readily compute the Euclidean projection onto the set $\Omega$, namely, $\Pi_{\Omega}(v)=\arg\min_{w\in\Omega}\|w-v\|$. In this section we restrict our attention to the projected stochastic gradient method

A common application of the above iteration in machine learning is solving Tikhonov regularization problems. Specifically, the empirical risk is augmented with an additional regularization term,

where $f$ is as before a pre-specified loss function. We can assume without loss of generality that $f(0;\cdot)=1$. Then, the optimal solution of (3.6) must lie in a ball of radius $r$ about where $r=\sqrt{2/\mu}\,.$ This fact can be ascertained by plugging in $w=0$ and noting that the minimizer of (3.6) must have a smaller cost, thus $\frac{\mu}{2}\|w^{\star}\|^{2}\leq R_{S,\mu}[w^{\star}]\leq R_{S,\mu}=1.$ We can now define the set $\Omega$ to be the ball of radius $r$, in which case the projection is a simple scaling operation. Througout the rest of the section we replace $f(w;z)$ with its regularized form, namely,

which is strongly convex with parameter $\mu$. Similarly, we will overload the constant $L$ to by setting

Assume that the loss function $f(\cdot\,;z)$ is $\gamma$-strongly convex and $\beta$-smooth for all $z.$ Suppose we run the projected SGM iteration (3.5) with constant step size $\alpha\leq 1/\beta$ for $T$ steps. Then, SGM satisfies uniform stability with

The proof is analogous to that of Theorem 3.8 with a slightly different recurrence relation. We repeat the argument for completeness. Let $S$ and $S^{\prime}$ be two samples of size $n$ differing in only a single example. Consider the gradient updates $G_{1},\dots,G_{T}$ and $G_{1}^{\prime},\dots,G_{T}^{\prime}$ induced by running SGM on sample $S$ and $S^{\prime},$ respectively. Let $w_{T}$ and $w_{T}^{\prime}$ denote the corresponding outputs of SGM.

Denoting $\delta_{T}=\|w_{T}-w_{T}^{\prime}\|$ and appealing to the boundedness of the gradient of $f,$ we have

Observe that at step $t,$ with probability $1-1/n,$ the example selected by SGM is the same in both $S$ and $S^{\prime}.$ In this case we have that $G_{t}=G_{t}^{\prime}$. At this stage, note that

because Euclidean projection does not increase the distance between projected points (see Lemma 4.6 below for a generalization of this fact). We can now apply the following useful simplification of Lemma 3.7.0 if $\alpha\leq 1/\beta$: since $\frac{2\alpha\beta\gamma}{\beta+\gamma}\geq\alpha\gamma$ and $\alpha\gamma\leq 1$, $G_{f,\alpha}$ is $(1-\alpha\gamma)$-expansive. With probability $1/n$ the selected example is different in which case we use that both $G_{t}$ and $G_{t}^{\prime}$ are $\alpha{}M$-bounded as a consequence of Lemma 3.3. Hence, we can apply Lemma 2.5 and linearity of expectation to conclude that for every $t,$

Plugging the above inequality into equation (3.3), we obtain

Since this bounds holds for all $S,S^{\prime}$ and $z,$ the lemma follows. ∎

We would like to note that a nearly identical result holds for a “staircase” decaying step-size that is also popular in machine learning and stochastic optimization.

Assume that the loss function $f(\cdot\,;z)\in$ is $\gamma$-strongly convex has gradients bounded by $L$ as in (3.7), and is $\beta$-smooth function for all $z.$ Suppose we run SGM with step sizes $\alpha_{t}=\frac{1}{\gamma t}$. Then, SGM has uniform stability of

where $\rho=\sup_{w\in\Omega}\sup_{z}f(w;z)$.

Note that once $t>\frac{\beta}{\gamma}$, the iterates are contractive with contractivity $1-\alpha_{t}\gamma\leq 1-\frac{1}{t}$. Thus, for $t\geq t_{0}:=\frac{\beta}{\gamma}$, we have

Assuming that $\delta_{t_{0}}=0$ and expanding this recursion, we find:

Now, the result follows from Lemma 3.11 with the fact that $t_{0}=\frac{\beta}{\gamma}$. ∎

5 Non-convex optimization

In this section we prove stability results for stochastic gradient methods that do not require convexity. We will still assume that the objective function is smooth and Lipschitz as defined previously.

The crux of the proof is to observe that SGM typically makes several steps before it even encounters the one example on which two data sets in the stability analysis differ.

Assume that the loss function $f(\cdot\,;z)$ is nonnegative and $L$-Lipschitz for all $z.$ Let $S$ and $S^{\prime}$ be two samples of size $n$ differing in only a single example. Denote by $w_{T}$ and $w_{T}^{\prime}$ the output of $T$ steps of SGM on $S$ and $S^{\prime},$ respectively. Then, for every $z\in Z$ and every $t_{0}\in\{0,1,\dots,n\},$ under both the random update rule and the random permutation rule, we have

Let $S$ and $S^{\prime}$ be two samples of size $n$ differing in only a single example, and let $z\in Z$ be an arbitrary example. Consider running SGM on sample $S$ and $S^{\prime}$, respectively. As stated, $w_{T}$ and $w_{T}^{\prime}$ denote the corresponding outputs of SGM. Let ${\cal E}=\mathbf{1}\!\left[\delta_{t_{0}}=0\right]$ denote the event that $\delta_{t_{0}}=0.$ We have,

The second inequality follows from the Lipschitz assumption.

Under the random permutation rule, $I$ is a uniformly random number in $\{1,\dots,n\}$ and therefore

Assume that $f(\cdot;z)\in$ is an $L$-Lipschitz and $\beta$-smooth loss function for every $z.$ Suppose that we run SGM for $T$ steps with monotonically non-increasing step sizes $\alpha_{t}\leq c/t.$ Then, SGM has uniform stability with

In particular, omitting constant factors that depend on $\beta,$ $c,$ and $L,$ we get

Let $S$ and $S^{\prime}$ be two samples of size $n$ differing in only a single example. Consider the gradient updates $G_{1},\dots,G_{T}$ and $G_{1}^{\prime},\dots,G_{T}^{\prime}$ induced by running SGM on sample $S$ and $S^{\prime},$ respectively. Let $w_{T}$ and $w_{T}^{\prime}$ denote the corresponding outputs of SGM.

By Lemma 3.11, we have for every $t_{0}\in\{1,\dots,n\},$

Toward this goal, observe that at step $t,$ with probability $1-1/n,$ the example selected by SGM is the same in both $S$ and $S^{\prime}.$ In this case we have that $G_{t}=G_{t}^{\prime}$ and we can use the $(1+\alpha_{t}\beta)$-expansivity of the update rule $G_{t}$ which follows from our smoothness assumption via Lemma 3.7.0. With probability $1/n$ the selected example is different in which case we use that both $G_{t}$ and $G_{t}^{\prime}$ are $\alpha_{t}L$-bounded as a consequence of Lemma 3.3.

Hence, we can apply Lemma 2.5 and linearity of expectation to conclude that for every $t\geq t_{0},$

Here we used that $1+x\leq\exp(x)$ for all $x.$

Using the fact that $\Delta_{t_{0}}=0,$ we can unwind this recurrence relation from $T$ down to $t_{0}+1.$ This gives

Letting $q=\beta c,$ the right hand side is approximately minimized when

Since the bound we just derived holds for all $S,S^{\prime}$ and $z,$ we immediately get the claimed upper bound on the uniform stability. ∎

Stability-inducing operations

In light of our results, it makes sense to analyse for operations that increase the stability of the stochastic gradient method. We show in this section that pleasingly several popular heuristics and methods indeed improve the stability of SGM. Our rather straightforward analyses both strengthen the bounds we previously obtained and help to provide an explanation for the empirical success of these methods.

Weight decay is a simple and effective method that often improves generalization .

Let $f\colon\Omega\to\Omega,$ be a differentiable function. We define the gradient update with weight decay at rate $\mu$ as $G_{f,\mu,\alpha}(w)=(1-\alpha\mu)w-\alpha\nabla f(w).$

Assume that $f$ is $\beta$-smooth. Then, $G_{f,\mu,\alpha}$ is $(1+\alpha(\beta-\mu))$-expansive.

Let $G=G_{f,\mu,\alpha}.$ By triangle inequality and our smoothness assumption,

The above lemma shows as that a regularization parameter $\mu$ counters a smoothness parameter $\beta.$ Once $r>\beta,$ the gradient update with decay becomes contractive. Any theorem we proved in previous sections that has a dependence on $\beta$ leads to a corresponding theorem for stochastic gradient with weight decay in which $\beta$ is replaced with $\beta-\mu.$

Gradient Clipping.

It is common when training deep neural networks to enforce bounds on the norm of the gradients encountered by SGD. This is often done by either truncation, scaling, or dropping of examples that cause an exceptionally large value of the gradient norm. Any such heuristic directly leads to a bound on the Lipschitz parameter $L$ that appears in our bounds. It is also easy to introduce a varying Lipschitz parameter $L_{t}$ to account for possibly different values.

Dropout.

Dropout is a popular and effective heuristic for preventing large neural networks from overfitting. Here we prove that, indeed, dropout improves all of our stability bounds generically. From the point of view of stochastic gradient descent, dropout is equivalent to setting a fraction of the gradient weights to zero. That is, instead of updating with a stochastic gradient $\nabla f(w;z)$ we instead update with a perturbed gradient $D\nabla f(w;z)$ which is is typically identical to $\nabla f(w;z)$ in some of the coordinates and equal to on the remaining coordinates, although our definition is a fair bit more general.

As expected, dropout improves the effective Lipschitz constant of the objective function.

Assume that $f$ is $L$-Lipschitz. Then, the dropout update $DG_{f,\alpha}$ with dropout rate $s$ is $(s\alpha L)$-bounded.

By our Lipschitz assumption and linearity of expectation,

From this lemma we can obtain various corollaries by replacing $L$ with $sL$ in our theorems.

Projections and Proximal Steps.

Related to regularization, there are many popular updates which follow a stochastic gradient update with a projection onto a set or some statistical shrinkage operation. The vast majority of these operations can be understood as applying a proximal-point operation associated with a convex function. Similar to the gradient operation, we can define the proximal update rule.

An elementary proof of the following Lemma, due to Rockafellar , can be found in the appendix.

If $f$ is convex, the proximal update (4.1) is $1$-expansive.

In particular, this Lemma implies that the Euclidean projection onto a convex set is $1$-expansive. Note that in many important cases, proximal operators are actually contractive. That is, they are $\eta$-expansive with $\eta<1$. An notable example is when $f(\cdot)$ is the Euclidean norm for which the update rule is $\eta$-expansive with $\eta=(1+\alpha)^{-1}$. So stability can be induced by the choice of an appropriate prox-operation, which can always be interpreted as some form of regularization.

Model Averaging.

Model averaging refers to the idea of averaging out the iterates $w_{t}$ obtained by a run of SGD. In convex optimization, model averaging is sometimes observed to lead to better empirical performance of SGM and closely replated updates such as the Perceptron . Here we show that model averaging improves our bound for the convex optimization by a constant factor.

Let $\bar{w}_{T}=\frac{1}{T}\displaystyle\sum_{t=1}^{T}w_{t}$ denoet the average of the stochastic gradient iterates. Since

Using Lemma 3.8, the deviation between $\bar{w}_{t}$ and $\bar{w}_{t}^{\prime}$ obeys

Here the expectation is taken over the algorithm and hence the claim follows by our definition of uniform stability. ∎

Convex risk minimization

We now outline how our generalization bounds lead to bounds on the population risk achieved by SGM in the convex setting. We restrict our attention to the convex case where we can contrast against known results. The main feature of our results is that we show that one can achieve bounds comparable or perhaps better than known results on stochastic gradient for risk minimization by running for multiple passes over the data set.

The key to the analysis in this section is to decompose the risk estimates into an optimization error term and a stability term. The optimization error designates how closely we optimize the empirical risk or a proxy of the empirical risk. By optimizing with stochastic gradient, we will be able to balance this optimization accuracy against how well we generalize. These results are inspired by the work of Bousquet and Bottou who provided similar analyses for SGM based on uniform convergence . However, our stability results will yield sharper bounds.

Throughout this section, our risk decomposition works as follows. We define the optimization error to be the gap between the empirical risk and minimum empirical risk in expectation:

By Theorem 2.2, the expected risk of a $w$ output by SGM is bounded as

In general, the optimization error decreases with the number of SGM iterations while the stability increases. Balancing these two terms will thus provide a reasonable excess risk against the empirical risk minimizer. Note that our analysis involves the expected minimum empirical risk which could be considerably smaller than the minimum risk. However, as we now show, it can never be larger.

To analyze the optimization error, we will make use of a classical result due to Nemirovski and Yudin .

Assume we run stochastic gradient descent with constant stepsize $\alpha$ on a convex function

Assume further that $\|\nabla f(w;z)\|\leq L$ and $\|w_{0}-w_{\star}\|\leq D$ for some minimizer $w_{\star}$ of $R$. Let $\bar{w}_{T}$ denote the average of the $T$ iterates of the algorithm. Then we have

The upper bound stated in the previous theorem is known to be tight even if the function is $\beta$-smooth

If we plug in the population risk for $J$ in the previous theorem, we directly obtain a generalization bound for SGM that holds when we make a single pass over the data. The theorem requires fresh samples from the distribution in each update step of SGM. Hence, given $n$ data points, we cannot make more than $n$ steps, and each sample must not be used more than once.

We now contrast this bound with what follows from our results.

Let $S=(z_{1},\dots,z_{n})$ be a sample of size $n.$ Let $f$ be a $\beta$-smooth convex loss function satisfying $\|\nabla f(w,z)\|\leq L$ and let $w_{\star}^{S}$ be a minimizer of the empirical risk $R_{S}[w]=\frac{1}{n}\sum_{i=1}^{n}f(w;z_{i}).$ Suppose we run $T$ steps of SGM with suitably chosen step size from a starting point $w_{0}$ that satisfies $\|w_{0}-w_{\star}^{S}\|\leq D.$ Then, the average $\bar{w}_{T}$ over the iterates satisfies

On the one hand, applying Theorem 5.2 to the empirical risk $R_{S},$ we get

Here, $w_{\star}^{S}$ is an empirical risk minimizer. On the other hand, by our stability bound from Theorem 4.7,

Combining these two inequalities we have,

yields the bound provided in the proposition. ∎

Note that the bound from our stability analysis is not directly comparable to Corollary 5.3 as we are comparing against the expected minimum empirical risk rather than the minimum risk. Lemma 5.1 implies that the excess risk in our bound is at most worse by a factor of $\sqrt{3}$ compared with Corollary 5.3 when $T=n$. Moreover, the excess risk in our bound tends to a factor merely $\sqrt{2}$ larger than the Nemirovski-Yudin bound as $T$ goes to infinity. In contrast, the classical bound does not apply when $T>n.$

Experimental Evaluation

The goal of our experiments is to isolate the effect of training time, measured in number of steps, on the stability of SGM. We evaluated broadly a variety of neural network architectures and varying step sizes on a number of different datasets.

To measure algorithmic stability we consider two proxies. The first is the Euclidean distance between the parameters of two identical models trained on the datasets which differ by a single example. In all of our proofs, we use slow growth of this parameter distance as a way to prove stability. Note that it is not necessary for this parameter distance to grow slowly in order for our models to be algorithmically stable. This is a strictly stronger notion. Our second weaker proxy is to measure the generalization error directly in terms of the absolute different between the test error and training error of the model.

We analyzed four standard machine learning datasets each with their own corresponding deep architecture. We studied the LeNet architecture for MNIST, the cuda-convnet architecture for CIFAR-10, the AlexNet model for ImageNet, and the LSTM model for the Penn Treebank Language Model (PTB). Full details of our architectures and training procedures can be found below.

In all cases, we ran the following experiment. We choose a random example from the training set and remove it. The remaining examples constitute our set $S$. Then we create a set $S^{\prime}$ by replacing a random element of $S$ with the element we deleted. We train stochastic gradient descent with the same random seed on datasets $S$ and $S^{\prime}$. We record the Euclidean distance between the individual layers in the neural network after every $100$ SGM updates. We also record the training and testing errors once per epoch.

To varying degrees, our experiments show four primary findings:

Typically, halving the step size roughly halves the generalization error. This behavior is fairly consistent for both generalization error defined with respect to classification accuracy and cross entropy (the loss function used for training). It thus suggests that there is an intrinsic linear dependence on the step size in the generalization error. The linear relationship between generalization error and step-size is quite pronounced in the Cifar10 experiments, as shown in Figure 1.

We evaluate the Euclidean distance between the parameters of two models trained on two copies of the data differing in a random substitution. We observe that the parameter distance grows sub-linearly even in cases where our theory currently uses an exponential bound. This shows that our bounds are pessimistic.

There is a close correspondence between the parameter distance and generalization error. A priori, it could have been the case that the generalization error is small even though the parameter distance is large. Our experiments show that these two quantities often move in tandem and seem to be closely related.

When measuring parameter distance it is indeed important that SGM does not immediately encounter the random substitution, but only after some progress in training has occurred. If we artificially place the corrupted data point at the first step of SGM, the parameter distance can grow significantly faster subsequently. This effect is most pronounced in the ImageNet experiments, as displayed in Figure 7.

We evaluated convolutional neural networks for image classification on three datasets: MNIST, Cifar10 and ImageNet.

Starting with Cifar10, we chose a standard model consisting of three convolutional layers each followed by a pooling operation. This model roughly corresponds to that proposed by Krizhevsky et al. and available in the “cudaconvnet” codehttps://code.google.com/archive/p/cuda-convnet. However, to make the experiments more interpretable, we avoid all forms of regularization such as weight decay or dropout. We also do not employ data augmentation even though this would greatly improve the ultimate test accuracy of the model. Additionally, we use only constant step sizes in our experiments. With these restrictions the model we use converges to below $20\%$ test error. While this is not state of the art on Cifar10, our goal is not to optimize test accuracy but rather a simple, interpretable experimental setup.

2 Convolutional neural nets on MNIST

The situation on MNIST is largely analogous to what we saw on Cifar10. We trained a LeNet inspired model with two convolutional layers and one fully-connected layer. The first and second convolutional layers have 20 and 50 hidden units respectively. This model is much smaller and converges significantly faster than the Cifar10 models, typically achieving best test error in five epochs. We trained with minibatch size 60. As a result, the amount of overfitting is smaller as shown in Figure 5.

In the case of MNIST, we also repeated our experiments after replacing the usual cross entropy objective with a squared loss objective. The results are displayed in Figure 6. It turned out that this does not harm convergence at all, while leading to somewhat smaller generalization error and parameter divergence.

3 Convolutional neural nets on ImageNet

On ImageNet, we trained the standard AlexNet architecture using data augmentation, regularization, and dropout. Unlike in the case of Cifar10, we were unable to find a setting of hyperparameters that yielded reasonable performance without using these techniques. However, for Figure 8 , we did not use data-augmentation to exaggerate the effects of overfitting and demonstrate the impact scaling the model-size. This figure demonstrates that the model-size appears to be a second-order effect with regards to generalization error, and step-size has a considerably stronger impact.

4 Recurrent neural networks with LSTM

We also examined the stability of recurrent neural networks. Recurrent models have a considerably different connectivity pattern than their convolutional counterparts. Specifically, we looked at an LSTM architecture that was used by Zaremba et al. for language modeling . We focused on word-level prediction experiments using the Penn Tree Bank (PTB) , consisting of 929,000 training words, 73,000 validation words, and 82,000 test words. PTB has 10,000 words in its vocabularyThe data can be accessed at the URL http://www.fit.vutbr.cz/~imikolov/rnnlm/simple-examples.tgz. Following Zaremba et al., we trained regularized LSTMs with two layers that were unrolled for 20 steps. We initialize the hidden states to zero. We trained with minibatch size 20. The LSTM has 200 units per layer and its parameters are initialized to have mean zero and standard deviation of 0.1. We did not use dropout to enhance reproducibility. Dropout would only increase the stability of our models. The results are displayed in Figure 9.

Future Work and Open Problems

Our analysis parts from much previous work in that we directly analyze the generalization performance of an algorithm rather than the solution of an optimization problem. In doing so we build on the toolkit usually used to prove that algorithms converge in objective value.

This approach could be more powerful than analyzing optimality conditions, as it may be easier to understand how each data point affects a procedure rather than an optimal solution. It also has the advantage that the generalization bound holds even if the algorithm fails to find a unique optimal solution as is common in non-convex problems.

In addition to this broader perspective on algorithms for learning, there are many exciting theoretical and empirical directions that we intend to pursue in future work.

Stability of the gradient method.

In the non-convex case, we are unable to prove any reasonable form of stability at all. In fact, gradient descent is not uniformly stable as it does not enjoy the “burn-in” period of SGM as illustrated in Figure 10.

Poor generalization behavior of gradient descent has been observed in practice, but lower bounds for this approach are necessary to rule out a stable implementation for non-convex machine learning.

Acceleration and momentum.

We have described how many of the best practices in neural net training can be understood as stability inducing operations. One very important technique that we did not discuss is momentum. In momentum methods, the update is a linear combination of the current iterate and the previous direction. For convex problems, momentum is known to decrease the number of iterations required by stochastic gradient descent . For general nonlinear problems, is believed to decrease the number of iterations required to achieve low-training error . However, it is not clear that momentum adds stability. Indeed, in the case of convex optimization, momentum methods are less robust to noise than gradient methods . Thus, it is possible that momentum speeds up training but adversely impacts generalization.

Model Selection.

Another related avenue that bridges theory and practice is using stability as a method for model selection. In particular, our results imply that the models that train the fastest also generalize the best. This suggests that a heuristic for model selection would be to run many different parameter settings and choose the model which results in the lowest training error most quickly. This idea is relatively simple to try in practice, and ideas from bandit optimization can be applied to efficiently search with this heuristic cost . From the theoretical perspective, understanding the sensibility of this heuristic would require understanding lower bounds for generalizability. Are there necessary conditions which state that models which take a long training time by SGM generalize less well than those with short training times?

High capacity models that train quickly.

If the models can be trained quickly via stochastic gradient, our results prove that these models will generalize. However, this manuscript provides no guidance as to how to build a model where training is stable and training error is low. Designing a family of models which both has high capacity and can be trained quickly would be of significant theoretical and practical interest.

Indeed, the capacity of models trained in current practice steadily increases as growing computational power makes it possible to effectively train larger models. It is not uncommon for some models, such as large neural networks, to have more free parameters than the size of the sample yet have rather small generalization error . In fact, sometimes increasing the model capacity even seems to decrease the generalization error . Is it possible to understand this phenomena via stability? How can we find models which provably both have high capacity and train quickly?

Algorithm Design.

Finally, we note that stability may also provide new ideas for designing learning rules. There are a variety of successful methods in machine learning and signal processing that do not compute an exact stochastic gradient, yet are known to find quality stationary points in theory and practice . Do the ideas developed in this paper provide new insights into how to design learning rules that accelerate the convergence and improve the generalization of SGM?

Acknowledgements

The authors would like to thank Martin Abadi, Samy Bengio, Thomas Breuel, John Duchi, Vineet Gupta, Kevin Jamieson, Kenneth Marino, Giorgio Patrini, John Platt, Eric Price, Ludwig Schmidt, Nati Srebro, Ilya Sutskever, and Oriol Vinyals for their insightful feedback and helpful suggestions.

References

Appendix A Elementary properties of convex functions

Let $G=G_{f,\alpha}.$ By triangle inequality and our smoothness assumption,

Convexity and $\beta$-smoothness implies that the gradients are co-coercive, namely

First, note that if $f$ is $\gamma$ strongly convex, then $\varphi(w)=f(w)-\tfrac{\gamma}{2}\|w\|^{2}$ is convex with $(\beta-\gamma)$-smooth. Hence, applying (A.1) to $\varphi$ yields the inequality

With our assumption that $\alpha\leq\frac{2}{\beta+\gamma}$, this implies

The lemma follows by applying the inequality $\sqrt{1-x}\leq 1-x/2$ which holds for $x\in$. ∎

This proof is due to Rockafellar . Define

This is the proximal mapping associated with $f$. Define the map $Q_{\nu}(w):=w-P_{\nu}(w)$. Then, by the optimality conditions associated with (A.2), we have