Generalization Error Bounds for Noisy, Iterative Algorithms

Introduction

Many popular machine learning applications may be cast in the framework of empirical risk minimization (ERM) . This risk is defined as the expected value of an appropriate loss function, where the expectation is taken over a population. Rather than minimizing the risk directly, ERM proceeds by minimizing the empirical average of the loss function evaluated on the finite sample of data points contained in the training set . In addition to obtaining a computationally efficient, near-optimal solution to the ERM problem, it is therefore necessary to quantify how much the empirical risk deviates from the true risk of the loss function, which in turn dictates the closeness of the ERM estimate to the underlying parameter of the data-generating distribution.

In this paper, we focus on a family of iterative ERM algorithms, and derive generalization error bounds for the parameter estimates obtained from such algorithms. A unifying characteristic of the iterative algorithms considered in our paper is that each successive update includes the addition of noise, which prevents the learning algorithm from overfitting to the training data. Furthermore, the iterates of the algorithm are related via a Markov structure, and the difference between successive updates (disregarding the noise term) is assumed to be bounded. One popular learning algorithm of this nature is stochastic gradient Langevin dynamics (SGLD)—which may be viewed as a version of stochastic gradient descent (SGD) that injects Gaussian noise at each iteration—applied to a loss function with bounded gradients. Our approach leverages recent results that bound the generalization error using the mutual information between the input data set and the output parameter estimates . Importantly, this technique allows us to apply the chain rule of mutual information and leads to a simple analysis that extends to estimates that are obtained as an arbitrary function of the iterates of the algorithm. The sampling strategy may also be data-dependent and allowed to vary over time, but should be agnostic to the parameters.

Generalization properties of SGD have recently been derived using a different approach involving algorithmic stability . The main idea is that learning algorithms that change by a small bounded amount with the addition or removal of a single data point must also generalize fairly well . However, the arguments employed to show that SGD is a stable algorithm crucially rely on the fact that the updates are obtained using bounded gradient steps. Mou et al. provide generalization error bounds for SGLD by relating stability to the squared Hellinger distance, and bounding the latter quantity. Although their generalization error bounds are tighter than ours in certain cases, our approach based on a purely information-theoretic notion of stability (i.e., mutual information) allows us to consider much more general classes of updates and final outputs, including averages of iterates; furthermore, the algorithms analyzed in our framework may perform iterative updates with respect to a non-uniform sampling scheme on the training data set.

The remainder of the paper is organized as follows: In Section 2, we introduce the notation and assumptions to be used in our paper. In Section 3, we present the main result bounding the mutual information between inputs and outputs for our class of iterative learning algorithms, and derive generalization error bounds in expectation and with high probability. In Section 4, we provide illustrative examples bounding the generalization error of various noisy algorithms. We conclude with a discussion of related open problems. Detailed proofs of supporting lemmas are contained in the Appendix.

Problem setting

In the framework of ERM, we are interested in the empirical risk, defined to be the empirical average of the loss function computed with respect to the training data:

The excess risk, defined as the difference between the expected loss incurred by the algorithm and the true minimum of the risk, may be decomposed as follows:

denotes the optimization error incurred by the algorithm in minimizing the empirical risk.

2 Generalization error bounds

The idea of bounding generalization error by the mutual information $I(W;S)$ between the input and output of an ERM algorithm was first proposed by Russo and Zou and further investigated by Xu and Raginsky . We now describe their results, which will be instrumental in our work. Recall the following definition:

A random variable $X$ is $R$-sub-Gaussian if the following inequality holds:

We will assume that the loss function is uniformly sub-Gaussian in the second argument over the space $\mathcal{W}$:

Under Assumption 1, the following bound holds:

In other words, the generalization error is controlled by the mutual information, supporting the intuition that an algorithm without heavy dependence on the data will avoid overfitting.

3 Class of learning algorithms

We now define the types of ERM algorithms to be studied in our paper. We will focus on algorithms that proceed by iteratively updating a parameter estimate based on samples drawn from the data set $S$. Our theory is applicable to algorithms that make noisy, bounded updates on each step, such as the SGLD algorithm applied to a loss function with uniformly bounded gradients.

Our analysis does not actually require the noise vectors $\{\xi_{t}\}$ to be Gaussian, as long as they are drawn from a continuous distribution. The proofs would continue to hold with minimal modification, but would lead to sub-optimal bounds—indeed, a careful examination of our proofs shows that Gaussian noise produces the tightest bounds, because Gaussian noise has the maximum entropy for a fixed variance. Our results also generalize to settings where $Z_{t}$ may be a collection of data points drawn from $S$ and $F$ is computed with respect to all the data points (e.g., a mini-batched version of SGD), provided the sampling strategy satisfies the Markov structure imposed by Assumption 3 below.

For $t\geq 0$, let $W^{(t)}:=(W_{1},\ldots,W_{t})$ and $Z^{(t)}:=(Z_{1},\ldots,Z_{t})$. We impose the following assumptions on $g$, $F$, and the dependency structure between the $W$’s and $Z$’s:

The updates are bounded; i.e., $\sup_{w\in\mathcal{W},z\in\mathcal{Z}}\|F(w,z)\|_{2}\leq L$, for some $L>0$.

The sampling strategy is agnostic to the previous iterates of the parameter vectors:

where $T$ denotes the final iterate. We may represent the dependence structure defined by our class of algorithms in the form of a graphical model (see Figure 1 in the Appendix).

Importantly, we do not impose any further restrictions on the form of the updates or the sampling strategy; in particular, $Z_{t}$ need not be drawn uniformly from the data set $S$, and may even depend on past iterates $\{Z_{s}\}_{s<t}$, as in the case of sampling without replacement. Some examples of iterative algorithms where the probability of sampling a data point $z_{i}$ depends on the value of $z_{i}$ may be found in Zhao and Zhang or Needell et al. —such sampling strategies are also covered by our theory. However, note that $Z_{t}$ must be independent of the parameter iterates $\{W_{s}\}_{s<t}$, since if edges exist between $W_{t}$ and any $Z_{s}$ such that $s>t$, equation (5) will not hold. Intuitively, if the sampled data point adapts to current iterates of the parameter vector $W_{t}$, the algorithm may be prone to over-fitting and may not generalize.

The output of our estimation algorithm is defined to be an arbitrary function of the $T$ iterates: $W=f(W^{(T)})$. Some common examples appearing in the ERM literature include (i) the mean: $f(W^{(T)})=\frac{1}{T}\sum_{t=1}^{T}{W_{t}}$; (ii) the last iterate: $f(W^{(T)})=W_{T}$; or (iii) suffix averaging, and variants thereof .

Main results

We now derive an upper bound on $I(S;W)$ for the class of iterative algorithms described in Section 2, from which we obtain bounds on the generalization error.

The mutual information satisfies the bound

where the inequality follows from Lemma 2 and the last equality comes from the chain rule of mutual information.

where equality $(a)$ follows from Lemma 3 and Lemma 4 in the Appendix, whereas inequality $(b)$ follows from Lemma 5.

We may obtain bounds without a log term by using the fact that $\log(1+x)\leq\frac{x}{\sqrt{1+x}}<x$, $\forall x>0$. ∎

2 Consequences

We now use this bound on mutual information from Theorem 1 to derive bounds on the generalization error, first in expectation and then with high probability. The first bound follows directly from Theorem 1 and Lemma 1:

The generalization error of our class of iterative algorithms is bounded by

Similarly, Theorem 3 in Xu and Raginsky implies a generalization error bound that holds with high probability:

[High-probability bound] Let $I(S;W)\leq\epsilon$. Then by Theorem 1, $\epsilon$ can be equal to $\sum_{t=1}^{T}\frac{d}{2}\log\left(1+\frac{\eta_{t}^{2}L^{2}}{d\sigma_{t}^{2}}\right)$. For any $\alpha>0$ and $0<\beta\leq 1$, if $n>\frac{8R^{2}}{\alpha^{2}}\left(\frac{\epsilon}{\beta}+\log(\frac{2}{\beta})\right)$, we have

where the probability is with respect to $S\sim\mu^{\otimes n}$ and $W$.

Examples

We now apply the corollaries in Section 3.2 to obtain generalization error bounds for various algorithms.

the noise variance is set to be $\sigma_{t}=\sqrt{\eta_{t}}$,

the algorithm is run for $K$ epochs; i.e., $T=nK$,

for a constant $c>0$, the stepsizes are $\eta_{t}=\frac{c}{t}$.

For a given choice of $\{\beta,\alpha\}$, taking $n\geq\frac{64R^{4}}{\alpha^{4}}\left(\log(\frac{2}{\beta})\right)^{2}$ ensures inequality (9), provided that we run $K\leq\frac{1}{ne}\left(\frac{2}{\beta}^{\frac{2(\sqrt{n}-1)\beta}{cL^{2}}}\right)$ epochs. For more details, see Lemma 6 in Appendix B.

Bounds in expectation

Using the identity $\sum_{t=1}^{T}\frac{1}{t}\leq\log(T)+1$, we obtain the following bound:

Note that Mou et al. achieve a tighter bound on generalization error of the order $\mathcal{O}\left(\frac{1}{n}\right)$, but their bound is only applicable to the last iterate $W_{T}$ of SGLD and a uniform sampling strategy.

Convex risk minimization

where $G=\sup_{w,S}\|w_{0}-w^{*}_{S}\|_{2}$. By inequalities (1), (4.3), and (10), we then have

Setting $\sigma=\frac{G}{\sqrt{dT}}$ and $\eta=\sqrt{\frac{G^{2}}{TL(\frac{L}{2}+\frac{R\sqrt{d}T}{\sqrt{n}G})}}$, we obtain

2 Perturbed SGD

Due to the requirement that an independent noise term $\xi_{t}$ is present in each update, our results on generalization error may not be applied to SGD. On the other hand, our framework does apply to noisy versions of SGD, which have recently drawn interest in the optimization literature due to their ability to escape saddle points efficiently . For a stepsize parameter $\eta>0$, updates of the perturbed SGD algorithm take the following form :

where $\xi_{t}$ are i.i.d. noise terms sampled uniformly from the unit sphere. Hence, noise is added to each gradient. Unfortunately, our techniques cannot be applied to this exact setting because $\xi_{t}$ has a degenerate distribution concentrated on the sphere. For large enough $d$, choosing $\xi_{t}$ on the unit sphere is almost equivalent to choosing it inside the unit ball. If $\xi_{t}$ is chosen uniformly in the unit ball (cf. the perturbed SGD formulation in Jin et al. ), our methods yield the following bound:

This is because $\|W_{t}-W_{t-1}\|_{2}\leq\eta(L+1)$, so we may bound $h(W_{t}|W_{t-1})$ by the entropy of the uniform distribution on the $d$-dimensional ball of radius $\eta(L+1)$. Also, $h(W_{t}|W_{t-1},Z_{t})$ is simply the entropy of the uniform distribution on the $d$-dimensional ball of radius $\eta$. This shows that $I(W_{t};Z_{t}|W_{t-1})\leq d\log(1+L)$, so

3 Noisy momentum

In this section, we show how we can develop bounds for momentum-like algorithms in addition to SGLD. We consider an algorithm similar to the SGHMC algorithm . Every iteration $t$ involves an extra parameter vector $V_{t}$, which represents the “velocity” of $W_{t}$. We analyze a modified SGHMC algorithm, where we add the (independent and Gaussian) noise $\xi^{\prime}_{t}$ to the velocity, as well. This leads to the update equations

Note that if the gradients are upper-bounded by $L$, we have $\sup_{v,w\in\mathcal{W},z\in\mathcal{Z}}\|F((v,w),z)\|_{2}\leq\sqrt{2}L$. Using Theorem 1, we then arrive at the following bound:

Note that it is twice the bound on the mutual information appearing in Theorem 1. We may then apply the results in Section 3.2 to obtain bounds on the generalization error:

4 Accelerated gradient descent

Finally, we consider a noisy version of the accelerated gradient descent method of Nesterov , where we again add independent noise to both the velocity and parameter vectors at each iteration. This leads to the update equations

Conclusion

In this paper, we have demonstrated that mutual information is a very effective tool for bounding the generalization error of a large class of iterative ERM algorithms. The simplicity of our analysis is due to properties such as the data processing inequality and the chain rule of mutual information. However, entropy and mutual information also have certain shortcomings that limit the scope of our analysis, particularly concerning the sensitivity of entropy with respect to degenerate random variables. In some instances, mutual information-based bounds become very weak or even inapplicable. For example, if we were to analyze the SGD algorithm rather than SGLD, or add noise that is degenerate, such as the uniform distribution on a sphere , the mutual information $I(W;S)$ would be $+\infty$, leading to meaningless generalization error bounds. It would be interesting to develop information-theoretic strategies that could bound the generalization error for such algorithms, as well. Finally, note that we have only provided upper bounds for the generalization error—having a large $I(W;S)$ does not necessarily mean that an algorithm is overfitting, since our upper bound might be loose. Deriving lower bounds on the generalization error appears to be a challenging problem that could benefit from an information-theoretic approach, as well.

References

Appendix A Proofs of supporting lemmas to Theorem 1

We now prove the lemmas employed in the proof of Theorem 1.

where $Z^{(T)\backslash\{t\}}:=(Z_{1},\ldots,Z_{t-1},Z_{t+1},\ldots,Z_{T})$. See equality (5). ∎

We now bound each of the terms in the final expression. First, note that conditioned on $W_{t-1}=w_{t-1}$, we have

since translation does not affect the entropy of a random variable. Also note that the random variables $\xi_{t}$ and $\eta_{t}F(w_{t-1},Z_{t})$ are independent, so we can upper-bound the expected squared-norm of $W_{t}-w_{t-1}$, as follows:

Since the above bound holds for all values $w_{t-1}$, we may integrate the bound to conclude that

This leads to the following desired bound:

Note that if the noise were non-Gaussian, we would have to replace $\frac{d}{2}\log(2\pi e\sigma_{t}^{2})$ by the entropy of the noise. ∎

Appendix B Details for the SGLD algorithm

In this Appendix, we include more details for the derivations concerning SGLD in Section 4.

For a given choice of $\{\beta,\alpha\}$, taking $n\geq\frac{64R^{4}}{\alpha^{4}}(\log(\frac{2}{\beta}))^{2}$ ensures inequality (9), provided that we run $K\leq\frac{1}{ne}\left(\frac{2}{\beta}^{\frac{2(\sqrt{n}-1)\beta}{cL^{2}}}\right)$ epochs.

If we show that for $K\leq\frac{1}{ne}\left(\frac{2}{\beta}^{\frac{2(\sqrt{n}-1)\beta}{cL^{2}}}\right)$, we have $I(S;W)\leq(\sqrt{n}-1)\beta\log\left(\frac{2}{\beta}\right)$, the proof will follow from Corollary 2. We have

B.2 Optimization error bounds

We now derive the bound on the optimization error $\epsilon_{\text{opt}}^{W}$ of SGLD.

If we run the SGLD algorithm on an $L$-Lipschitz convex loss function for $T$ time steps with parameters $\{\eta,\sigma\}$, we have the following bound on the empirical risk for the average of the iterates:

We follow the same notation as in the rest of the paper.

The first inequality follows from the convexity of the the loss fuction.

Moreover, since the noise is independent and the loss function is convex and $L$-Lipschitz, we have

Combining this bound with inequality (14) yields the desired result. ∎