Data-Dependent Stability of Stochastic Gradient Descent

Introduction

Stochastic gradient descent (SGD) has become one of the workhorses of modern machine learning. In particular, it is the optimization method of choice for training highly complex and non-convex models, such as neural networks. When it was observed that these models generalize better (suffer less from overfitting) than classical machine learning theory suggests, a large theoretical interest emerged to explain this phenomenon. Given that SGD at best finds a local minimum of the non-convex objective function, it has been argued that all such minima might be equally good. However, at the same time, a large body of empirical work and tricks of trade, such as early stopping, suggests that in practice one might not even reach a minimum, yet nevertheless observes excellent performance.

In this work we follow an alternative route that aims to directly analyze the generalization ability of SGD by studying how sensitive it is to small perturbations in the training set. This is known as algorithmic stability approach and was used recently to establish generalization bounds for both convex and non-convex learning settings. To do so they employed a rather restrictive notion of stability that does not depend on the data, but captures only intrinsic characteristics of the learning algorithm and global properties of the objective function. Consequently, their analysis results in worst-case guarantees that in some cases tend to be too pessimistic. As recently pointed out in , deep learning might indeed be such a case, as this notion of stability is insufficient to give deeper theoretical insights, and a less restrictive one is desirable.

As our main contribution in this work we establish that a data-dependent notion of algorithmic stability, very similar to the On-Average Stability , holds for SGD when applied to convex as well as non-convex learning problems. As a consequence we obtain new generalization bounds that depend on the data-generating distribution and the initialization point of an algorithm. For convex loss functions, the bound on the generalization error is essentially multiplicative in the risk at the initialization point when noise of stochastic gradient is not too high. For the non-convex loss functions, besides the risk, it is also critically controlled by the expected second-order information about the objective function at the initialization point. We further corroborate our findings empirically and show that, indeed, the data-dependent generalization bound is tighter than the worst-case counterpart on non-convex objective functions. Finally, the nature of the data-dependent bounds allows us to state optimistic bounds that switch to the faster rate of convergence subject to the vanishing empirical risk.

In particular, our findings justify the intuition that SGD is more stable in less curved areas of the objective function and link it to the generalization ability. This also backs up numerous empirical findings in the deep learning literature that solutions with low generalization error occur in less curved regions. At the same time, in pessimistic scenarios, our bounds are no worse than those of .

Finally, we exemplify an application of our bounds, and propose a simple yet principled transfer learning scheme for the convex and non-convex case, which is guaranteed to transfer from the best source of information. In addition, this approach can also be used to select a good initialization given a number of random starting positions. This is a theoretically sound alternative to the purely random commonly used in non-convex learning.

The rest of the paper is organized as follows. We revisit the connection between stability and generalization of SGD in Section 3 and introduce a data-dependent notion of stability in Section 4. We state the main results in Section 5, in particular, Theorem 3 for the convex case, and Theorem 4 for the non-convex one. Next we demonstrate empirically that the bound shown in Theorem 4 is tighter than the worst-case one in Section 5.2.1. Finally, we suggest application of these bounds by showcasing principled transfer learning approaches in Section 5.3, and we conclude in Section 6.

Related Work

Algorithmic stability has been a topic of interest in learning theory for a long time, however, the modern approach on the relationship between stability and generalization goes back to the milestone work of . They analyzed several notions of stability, which fall into two categories: distribution-free and distribution-dependent ones. The first category is usually called uniform stability and focuses on the intrinsic stability properties of an algorithm without regard to the data-generating distribution. Uniform stability was used to analyze many algorithms, including regularized ERM (ERM) , randomized aggregation schemes , and recently SGD by , and . Despite the fact that uniform stability has been shown to be sufficient to guarantee learnability, it can be too pessimistic, resulting in worst-case rates.

In this work we are interested in the data-dependent behavior of SGD, thus the emphasis will fall on the distribution-dependent notion of stability, known as on-average stability, explored throughly in . The attractive quality of this less restrictive stability type is that the resulting bounds are controlled by how stable the algorithm is under the data-generating distribution. For instance, in and , the on-average stability is related to the variance of an estimator. In [31, Sec. 13], the authors show risk bounds that depend on the expected empirical risk of a solution to the regularized ERM. In turn, one can exploit this fact to state improved optimistic risk bounds, for instance, ones that exhibit fast-rate regimes , or even to design enhanced algorithms that minimize these bounds in a data-driven way, e.g. by exploiting side information as in transfer and metric learning . Here, we mainly focus on the later direction in the context of SGD: how stable is SGD under the data-generating distribution given an initialization point? We also touch the former direction by taking advantage of our data-driven analysis and show optimistic bounds as a corollary.

We will study the on-average stability of SGD for both convex and non-convex loss functions. In the convex setting, we will relate stability to the risk at the initialization point, while previous data-driven stability arguments usually consider minimizers of convex ERM rather than a stochastic approximation . Beside convex problems, our work also covers the generalization ability of SGD on non-convex problems. Here, we borrow techniques of and extend them to the distribution-dependent setting. That said, while bounds of are stated in terms of worst-case quantities, ours reveal new connections to the data-dependent second-order information. These new insights also partially justify empirical observations in deep learning about the link between the curvature and the generalization error . At the same time, our work is an alternative to the theoretical studies of neural network objective functions , as we focus on the direct connection between the generalization and the curvature.

In this light, our work is also related to non-convex optimization by SGD. Literature on this subject typically studies rates of convergence to the stationary points , and ways to avoid saddles . However, unlike these works, and similarly to , we are interested in the generalization ability of SGD, and thanks to the stability approach, involvement of stationary points in our analysis is not necessary.

Finally, we propose an example application of our findings in TL (TL). For instance, by controlling the stability bound in a data-driven way, one can choose an initialization that leads to improved generalization. This is related to TL where one transfers from pre-trained models , especially popular in deep learning due to its data-demanding nature . Literature on this topic is mostly focused on the ERM setting and PAC-bounds, while our analysis of SGD yields such guarantees as a corollary.

Stability of SGD

First, we introduce definitions used in the rest of the paper.

We indicate an example space by $\mathcal{Z}$ and its member by $z\in\mathcal{Z}$. For instance, in a supervised setting $\mathcal{Z}=\mathcal{X}\times\mathcal{Y}$, such that $\mathcal{X}$ is the input and $\mathcal{Y}$ is the output space of a learning problem. We assume that training and testing examples are drawn iid from a probability distribution $\mathcal{D}$ over $\mathcal{Z}$. In particular, we will denote the training set as $S=\left\{z_{i}\right\}_{i=1}^{m}\sim\mathcal{D}^{m}$.

Finally, define $R^{\star}:=\inf_{\boldsymbol{w}\in\mathcal{H}}R(\boldsymbol{w})$.

2 Uniform Stability and Generalization

On an intuitive level, a learning algorithm is said to be stable whenever a small perturbation in the training set does not affect its outcome too much. Of course, there is a number of ways to formalize the perturbation and the extent of the change in the outcome, and we will discuss some of them below. The most important consequence of a stable algorithm is that it generalizes from the training set to the unseen data sampled from the same distribution. In other words, the difference between the risk $R(A_{S})$ and the empirical risk $\widehat{R}_{S}(A_{S})$ of the algorithm’s output is controlled by the quantity that captures how stable the algorithm is. So, to observe good performance, or a decreasing true risk, we must have a stable algorithm and decreasing empirical risk (training error), which usually comes by design of the algorithm. In this work we focus on the stability of the SGD (SGD) algorithm, and thus, as a consequence, we study its generalization ability.

Recently, used a stability argument to prove generalization bounds for learning with SGD. Specifically, the authors extended the notion of the uniform stability originally proposed by , to accommodate randomized algorithms.

A randomized algorithm $A$ is $\epsilon$-uniformly stable if for all datasets $S,{S^{(i)}}\in\mathcal{Z}^{m}$ such that $S$ and ${S^{(i)}}$ differ in the $i$-th example, we have

Since SGD is a randomized algorithm, we have to cope with two sources of randomness: the data-generating process and the randomization of the algorithm $A$ itself, hence we have statements in expectation. The following theorem of shows that the uniform stability implies generalization in expectation.

Let $A$ be $\epsilon$-uniformly stable. Then,

Thus it suffices to characterize the uniform stability of an algorithm to state a generalization bound. In particular, showed generalization bounds for SGD under different assumptions on the loss function $f$. Despite that these results hold in expectation, other forms of generalization bounds, such as high-probability ones, can be derived from the above .

Apart from SGD, uniform stability has been used before to prove generalization bounds for many learning algorithms . However, these bounds typically suggest worst-case generalization rates, and rather reflect intrinsic stability properties of an algorithm. In other words, uniform stability is oblivious to the data-generating process and any other side information, which might reveal scenarios where generalization occurs at a faster rate. In turn, these insights could motivate the design of improved learning algorithms. In the following we address some limitations of analysis through uniform stability by using a less restrictive notion of stability. We extend the setting of by proving data-dependent stability bounds for convex and non-convex loss functions. In addition, we also take into account the initialization point of an algorithm as a form of supplementary information, and we dedicate special attention to its interplay with the data-generating distribution. Finally, we discuss situations where one can explicitly control the stability of SGD in a data-dependent way.

Data-dependent Stability Bounds for SGD

In this section we describe a notion of data-dependent algorithmic stability, that allows us to state generalization bounds which depend not only on the properties of the learning algorithm, but also on the additional parameters of the algorithm. We indicate such additional parameters by $\theta$, and therefore we denote stability as a function $\epsilon(\theta)$. In particular, in the following we will be interested in scenarios where $\theta$ describes the data-generating distribution and the initialization point of SGD.

A randomized algorithm $A$ is $\epsilon(\theta)$-on-average stable if it is true that

where $S\mathbin{\overset{\text{iid}}{\sim}}\mathcal{D}^{m}$ and ${S^{(i)}}$ is its copy with $i$-th example replaced by $z\mathbin{\overset{\text{iid}}{\sim}}\mathcal{D}$.

Our definition of on-average stability resembles the notion introduced by . The difference lies in the fact that we take supremum over index of replaced example. A similar notion was also used by and later by for analysis of a randomized aggregation schemes, however their definition involves absolute difference of losses. The dependence on $\theta$ also bears similarity to recent work of , however, there, it is used in the context of uniform stability. The following theorem shows that on-average - stable random algorithm is guaranteed to generalize in expectation.

Let an algorithm $A$ be $\epsilon(\theta)$-on-average stable. Then,

Main Results

Before presenting our main results in this section, we discuss algorithmic details and assumptions. We will study the following variant of SGD: given a training set $S=\{z_{i}\}_{i=1}^{m}\mathbin{\overset{\text{iid}}{\sim}}\mathcal{D}^{m}$, step sizes $\left\{\alpha_{t}\right\}_{t=1}^{T}$, random indices $I=\{j_{t}\}_{t=1}^{T}$, and an initialization point $\boldsymbol{w}_{1}$, perform updates

for $T\leq m$ steps. Moreover we will use the notation $\boldsymbol{w}_{S,t}$ to indicate the output of SGD ran on a training set $S$, at step $t$. We assume that the indices in $I$ are sampled from the uniform distribution over $[m]$ without replacement, and that this is the only source of randomness for SGD. In practice this corresponds to permuting the training set before making a pass through it, as it is commonly done in practical applications. We also assume that the variance of stochastic gradients obeys

Next, we introduce statements about the loss functions $f$ used in the following.

A loss function $f$ is $L$-Lipschitz if $\|\nabla f(\boldsymbol{w},z)\|\leq L$, $\forall\boldsymbol{w}\in\mathcal{H}$ and $\forall z\in\mathcal{Z}$. Note that this also implies that $|f(\boldsymbol{w},z)-f(\boldsymbol{v},z)|\leq L\|\boldsymbol{w}-\boldsymbol{v}\|~.$

A loss function is $\beta$-smooth if $\forall\boldsymbol{w},\boldsymbol{v}\in\mathcal{H}$ and $\forall z\in\mathcal{Z}$, $\|\nabla f(\boldsymbol{w},z)-\nabla f(\boldsymbol{v},z)\|\leq\beta\|\boldsymbol{w}-\boldsymbol{v}\|~,$ which also implies $f(\boldsymbol{w},z)-f(\boldsymbol{v},z)\leq\nabla f(\boldsymbol{v},z)^{\top}(\boldsymbol{w}-\boldsymbol{v})+\frac{\beta}{2}\|\boldsymbol{w}-\boldsymbol{v}\|^{2}~.$

A loss function $f$ has a $\rho$-Lipschitz Hessian if $\forall\boldsymbol{w},\boldsymbol{v}\in\mathcal{H}$ and $\forall z\in\mathcal{Z}$, $\|\nabla^{2}f(\boldsymbol{w},z)-\nabla^{2}f(\boldsymbol{v},z)\|_{2}\leq\rho\|\boldsymbol{w}-\boldsymbol{v}\|~.$

The last condition is occasionally used in analysis of SGD and holds whenever $f$ has a bounded third derivative. All presented theorems assume that the loss function used by SGD is non-negative, Lipschitz, and $\beta$-smooth. Examples of such commonly used loss functions are the logistic/softmax losses and neural networks with sigmoid activations. Convexity of loss functions or Lipschitzness of Hessians will only be required for some results, and we will denote it explicitly when necessary. Proofs for all the statements in this section are given in the supplementary material.

First, we present a new and data-dependent stability result for convex losses.

Assume that $f$ is convex, and that SGD’s step sizes satisfy $\alpha_{t}=\frac{c}{\sqrt{t}}\leq\frac{1}{\beta},~\forall t\in[T]$. Then SGD is $\epsilon(\mathcal{D},\boldsymbol{w}_{1})$-on-average stable with

Under the same assumptions, taking step size of order $\mathcal{O}(1/\sqrt{t})$, showed a uniform stability bound $\epsilon=\mathcal{O}(\sqrt{T/m})$. Our bound differs since it involves a multiplicative risk at the initialization point. Thus, our bound corroborates the intuition that whenever we start at a good location of the objective function, the algorithm is more stable and thus generalizes better. However, this is only the case, whenever the variance of stochastic gradient $\sigma^{2}$ is not too large. In the extreme case, deterministic case, and of $R(\boldsymbol{w}_{1})=0$, the theorem confirms that SGD, in expectation, does not need to make any updates and is therefore perfectly stable. On the other hand, when the variance $\sigma^{2}$ is large enough to make the second summand in Theorem 3 dominant, the bound does not offer improvement compared to . Note, that a result of this type cannot be obtained through the more restrictive uniform stability, precisely because such bounds on the stability must hold even for a worst-case choice of data distribution and initialization. In contrast, the notion of stability we employ depends on the data-generating distribution, which allowed us to introduce dependency on the risk.

Furthermore, consider that we start at arbitrary location $\boldsymbol{w}_{1}$: assuming that the loss function is bounded for a concrete $\mathcal{H}$ and $\mathcal{Z}$, the rate of our bound up to a constant is no worse than that of . Finally, one can always tighten this result by taking the minimum of two bounds.

2 Non-convex Losses

Now we state a new stability result for non-convex losses.

Assume that $f(\cdot,z)\in$ and has a $\rho$-Lipschitz Hessian, and that step sizes of a form $\alpha_{t}=\frac{c}{t}$ satisfy $c\leq\min\left\{\frac{1}{\beta},\frac{1}{4(2\beta\ln(T))^{2}}\right\}$. Then SGD is $\epsilon(\mathcal{D},\boldsymbol{w}_{1})$-on-average stable with

Theorem 4 immediately implies following statement that further reinforces the effect of the initialization point on the generalization error, assuming that $\E_{S}[R(A_{S})]\leq R(\boldsymbol{w}_{1})$.

Under conditions of Theorem 4 we have that SGD is $\epsilon(\mathcal{D},\boldsymbol{w}_{1})$-on-average stable with

We take a moment to discuss the role of the risk term in $\left(R(\boldsymbol{w}_{1})\cdot T\right)^{\frac{c\gamma}{1+c\gamma}}$. Observe that $\epsilon(\mathcal{D},\boldsymbol{w}_{1})\rightarrow 0$ as $R(\boldsymbol{w}_{1})\rightarrow 0$, in other words, the generalization error approaches zero as the risk of the initialization point vanishes. This is an intuitive behavior, however, uniform stability does not capture this due to its distribution-free nature. Finally, we note that [15, Theorem 3.8] showed a bound similar to (1), however, in place of $\gamma$ their bound has a Lipschitz constant of the gradient. The crucial difference lies in term $\gamma$ which is now not merely a Lipschitz constant, but rather depends on the data-generating distribution and initialization point of SGD. We compare to their bound by considering the worst case scenario, namely, that SGD is initialized in a point with high curvature, or altogether, that the objective function is highly curved everywhere. Then, at least our bound is no worse than the one of , since $\gamma\leq\beta$.

Theorem 4 also allows us to prove an optimistic generalization bound for learning with SGD on non-convex objectives.

Under conditions of Theorem 4 we have that the output of SGD obeys

An important consequence of Corollary 2, is that for a vanishing expected empirical risk, in particular for $\E_{S,A}[\widehat{R}_{S}(A_{S})]=\mathcal{O}\left(\frac{T^{c\gamma}}{m^{1+c\gamma}}\right)$, the generalization error behaves as $\mathcal{O}\left(\frac{T^{c\gamma}}{m^{1+c\gamma}}\right)$. Considering the full pass, that is $m=\mathcal{O}(T)$, we have an optimistic generalization error of order $\mathcal{O}\left(1/m\right)$ instead of $\mathcal{O}(m^{-\frac{1}{1+c\gamma}})$. We note that PAC bounds with similar optimistic message (although not directly comparable), but without curvature information can also be obtained through empirical Bernstein bounds as in . However, a PAC bound does not suggest a way to minimize non-convex empirical risk in general, where, on the other hand, SGD is known to work reasonably well.

Next we empirically assess the tightness of our non-convex generalization bounds on real data. In the following experiment we train a neural network with three convolutional layers interlaced with max-pooling, followed by the fully connected layer with $16$ units, on the MNIST dataset. This totals in a model with $18$K parameters.

Figure 1 compares our data-dependent bound (1) to the distribution-free one of [15, Theorem 3.8]. As as a reference we also include an empirical estimate of the generalization error taken as an absolute difference of the validation and training average losses. Since our bound also depends on the initialization point, we plot (1) for multiple “warm-starts”, ie.with SGD initialized from a pre-trained position. We consider $7$ such warm-starts at every $200$ steps, and report data-dependent quantities used to compute (1) just beneath the graph. Our first observation is that, clearly, the data-dependent bound gives tighter estimate, by roughly one order of magnitude. Second, simulating start from a pre-trained position suggests even tighter estimates: we suspect that this is due to decreasing validation error which is used as an empirical estimate for $R(\boldsymbol{w}_{1})$ which affects bound (1).

We compute an empirical estimate of the expected Hessian spectral norm by the power iteration method using an efficient Hessian-vector multiplication method . Since bounds depend on constants $L$, $\beta$, and $\rho$, we estimate them by tracking maximal values of the gradient and Hessian norms throughout optimization. We compute bounds with estimates $\widehat{L}=78.72$, $\widehat{\beta}=1692.28$, $\widehat{\rho}=3823.73$, and $c=10^{-3}$.

3 Application to Transfer Learning

One example application of data-dependent bounds presented before lies in TL (TL), where we are interested in achieving faster generalization on a target task by exploiting side information that originates from different but related source tasks. The literature on TL explored many ways to do so, and here we will focus on the one that is most compatible with our bounds. More formally, suppose that the target task at hand is characterized by a joint probability distribution $\mathcal{D}$, and as before we have a training set $S\mathbin{\overset{\text{iid}}{\sim}}{\mathcal{D}}^{m}$. Some TL approaches also assume access to the data sampled from the distributions associated with the source tasks. Here we follow a conservative approach – instead of the source data, we receive a set of source hypotheses $\left\{\boldsymbol{w}^{\text{src}}_{k}\right\}_{k=1}^{K}\subset\mathcal{H}$, trained on the source tasks. The goal of a learner is to come up with a target hypothesis, which in the optimistic scenario generalizes better by relying on source hypotheses. In the TL literature this is known as HTL (HTL) , that is, we transfer from the source hypotheses which act as a proxy to the source tasks and the risk $R(\boldsymbol{w}^{\text{src}}_{k})$ quantifies how much source and target tasks are related. In the following we will consider SGD for HTL, where the source hypotheses act as initialization points. First, consider learning with convex losses: Theorem 3 depends on $R(\boldsymbol{w}_{1})$, thus it immediately quantifies the relatedness of source and target tasks. So it is enough to pick the point that minimizes the stability bound to transfer from the most related source. Then, bounding $R(\boldsymbol{w}^{\text{src}}_{k})$ by $\widehat{R}_{S}(\boldsymbol{w}^{\text{src}}_{k})$ through Hoeffding bound along with union bound gives with high probability that

Hence, the most related source is the one that simply minimizes empirical risk. Similar conclusions where drawn in HTL literature, albeit in the context of ERM. Matters are slightly more complicated in the non-convex case. We take a similar approach, however, now we minimize stability bound (3), and for the sake of simplicity assume that we make a full pass over the data, so $T=m$. Minimizing the following empirical upper bound select the best source.

Let $\widehat{\gamma}_{k}^{\pm}=\Theta\Big(\frac{1}{m}\sum_{i=1}^{m}\|\nabla^{2}f(\boldsymbol{w}^{\text{src}}_{k},z_{i})\|_{2}+\sqrt{\widehat{R}_{S}(\boldsymbol{w}^{\text{src}}_{k})}\pm\sqrt{\log(K)/m}\Big)$. Then with high probability the generalization error of $\boldsymbol{w}^{\text{src}}_{k}$ is bounded by

Note that $\widehat{\gamma}_{k}^{\pm}$ involves estimation of the spectral norm of the Hessian, which is computationally cheaper to evaluate compared to the complete Hessian matrix . This is particularly relevant for deep learning, where computation of the Hessian matrix can be prohibitively expensive.

Conclusions and Future Work

In this work we proved data-dependent stability bounds for SGD and revisited its generalization ability. We presented novel bounds for convex and non-convex smooth loss functions, partially controlled by data-dependent quantities, while previous stability bounds for SGD were derived through the worst-case analysis. In particular, for non-convex learning, we demonstrated theoretically that generalization of SGD is heavily affected by the expected curvature around the initialization point. We demonstrated empirically that our bound is indeed tighter compared to the uniform one. In addition, our data-dependent analysis also allowed us to show optimistic bounds on the generalization error of SGD, which exhibit fast rates subject to the vanishing empirical risk of the algorithm’s output.

In future work we further intend to explore our theoretical findings experimentally and evaluate the feasibility of the transfer learning based on the second-order information. Another direction lies in making our bounds adaptive. So far we have presented bounds that have data-dependent components, however the step size cannot be adjusted depending on the data, e.g. as in . This was partially addressed by , albeit in the context of uniform stability, and we plan to extend this idea to the context of data-dependent stability.

References

Acknowledgments

This work was in parts funded by the European Research Council (ERC) under the European Union’s Horizon 2020 research and innovation programme (grant agreement no 637076). This work was in parts funded by the European Research Council under the European Union’s Seventh Framework Programme (FP7/2007-2013)/ERC grant agreement no 308036.

Appendix A Proofs

In this section we present proofs of all the statements.

Indicate by $S=\{z_{i}\}_{i=1}^{m}$ and $S^{\prime}=\{z^{\prime}_{i}\}_{i=1}^{m}$ independent training sets sampled i.i.d. from $\mathcal{D}$, and let ${S^{(i)}}=\left\{z_{1},\ldots,z_{i-1},z_{i}^{\prime},z_{i+1},\ldots,z_{m}\right\}$, such that $z_{i}^{\prime}\mathbin{\overset{\text{iid}}{\sim}}\mathcal{D}$. We relate expected empirical risk and expected risk by

Renaming $z^{\prime}_{i}$ as $z$ and taking $\sup$ over $i$ we get that

We say that the SGD gradient update rule is an operator $G_{t}~:~\mathcal{H}~\mapsto~\mathcal{H}$, such that

and it is also a function of the training set $S$ and a random index set $I$. Then, $\boldsymbol{w}_{t+1}=G_{t}(\boldsymbol{w}_{t})$, throughout $t=1,\ldots,T$. Recall the use of notation $\boldsymbol{w}_{S,t}$ to indicate the output of SGD ran on a training set $S$, at step $t$, and define

Next, we summarize a few instrumental facts about $G_{t}$ and few statements about the loss functions used in our proofs.

A gradient update rule is $\eta$-expansive if for all $\boldsymbol{w},\boldsymbol{v}$,

The following lemma characterizes expansiveness for the gradient update rule under different assumptions on $f$.

Assume that $f$ is $\beta$-smooth. Then, we have that:

$G_{t}$ is $(1+\alpha_{t}\beta)$-expansive,

If $f$ in addition is convex, then, for any $\alpha_{t}\leq\frac{2}{\beta}$, the gradient update rule $G_{t}$ is $1$-expansive.

An important consequence of $\beta$-smoothness of $f$ is self-boundedness , which we will use on many occasions.

For $\beta$-smooth non-negative function $f$ we have that

Self-boundedness in turn implies the following boundedness of a gradient update rule.

Assume that $f$ is $\beta$-smooth and non-negative. Then,

and also by Lipschitzness of $f$, $\|\alpha_{t}\nabla f(\boldsymbol{w},z_{j_{t}})\|\leq\alpha_{t}L$. ∎

Next we introduce a bound that relates the risk of the output at step $t$ to the risk of the initialization point $\boldsymbol{w}_{1}$ through the variance of the gradient. Given an appropriate choice of step size, this bound will be crucial at stating stability bounds that depend on the risk at $\boldsymbol{w}_{1}$. The proof idea is similar to the one of . In particular, it does not require convexity of the loss function.

Suppose SGD is ran with step sizes $\alpha_{1},\ldots,\alpha_{t-1}\leq\frac{1}{\beta}$ w.r.t. the $\beta$-smooth loss $f$. Then we have that

For brevity denote $f_{k}(\boldsymbol{w})\equiv f(\boldsymbol{w},z_{j_{k}})$. By $\beta$-smoothness of $R$ and recalling that the SGD update rule $\boldsymbol{w}_{k+1}=\boldsymbol{w}_{k}-\alpha_{k}\nabla f_{k}(\boldsymbol{w}_{k})$, we have

Taking expectation w.r.t. $S$ on both sides, recalling that $\E_{z_{k}}[\nabla f_{k}(\boldsymbol{w}_{k})]=\nabla R(\boldsymbol{w}_{k})$ and rearranging terms we get

and summing above over $k=1,\ldots,t-1$ we get the statement. ∎

Suppose SGD is ran with step sizes $\alpha_{1},\ldots,\alpha_{t-1}\leq\frac{1}{\beta}$ on the $\beta$-smooth loss $f$. Assume that the variance of stochastic gradients obeys

Now we invoke the stationary-point argument to bound the first term above as

The following lemma is similar to Lemma 3.11 of , and is instrumental in bounding the stability of SGD. However, we make an adjustment and state it in expectation over the data. Note that it does not require convexity of the loss function.

Assume that the loss function $f(\cdot,z)\in$ is $L$-Lipschitz for all $z$. Then, for every $t_{0}\in\{0,1,2,\ldots m\}$ we have that,

We proceed with elementary decomposition, Lipschitzness of $f$, and using the fact that $f$ is non-negative to have that

Taking expectation w.r.t. algorithm randomization, we get that

Now, focus on the r.h.s. above. Recall that we assume randomization by sampling from the uniform distribution over $[m]$ without replacement, and denote a realization by $\left\{j_{i}\right\}_{i=1}^{m}$. Then, we can always express our randomization as permutation function $\pi_{A}(S)=\left\{z_{j_{i}}\right\}_{i=1}^{m}$. In addition, introduce an algorithm $\text{GD}:\mathcal{Z}^{m}\mapsto\mathcal{H}$, which is identical to $A$, except that it passes over the training set $S$ sequentially without randomization. That said, we have that

Now observe that for any realization of $A$, $\E_{S,z}\left[f(\text{GD}_{\pi_{A}(S)},z)\right]=\E_{A}\E_{S,z}\left[f(A_{S},z)\right]$ because expectation w.r.t. $S$ and $z$ does not change under our randomization Strictly speaking we could omit $\E_{A}[\cdot]$ and consider any randomization by reshuffling, but we keep expectation for the sake of clarity.. Thus, we have that

Now assuming that $\tau_{A}$ is uniformly distributed over $[m]$ we have that

Putting this together with (10) and (11), we finally get that

We spend a moment to highlight the role of conditional expectation in (9). Observe that we could naively bound (8) by the Lipschitzness of $f$, but Lemma 5 follows a more careful argument. First note that $t_{0}$ is a free parameter. The expected distance in (9) between SGD outputs $\boldsymbol{w}_{S,t}$ and $\boldsymbol{w}_{{S^{(i)}},t}$ is conditioned on the fact that at step $t_{0}$ outputs of SGD are still the same. This means that the perturbed point is encountered after $t_{0}$. Then, the conditional expectation should be a decreasing function of $t_{0}$: the later the perturbation occurs, the smaller deviation between $\boldsymbol{w}_{S,t}$ and $\boldsymbol{w}_{{S^{(i)}},t}$ we should expect. Later we use this fact to minimize the bound (9) over $t_{0}$.

A.2 Convex Losses

In this section we prove on-average stability for loss functions that are non-negative, $\beta$-smooth, and convex.

Assume that $f$ is convex, and that SGD’s is ran with step sizes $\left\{\alpha_{t}\right\}_{t=1}^{T}$. Then, for every $t_{0}\in\{0,1,2,\ldots m\}$, SGD is $\epsilon(\mathcal{D},\boldsymbol{w}_{1})$-on-average stable with

For brevity denote $\Delta_{t}(S,z):=\E_{A}\left[\delta_{t}(S,z)\ |\ \delta_{t_{0}}(S,z)=0\right]$. We start by applying Lemma 5:

Our goal is to bound the first term on the r.h.s. as a decreasing function of $t_{0}$, so that eventually we can minimize the bound w.r.t. $t_{0}$. At this point we focus on the first term, and the proof partially follows the outline of the proof of Theorem 3.7 in . The strategy will be to establish the bound on $\Delta_{T}(S,z)$ by using a recursive argument. In fact we will state the bound on $\Delta_{t+1}(S,z)$ in terms of $\Delta_{t}(S,z)$ and then unravel the recursion. Finally, we will take expectation w.r.t. the data after we obtain the bound by recursion.

To do so, we distinguish two cases: 1) SGD encounters a perturbed point at step $t$, that is $t=i$, and 2) the current point is the same in $S$ and ${S^{(i)}}$, so $t\neq i$. For the first case, we will use data-dependent boundedness of the gradient update rule, Corollary 3, that is

To handle the second case, we will use the expansiveness of the gradient update rule, Lemma 1, which states that for convex loss functions, the gradient update rule is $1$-expansive, so $\delta_{t+1}(S,z)\leq\delta_{t}(S,z)$. Considering both cases of example selection, and noting that SGD encounters the perturbation w.p. $\frac{1}{m}$, we write $\E_{A}$ for a step $t$ as

Unraveling the recursion from $T$ to $t_{0}$ and plugging the above into (14) yields

Next statement is a simple consequence of Theorem 5 and Lemma 4.

Bounding the sum using Lemma 4 recalling that $\alpha_{t}=c/\sqrt{t}$, we get

Combining above with (15) completes the proof. ∎

A.3 Non-convex Losses

Our proof of a stability bound for non-convex loss functions, Theorem 4 (in the submission file), follows a general outline of [15, Theorem 3.8]. Namely, the outputs of SGD run on a training set $S$ and its perturbed version ${S^{(i)}}$ will not differ too much, because by the time a perturbation is encountered, the step size has already decayed enough. So, on the one hand, stabilization is enforced by the diminishing the step size, and on the other hand, by how much updates expand the distance between the gradients after the perturbation. Since work with uniform stability, they capture the expansiveness of post-perturbation update by the Lipschitzness of the gradient. In combination with a recursive argument, their bound has exponential dependency on the Lipschitz constant of the gradient. We argue that the Lipschitz continuity of the gradient can be too pessimistic in general. Instead, we rely on a local data-driven argument: considering that we initialize SGD at point $\boldsymbol{w}_{1}$, how much do updates expand the gradient under the distribution of interest? The following crucial lemma characterizes such behavior in terms of the curvature at $\boldsymbol{w}_{1}$.

Assume that the loss function $f(\cdot,z)$ is $\beta$-smooth and that its Hessian is $\rho$-Lipschitz. Then,

Recall that the randomness of the algorithm is realized through sampling without replacement from the uniform distribution over $[m]$. Apart from that we will not be concerned with the randomness of the algorithm, and given the set of random variables $\{j_{i}\}_{i=1}^{m}$, for brevity we will use indexing notation $z_{1},z_{2},\ldots,z_{m}$ to indicate $z_{j_{1}},z_{j_{2}},\ldots,z_{j_{m}}$. Next, let ${S^{(i)}}=\left\{z_{i}^{\prime}\right\}_{i=1}^{m}$, and introduce a shorthand notation $f_{k}(\boldsymbol{w})=f(\boldsymbol{w},z_{k})$ and $f_{k^{\prime}}(\boldsymbol{w})=f(\boldsymbol{w},z^{\prime}_{k})$. We start by applying triangle inequality to get

In the following we will focus on the second term of r.h.s. above. Given SGD outputs $\boldsymbol{w}_{S,t}$ and $\boldsymbol{w}_{{S^{(i)}},t}$ with $t>i$, our goal here is to establish how much do gradients grow apart with every new update. This behavior can be characterized assuming that gradient is Lipschitz continuous, however, we conduct a local analysis. Specifically, we observe how much do updates expand gradients, given that we start at some point $\boldsymbol{w}_{1}$ under the data-generating distribution. So, instead of the Lipschitz constant, expansiveness rather depends on the curvature around $\boldsymbol{w}_{1}$. On the other hand, we are dealing with outputs at an arbitrary time step $t$, and therefore we first have to relate them to the initialization point $\boldsymbol{w}_{1}$. We do so by using the gradient update rule and telescopic sums, and conclude that this relationship is controlled by the sum of gradient norms along the update path. We further establish that this sum is controlled by the risk of $\boldsymbol{w}_{1}$ up to the noise of stochastic gradients, through stationary-point result of Lemma 4. Thus, the proof consists of two parts: 1) Decomposition into curvature and gradients along the update path, and 2) bounding those gradients.

Introduce $\boldsymbol{\delta}_{t}:=\boldsymbol{w}_{{S^{(i)}},t}-\boldsymbol{w}_{S,t}$. By Taylor theorem we get that

Taking norm on both sides, applying triangle inequality, Cauchy-Schwartz inequality, and assuming that Hessians are $\rho$-Lipschitz we obtain

Using telescoping sums and SGD update rule we get that

Plugging above into the integral of (17) we have

Plugging this result back into (17) completes the proof of the first statement. The second statement comes from Lemma 4 with $\alpha_{t}=c/t$. ∎ Next, we need the following statement to prove our stability bound.

Let $Z$ be a zero-mean real-valued r.v., such that $|Z|\leq b$ and $\E[Z^{2}]\leq\sigma^{2}$. Then for all $|c|\leq\frac{1}{2b}$, we have that $\E\left[e^{cZ}\right]\leq e^{c^{2}\sigma^{2}}~.$

Stated inequality is a consequence of a Bernstein-type inequality for moment generating functions, Theorem 2.10 in . Observe that zero-centered r.v. $Z$ bounded by $b$ satisfies Bernstein’s condition, that is

This in turn satisfies condition for Bernstein-type inequality stating that

Choosing $|c|\leq\frac{1}{2b}$ verifies the statement. ∎

Now we are ready to prove Theorem 4, which bounds the $\epsilon(\mathcal{D},\boldsymbol{w}_{1})$-on-average stability of SGD.

Most of the proof is dedicated to bounding the first term in (18). We deal with this similarly as in . Specifically, we state the bound on $\Delta_{T}(S,z)$ by using a recursion. In our case, however, we also have an expectation w.r.t. the data, and to avoid complications with dependencies, we first unroll the recursion for the random quantities, and only then take the expectation. At this point the proof crucially relies on the product of exponentials arising from the recursion, and all relevant random quantities end up inside of them. We alleviate this by Proposition 2. Finally, we conclude by minimizing (18) w.r.t. $t_{0}$. Thus we have three steps: 1) recursion, 2) bounding $\E[\exp(\cdots)]$, and 3) tuning of $t_{0}$.

We begin by stating the bound on $\Delta_{T}(S,z)$ by recursion. Thus we will first state the bound on $\Delta_{t+1}(S,z)$ in terms of $\Delta_{t}(S,z)$, and other relevant quantities and then unravel the recursion. As in the convex case, we distinguish two cases: 1) SGD encounters the perturbed point at step $t$, that is $t=i$, and 2) the current point is the same in $S$ and ${S^{(i)}}$, so $t\neq i$. For the first case, we will use worst-case boundedness of $G_{t}$, Corollary 3, that is, $\|G_{t}(\boldsymbol{w}_{S,t})-G_{t}(\boldsymbol{w}_{{S^{(i)}},t})\|\leq\delta_{t}(S,z)+2\alpha_{t}L~.$ To handle the second case we will use Lemma 6, namely,

In addition, as a safety measure we will also take into account that the gradient update rule is at most $(1+\alpha_{t}\beta)$-expansive by Lemma 1. So we will work with the function $\psi_{t}(S,z):=\min\left\{\xi_{t}(S,z),\beta\right\}$ instead of $\xi_{t}(S,z)$. and decompose the expectation w.r.t. $A$ for a step $t$. Noting that SGD encounters the perturbed example with probability $\frac{1}{m}$,

where the last inequality follows from $1+x\leq\exp(x)$. This inequality is not overly loose for $x\in$, and, in our case it becomes instrumental in handling the recursion.

Now, observe that relation $x_{t+1}\leq a_{t}x_{t}+b_{t}$ with $x_{t_{0}}=0$ unwinds from $T$ to $t_{0}$ as $x_{T}\leq\sum_{t=t_{0}+1}^{T}b_{t}\prod_{k=t+1}^{T}a_{k}$. Consequently, having $\Delta_{t_{0}}(S,z)=0$, we unwind (19) to get

We take expectation w.r.t. $S$ and $z$ on both sides and focus on the expectation of the exponential in (20). First, introduce $\mu_{k}:=\E_{S,z}[\psi_{k}(S,z)]$, and proceed as

Observe that zero-mean version of $\psi_{k}(S,z)$ is bounded as

and assume the setting of $c$ as $c\leq\frac{1}{2(2\beta\ln(T))^{2}}$. By Proposition 2, we have

Next, we give an upper-bound on $\mu_{k}$, that is $\mu_{k}\leq\min\left\{\beta,\E_{S,z}[\xi_{k}(S,z)]\right\}$. Finally, we bound $\E_{S,z}[\xi_{k}(S,z)]$ using the second result of Lemma 6, which holds for any $k\in[T]$, to get that $\mu_{k}\leq\gamma$, with $\gamma$ defined in the statement of the theorem.

Now we turn our attention back to (20). Considering that we took an expectation w.r.t. the data, we use (22) and the fact that $\mu_{k}\leq\gamma$ to get that

minimizes (23). Plugging $t_{0}$ back we get that (23) equals to

A.3.1 Optimistic Rates for Learning with Non-convex Loss Functions

Next we will prove an optimistic bound based on Theorem 4, in other words, the bound that demonstrates fast convergence rate subject to the vanishing empirical risk. First we will need the following technical statement.

[7, Lemma 7.2] Let $c_{1},c_{2},\ldots,c_{l}>0$ and $s>q_{1}>q_{2}>\ldots>q_{l-1}>0$. Then the equation

has a unique positive solution $x^{\star}$. In addition,

Next we prove a useful technical lemma similarly as in [25, Lemma 7].

Let $a,c>0$ and $0<\alpha<1$. Then the inequality

Consider a function $h(x)=x-ax^{\alpha}-c$. Applying Lemma 7 with $s=1$, $l=2$, $c_{1}=a$, $c_{2}=c$, and $q_{1}=\alpha$ we get that $h(x)=0$ has a unique positive solution $x^{\star}$ and

Moreover, the inequality $h(x)\leq 0$ is verified for $x=0$, and $\lim_{x\rightarrow+\infty}h(x)=+\infty$, so we have that $h(x)\leq 0$ implies $x\leq x^{\star}$. Now, using this fact and the fact that $h(x^{\star})=0$, we have that

and upper-bounding $x^{\star}$ by (24) we finally have

Consider Theorem 4 and observe that it verifies condition of Lemma 8 with $x=\E_{S,A}\left[R(A_{S})\right]$, $c=\E_{S,A}\left[\widehat{R}_{S}(A_{S})\right]$, $\alpha=\frac{c\gamma}{1+c\gamma}$, and

Note that $\alpha/(1-\alpha)=c\gamma$ and $1/(1-\alpha)=1+c\gamma$. Then, we obtain that

Consider minimizing the bound given by Corollary 1 (in the submission file) over a discrete set of source hypotheses $\left\{\boldsymbol{w}^{\text{src}}_{k}\right\}_{k=1}^{K}$,

By Hoeffding inequality, with high probability, we have that $|\gamma_{k}-\widehat{\gamma}_{k}|\leq\mathcal{O}\left(\frac{1}{\sqrt{m}}\right)$. Now we further upper bound (25) by upper bounding $R(\boldsymbol{w}^{\text{src}}_{k})$ and apply union bound to get

where $\widehat{\gamma}_{k}^{\pm}=\widehat{\gamma}_{k}\pm\frac{1}{\sqrt{m}}$. This completes the proof. ∎