Tightening Mutual Information Based Bounds on Generalization Error

TL;DR

This paper introduces a new, tighter information-theoretic bound on the generalization error of supervised learning algorithms, applicable under broad conditions and easily estimable in practice.

Contribution

It presents a novel mutual information-based bound that improves upon existing bounds in tightness and applicability, especially for noisy and iterative algorithms.

Findings

The new bound is tighter than previous bounds.

It applies to a broad class of loss functions.

The bound can be practically estimated from data.

Abstract

An information-theoretic upper bound on the generalization error of supervised learning algorithms is derived. The bound is constructed in terms of the mutual information between each individual training sample and the output of the learning algorithm. The bound is derived under more general conditions on the loss function than in existing studies; nevertheless, it provides a tighter characterization of the generalization error. Examples of learning algorithms are provided to demonstrate the the tightness of the bound, and to show that it has a broad range of applicability. Application to noisy and iterative algorithms, e.g., stochastic gradient Langevin dynamics (SGLD), is also studied, where the constructed bound provides a tighter characterization of the generalization error than existing results. Finally, it is demonstrated that, unlike existing bounds, which are difficult to compute and evaluate empirically, the proposed bound can be estimated easily in practice.

Figures

Equations

$$L_{\mu }(w)\triangleq {\mathbb{E}}_{Z\sim \mu }[\ell (w,Z)].$$

$$L_S(w)\triangleq \frac{1}{n}\sum _{i=1}^n\ell (w,Z_i).$$

$$\mathrm{gen}(\mu ,P_{W|S})\triangleq {\mathbb{E}}_{W,S}[L_{\mu }(W)-L_S(W)],$$

$$|\mathrm{gen}(\mu ,P_{W|S})|\le \sqrt{\frac{2R^2}{n}I(S;W)}.$$

$${\Lambda }_X(\lambda )\triangleq \log \mathbb{E}[e^{\lambda (X-\mathbb{E}X)}].$$

$$\psi^{*}(x)\triangleq\sup_{\lambda\in[0,b)}\big{(}\lambda x-\psi(\lambda)\big{)}.$$

$$\psi^{*-1}(y)=\inf_{\lambda\in(0,b)}\Big{(}\frac{y+\psi(\lambda)}{\lambda}\Big{)}.$$

$$\mathbb{E}[f(W,Z)]-\mathbb{E}[f(\tilde{W},\tilde{Z})]$$

$$\mathbb{E}[f(\tilde{W},\tilde{Z})]-\mathbb{E}[f(W,Z)]$$

$$D(P\|Q)=\sup_{g\in\mathcal{G}}\Big{\{}\mathbb{E}_{P}[g(X)]-\log\mathbb{E}_{Q}[e^{g(X)}]\Big{\}},$$

$$I(W;Z)$$

$$\ge \mathbb{E}[\lambda f(W,Z)]-\log \mathbb{E}[e^{\lambda f(\tilde{W},\tilde{Z})}]$$

$$\ge \lambda (\mathbb{E}[f(W,Z)]-\mathbb{E}[f(\tilde{W},\tilde{Z})])-{\psi }_{+}(\lambda ),$$

$${\Lambda }_f(\tilde{W};\tilde{Z})=\log \mathbb{E}[e^{\lambda (f(\tilde{W},\tilde{Z})-\mathbb{E}f(\tilde{W},\tilde{Z}))}]\le {\psi }_{+}(\lambda ),$$

$$D(P_{W,Z}\Vert$$

$$\ge \lambda (\mathbb{E}[f(W,Z)]-\mathbb{E}[f(\tilde{W},\tilde{Z})])-{\psi }_{-}(-\lambda ).$$

$$\mathbb{E}[f(W,Z)]-\mathbb{E}[f(\tilde{W},\tilde{Z})]$$

$$\displaystyle=\psi^{*-1}_{+}\big{(}I(W;Z)\big{)},$$

$$\mathbb{E}[f(\tilde{W},\tilde{Z})]-\mathbb{E}[f(W,Z)]$$

$$\displaystyle=\psi^{*-1}_{-}\big{(}I(W;Z)\big{)},$$

$$\displaystyle\mathrm{gen}(\mu,P_{W|S})\leq\frac{1}{n}\sum_{i=1}^{n}\psi^{*-1}_{-}\big{(}I(W;Z_{i})\big{)},$$

$$\displaystyle-\mathrm{gen}(\mu,P_{W|S})\leq\frac{1}{n}\sum_{i=1}^{n}\psi^{*-1}_{+}\big{(}I(W;Z_{i})\big{)}.$$

$$\mathrm{gen}(\mu ,P_{W|S})$$

$$\big{|}\mathrm{gen}(\mu,P_{W|S})\big{|}\leq\frac{1}{n}\sum_{i=1}^{n}\sqrt{2R^{2}I(W;Z_{i})}.$$

$$\frac{1}{n}\sum_{i=1}^{n}\psi^{*-1}\bigg{(}I(W;Z_{i})\bigg{)}\leq\psi^{*-1}\bigg{(}\frac{I(S;W)}{n}\bigg{)}.$$

$$I(W;S)=\sum _{i=1}^{n}I(W;Z_i|Z^{i-1})$$

$$I(W;Z_i|Z^{i-1})$$

$$=I(W,Z^{i-1};Z_i)$$

$$=I(W;Z_i)+I(Z^{i-1};Z_i|W)$$

$$\ge I(W;Z_i).$$

$$I(W;S)$$

$$L_S(w)\triangleq \frac{1}{n}\sum _{i=1}^n\Vert w-Z_i{\Vert }_2^2,w\in {\mathbb{R}}^d.$$

$$\mathrm{gen}(\mu ,P_{W|S})$$

$$I(W;Z_i)$$

Full text

I Introduction

Recent success of deep learning algorithms [2] has dramatically boosted their applications in various engineering and science domains, e.g., computer vision [3], natural language processing [4], autonomous driving [5], and health care [6]. A deep neural network trained using a sufficiently large amount of training data can achieve a small training error, while simultaneously performing well on unseen data, i.e., it generalizes well. However, we have yet to develop a satisfactory understanding of why deep learning algorithms generalize well.

Classical statistical learning approaches for analyzing the generalization capability of supervised learning algorithms can be mainly categorized into two groups. The first set of methods are based on measures of the complexity of the output hypothesis space, e.g., VC dimension and Rademacher complexity [7, 8]. However, these complexity measures usually scale exponentially with the depth of deep neural networks [9]. Moreover, these approaches do not take into consideration the regularization implicitly imposed by the algorithms used to train the neural networks, e.g., stochastic gradient descent [10, 11]. Thus, the generalization error bounds based on these complexity measures tend to be loose and do not explain why deep neural networks generalize well in practice. The second set of methods are based on exploiting properties of the learning algorithm, e.g., PAC-Bayesian bounds [12], uniform stability [13, 14], and compression bounds [15]. However, as discussed in [11, 16], these approaches do not exploit the fact that the generalization error depends strongly on the underlying true data-generating distribution. For example, if the labels are irrelevant to the input features, then the generalization error will be large for a deep neural network, since training error is usually small due to the large capacity of the network, but test error will be large due to the fact that there is no relationship between the input features and the label [11].

Recently, it was proposed in [17] and further studied in [18] and [19] that the metric of mutual information can be used to develop upper bounds on the generalization error of learning algorithms. Such an information-theoretic framework can handle a broader range of problems, and it could also address the aforementioned challenges of implicit regularization and dependence on data generating distribution. More importantly, it offers an information-theoretic point of view on how to improve the generalization capability of a learning algorithm, and this new perspective provides us with a better understanding of the generalization behavior of deep neural networks.

In this paper, we follow the information-theoretic framework proposed in [17, 18, 19]. Our main contribution is a tighter upper bound on the generalization error using the mutual information between an individual training sample and the output hypothesis of the learning algorithm. We show that compared to existing studies, our bound has a broader range of applicability, and can be considerably tighter.

We consider an instance space $\mathcal{Z}$, a hypothesis space $\mathcal{W}$, and a nonnegative loss function $\ell:\mathcal{W}\times\mathcal{Z}\to\mathbb{R}^{+}$. A training dataset $S=\{Z_{1},\cdots,Z_{n}\}$ that consists of $n$ i.i.d samples $Z_{i}\in\mathcal{Z}$ drawn from an unknown distribution $\mu$ is available. The goal of a supervised learning algorithm is to find an output hypothesis $w\in\mathcal{W}$ that minimizes the population risk:

[TABLE]

In practice, $\mu$ is unknown, and thus $L_{\mu}(w)$ cannot be computed directly. Instead, the empirical risk of $w$ on a training dataset $S$ is studied, which is defined as

[TABLE]

A learning algorithm can be characterized by a randomized mapping from the training dataset $S$ to a hypothesis $W$ according to a conditional distribution $P_{W|S}$.

In statistical learning theory, the (mean) generalization error111 The term “generalization error” of a learning algorithm is usually defined as the difference between the population risk and the training error without taking the expectation with respect to the randomness of the data and the learning algorithm. Here, we consider the expectation of the generalization error over the randomness of both the data and the learning algorithm. We will use the term “generalization error” throughout the paper, with the understanding that it is the “mean generalization error”. of a supervised learning algorithm is the expected difference between the population risk of the output hypothesis and its empirical risk on the training dataset:

[TABLE]

where the expectation is taken over the joint distribution $P_{S,W}=P_{S}\otimes P_{W|S}$. Note that $P_{W|S}$ will become degenerate if $W$ is a deterministic function of $S$. The generalization error is used to measure the extent to which the learning algorithm overfits the training data.

Main Contributions and Related Works

We first review the following lemma from [18], which provides an upper bound on the generalization error using the mutual information $I(S;W)$ between the training dataset $S$ and the output hypothesis $W$.

Lemma 1** ([18, Theorem 1]).**

Suppose $\ell(w,Z)$ is $R$-sub-Gaussian222A random variable $X$ is $R$-sub-Gaussian if $\log\mathbb{E}[e^{\lambda(X-\mathbb{E}X)}]\leq\frac{R^{2}\lambda^{2}}{2}$, $\forall\lambda\in\mathbb{R}$. under $Z\sim\mu$ for all $w\in\mathcal{W}$, then

[TABLE]

This mutual information based bound in (4) is related to the “on-average” stability (see, e.g., [20]), since it quantifies the overall dependence between the output of the learning algorithm and all the input training samples via $I(S;W)$. Note that $I(S;W)$ depends on the main components of a supervised learning problem, i.e., the hypothesis space $\mathcal{W}$, the learning algorithm $P_{W|S}$, and the data generating distribution $\mu$, in contrast to the traditional bounds based on VC dimension or the uniform stability, which only depend on one aspect of the learning problem. We also note that there is a connection between the mutual information based generalization bound and the PAC-Bayesian bound in [21], since both methods adopt the variational representation of relative entropy to establish the decoupling lemma. By further exploiting the structure of the hypothesis space and the dependency between the algorithm input and output, the authors of [19, 22] combined the chaining and mutual information methods, and obtained a tighter bound on the generalization error.

However, the bound in Lemma 1 and the chaining mutual information (CMI) bound in [19] both suffer from the following two shortcomings. First, for empirical risk minimization (ERM), if $W$ is the unique minimizer of $L_{S}(w)$ in $\mathcal{W}$, then $W$ is a deterministic function of $S$ and the mutual information $I(S;W)=\infty$. It can be shown that both bounds are not tight in this case. Second, both bounds assume that $\ell(w,Z)$ has a bounded cumulant generating function (CGF) under $Z\sim\mu$ for all $w\in\mathcal{W}$, which may not hold in many cases (see Section IV).

There has been some recent work on addressing these shortcomings of mutual information based bounds on generalization error by using other information-theoretical measures, e.g., Wasserstein distance [23, 24, 25], maximal leakage [26, 27] and total variation [28] to bound the generalization error. But the measures proposed in these papers are difficult to evaluate both analytically and empirically as we discuss in Section VI, which significantly undermines the usefulness of these results in practice.

In this paper, we get around the aforementioned shortcomings by combining the idea of point-wise stability [14, 23] with the information-theoretic framework introduced in [18]. Specifically, an algorithm is said to be point-wise stable if the expectation of the loss function $\ell(W,Z_{i})$ does not change too much with the replacement of any individual training sample $Z_{i}$, and if an algorithm is point-wise stable, then it generalizes well [14, 23]. Motivated by these facts, we tighten the mutual information based generalization error bound through a bound based on the individual sample mutual information (ISMI) $I(W;Z_{i})$. Compared with the bound in Lemma 1, and the CMI bound in [19], the ISMI bound is derived under a more general condition on the CGF of the loss function, is applicable to a broader range of problems, and can provide a tighter characterization of the generalization error.

The rest of the paper is organized as follows. In Section II, we provide some preliminary definitions and results for our analysis. In Section III, we introduce the individual sample mutual information generalization bound. In Section IV, we apply our method to bound the generalization errors of two learning problems with infinite $I(S;W)$. We show in the second example that our ISMI bound can be tighter than the CMI bound in [19], while the bound in Lemma 1 is infinity. In Section V, we improve the generalization error bound in [29] for SGLD algorithm using our method, which demonstrates that the ISMI bound is applicable to the noisy, iterative algorithms discussed in [29]. In Section VI, we provide an example where the ISMI bound can be evaluated empirically from the samples, while other existing bounds are difficult to estimate due to prohibitive computational complexity.

II Preliminaries

We use upper letters to denote random variables, and calligraphic upper letters to denote sets. For a random variable $X$ generated from a distribution $\mu$, we use $\mathbb{E}_{X\sim\mu}$ to denote the expectation taken over $X$ with distribution $\mu$. We write $I_{d}$ to denote the $d$-dimensional identity matrix. All the logarithms are the natural ones, and all the information measure units are nats. We use $\mu^{\otimes n}$ to denote the product distribution of $n$ copies of $\mu$.

Definition 1**.**

The cumulant generating function (CGF) of a random variable $X$ is defined as

[TABLE]

Assuming $\Lambda_{X}(\lambda)$ exists, it can be verified that $\Lambda_{X}(0)=\Lambda_{X}^{\prime}(0)=0$, and that it is convex.

Definition 2**.**

For a convex function $\psi$ defined on the interval $[0,b)$, where $0<b\leq\infty$, its Legendre dual $\psi^{*}$ is defined as

[TABLE]

The following lemma characterizes a useful property of the Legendre dual and its inverse function.

Lemma 2** ([30, Lemma 2.4]).**

Assume that $\psi(0)=\psi^{\prime}(0)=0$. Then $\psi^{*}(x)$ defined above is a non-negative convex and non-decreasing function on $[0,\infty)$ with $\psi^{*}(0)=0$. Moreover, its inverse function $\psi^{*-1}(y)=\inf\{x\geq 0:\psi^{*}(x)\geq y\}$ is concave, and can be written as

[TABLE]

For an $R$-sub-Gaussian random variable $X$, $\psi(\lambda)=\frac{R^{2}\lambda^{2}}{2}$ is an upper bound on $\Lambda_{X}(\lambda)$. Then by Lemma 2, $\psi^{*-1}(y)=\sqrt{2R^{2}y}$.

III Bounding Generalization Error via $I(W;Z_{i})$

In this section, we first generalize the decoupling lemma in [18, Lemma 1] to a different setting, and then tighten the bound on generalization error via the individual sample mutual information $I(W;Z_{i})$.

III-A General Decoupling Estimate

Consider a pair of random variables $W$ and $Z$ with joint distribution $P_{W,Z}$. Let $\widetilde{W}$ be an independent copy of $W$, and $\widetilde{Z}$ be an independent copy of $Z$, such that $P_{\widetilde{W}\widetilde{Z}}=P_{{W}}\otimes P_{{Z}}$. Suppose $f:\mathcal{W}\times\mathcal{Z}\to\mathbb{R}$ is a real-valued function. If the CGF $\Lambda_{f(\widetilde{W},\widetilde{Z})}(\lambda)$ of $f(\widetilde{W},\widetilde{Z})$ can be upper bounded by some function $\psi$ for $\lambda\in(b_{-},b_{+})$, we have the following theorem.

Theorem 1**.**

Assume that $\Lambda_{f(\widetilde{W},\widetilde{Z})}(\lambda)\leq\psi_{+}(\lambda)$ for $\lambda\in[0,b_{+})$, and $\Lambda_{f(\widetilde{W},\widetilde{Z})}(\lambda)\leq\psi_{-}(-\lambda)$ for $\lambda\in(b_{-},0]$ under distribution $P_{\widetilde{W}\widetilde{Z}}=P_{{W}}\otimes P_{{Z}}$, where $0<b_{+}\leq\infty$ and $-\infty\leq b_{-}<0$. Suppose that $\psi_{+}(\lambda)$ and $\psi_{-}(\lambda)$ are convex, and $\psi_{+}(0)=\psi_{+}^{\prime}(0)=\psi_{-}(0)=\psi_{-}^{\prime}(0)=0$. Then,

[TABLE]

Proof.

Consider the variational representation of the relative entropy between two probability measures $P$ and $Q$ defined on $\mathcal{X}$:

[TABLE]

where the supremum is over all measurable functions $\mathcal{G}=\{g:\mathcal{X}\to\mathbb{R},\ \mathrm{s.t.}\ \mathbb{E}_{Q}[e^{g(X)}]<\infty\}$, and equality is achieved when $g=\log\frac{dP}{dQ}$, where $\frac{dP}{dQ}$ is the Radon–Nikodym derivative. It then follows that $\forall\lambda\in[0,b_{+})$,

[TABLE]

where the last inequality follows from the assumption that

[TABLE]

for $\lambda\in[0,b_{+})$. Similarly, for $\lambda\in(b_{-},0]$, it follows that

[TABLE]

From (III-A) it follows that

[TABLE]

and from (III-A) it follows that

[TABLE]

where the equalities in (III-A) and (III-A) follow from Lemma 2. ∎

Theorem 1 provides a different characterization of the decoupling estimate than existing results. Specifically, it is assumed that the CGF of $f(w,Z)$ is bounded for all $w\in\mathcal{W}$ and $Z\sim\mu$ in [18, Lemma 1] and [31, Theorem 2], whereas in Theorem 1, it is assumed that the CGF of $f(\widetilde{W},\widetilde{Z})$ is bounded in expectation under $P_{\widetilde{W}\widetilde{Z}}=P_{{W}}\otimes P_{{Z}}$.

III-B Individual Sample Mutual Information Bound

Motivated by the idea of algorithmic stability, which measures how much an output hypothesis changes with the replacement of an individual training sample, we construct the following upper bound on the generalization error via $I(W;Z_{i})$.

Theorem 2**.**

Suppose $\ell(\widetilde{W},\widetilde{Z})$ satisfies $\Lambda_{\ell(\widetilde{W},\widetilde{Z})}(\lambda)\leq\psi_{+}(\lambda)$ for $\lambda\in[0,b_{+})$, and $\Lambda_{\ell(\widetilde{W},\widetilde{Z})}(\lambda)\leq\psi_{-}(-\lambda)$ for $\lambda\in(b_{-},0]$ under $P_{\widetilde{Z},\widetilde{W}}=\mu\otimes P_{W}$, where $0<b_{+}\leq\infty$ and $-\infty\leq b_{-}<0$. Then,

[TABLE]

Proof.

The generalization error can be written as follows:

[TABLE]

where $W$ and $Z_{i}$ in the second term are dependent with $P_{W,Z_{i}}=\mu\otimes P_{W|Z_{i}}$, and $W$ and $\widetilde{Z}$ in the first term are independent with the same marginal distributions. Applying Theorem 1 completes the proof. ∎

In the following proposition, we derive the ISMI bounds under two different sub-Gaussian assumptions.

Proposition 1**.**

Suppose that $\ell(w,Z)$ is $R$-sub-Gaussian under $Z\sim\mu$ for all $w\in\mathcal{W}$, then

[TABLE] 2. 2.

Suppose that $\ell(\widetilde{W},\widetilde{Z})$ is $R$-sub-Gaussian under distribution $P_{\widetilde{W}\widetilde{Z}}=P_{{W}}\otimes P_{{Z}}$, then

[TABLE]

Proof.

1. The generalization error can be written as in (18), where $W$ and $Z_{i}$ in the second term are dependent with $P_{W,Z_{i}}=\mu\otimes P_{W|Z_{i}}$, and $W$ and $\widetilde{Z}$ in the first term are independent whose marginal distributions are the same as those of $W$ and $Z_{i}$. The first inequality then follows from Lemma 1 by letting $S=Z_{i}$ and $n=1$, for each $i=1,\cdots,n$.

2. For an $R$-sub-Gaussian random variable, $\psi^{-1}_{+}(y)=\psi^{-1}_{-}(y)=\sqrt{2R^{2}y}$ is an upper bound on its CGF. The second inequality then follows from Theorem 2. ∎

Remark 1**.**

The condition that $\ell(w,Z)$ is $R$-sub-Gaussian under $Z\sim\mu$ for all $w\in\mathcal{W}$ in the first part of Proposition 1 is the same as the one in Lemma 1, which is not stronger than the condition in the second part of Proposition 1. An example was given in [32] for this argument. Specifically, consider $\mathcal{W}=\mathcal{Z}=\mathbb{R}$, with $\ell(w,z)=w+z$, and $(\widetilde{W},\widetilde{Z})\sim\text{Cauchy}\otimes\mathcal{N}(0,1)$. Then, $\ell(w,Z)$ is 1-sub-Gaussian for any $w\in\mathcal{W}$, whereas $\ell(\widetilde{W},\widetilde{Z})$ does not even have bounded absolute first moment.

The following proposition shows that the proposed ISMI bound is always no worse than the bound using $I(S;W)$ in Lemma 1 and [31, Theorem 2].

Proposition 2**.**

Suppose that $S=\{Z_{1},\cdots,Z_{n}\}$ consists of $n$ independent samples, and $\psi^{*-1}$ is a concave function, then

[TABLE]

Proof.

By the chain rule of mutual information,

[TABLE]

where $Z^{j}=\{Z_{1},\cdots,Z_{j}\}$. Note that $Z_{i}$ and $Z^{i-1}$ are independent, i.e., $I(Z_{i};Z^{i-1})=0$ , it then follows that

[TABLE]

Thus,

[TABLE]

and applying Jensen’s inequality completes the proof. ∎

Remark 2**.**

Under the sub-Gaussian condition (see sentence following Lemma 1), we can let \psi^{*-1}\big{(}y\big{)}=\sqrt{2R^{2}y}. Then by Proposition 2, the ISMI bound in Proposition 1 is always no worse than the bound based on $I(S;W)$ in Lemma 1.

Remark 3**.**

Following arguments similar to those used in the proof of $I(W;Z_{i})\leq I(W;Z_{i}|Z^{i-1})$, we can also show that $I(W;Z_{i})\leq I(W;Z_{i}|S^{-i})$, where $S^{-i}$ denotes the set obtained by deleting $Z_{i}$ from $S$. Therefore, the ISMI bound is always no worse than the bound based on $I(W;Z_{i}|S^{-i})$ in [23, Theorem 2].

In the next section, we will also show via several examples that the ISMI bound provides a more accurate characterization of the generalization error than the bound in Lemma 1 and the chaining bound in [19].

IV Examples with Infinite $I(W;S)$

In this section, we consider two examples of learning algorithms with infinite $I(W;S)$. We show that for these examples, the upper bound on generalization error in Lemma 1 blows up, whereas the ISMI bound in Theorem 2 still provides an accurate approximation. The details of the derivations of the bounds can be found in the Appendices.

IV-A Estimating the Mean

We first consider the problem of learning the mean of a Gaussian random vector $Z\sim\mathcal{N}(\mu,\sigma^{2}I_{d})$, which minimizes the square error $\ell(w,Z)\triangleq\|w-Z\|_{2}^{2}$. The empirical risk with $n$ i.i.d. samples is

[TABLE]

The empirical risk minimization (ERM) solution is the sample mean $W=\frac{1}{n}\sum_{i=1}^{n}Z_{i}$, which is deterministic given $S$. Its generalization error can be computed exactly as (see Appendix A):

[TABLE]

The bound in Lemma 1 is not applicable here due to the following two reasons: (1) $W$ is a deterministic function of $S$, and hence $I(S;W)=\infty$; and (2) since $Z$ is a Gaussian random vector, the loss function $\ell(w,Z)=\|w-Z\|_{2}^{2}$ is not sub-Gaussian for all $w\in\mathbb{R}^{d}$. Specifically, the variance of the loss function $\ell(w,Z)$ diverges as $\|w\|_{2}\to\infty$, which implies that a uniform upper bound on $\Lambda_{\ell(w,Z)}(\lambda)$, $\forall w\in\mathbb{R}^{d}$ does not exist.

We can get around both of these issues by applying the ISMI bound in Theorem 2. Since $W\sim\mathcal{N}(\mu,\frac{\sigma^{2}I_{d}}{n})$, the mutual information between each individual sample and the output hypothesis $I(W;Z_{i})$ can be computed exactly as (see Appendix A):

[TABLE]

In addition, since $W\sim\mathcal{N}(\mu,\frac{\sigma^{2}I_{d}}{n})$, it can be shown that $\ell(W,\widetilde{Z})\sim\sigma_{\ell}^{2}\chi^{2}_{d}$, where $\sigma_{\ell}^{2}\triangleq\frac{(n+1)\sigma^{2}}{n}$, and $\chi^{2}_{d}$ denotes the chi-squared distribution with $d$ degrees of freedom. Note that the expectation of $\chi^{2}_{d}$ distribution is $d$ and its moment generating function is $(1-2\lambda)^{d/2}$. Therefore, the CGF of $\ell(\widetilde{W},\widetilde{Z})$ is given by

[TABLE]

for $\lambda\in(-\infty,\frac{1}{2\sigma_{\ell}^{2}})$. Since $W$ is the ERM solution, it follows that $\mathrm{gen}(\mu,P_{W|S})\geq 0$, and we only need to consider the case $\lambda<0$. It can be shown that (see Appendix A):

[TABLE]

Then, $\psi_{-}^{*-1}(y)=2\sqrt{d\sigma_{\ell}^{4}y}.$ Combining the results in (27), we have

[TABLE]

As $n\to\infty$, the above bound is $\mathcal{O}\left(\frac{1}{\sqrt{n}}\right)$, which is sub-optimal compared to the true generalization error computed in (26). We should note that techniques based on VC dimension [7] and algorithmic stability [13] also yield bounds of $\mathcal{O}\left(\frac{1}{\sqrt{n}}\right)$.

IV-B Gaussian Process

In this subsection, we revisit the Gaussian process example studied in [19]. Let $\mathcal{W}=\{w\in\mathbb{R}^{2}:\|w\|_{2}=1\}$, and $Z\sim\mathcal{N}(0,I_{2})$ be a standard normal random vector in $\mathbb{R}^{2}$. The loss function is defined to be the following Gaussian process indexed by $w$:

[TABLE]

Note that the loss function333The loss function can be negative here. We ignore the non-negativity assumption of the loss function; this does not affect our analysis. $\ell(w,Z)$ is sub-Gaussian with parameter $R=1$ for all $w\in\mathcal{W}$. In addition, the output hypothesis $w\in\mathcal{W}$ can also be represented equivalently using the phase of $w$. In other words, we can let $\phi$ be the unique number in $[0,2\pi)$ such that $w=(\sin\phi,\cos\phi)$. For this problem, the empirical risk of a hypothesis $w\in\mathcal{W}$ is given by

[TABLE]

We consider two learning algorithms which are the same as the ones in [19]. The first is the ERM algorithm:

[TABLE]

The second is the ERM algorithm with additive noise:

[TABLE]

where the noise $\xi$ is independent of $S$, and has an atom with probability mass $\epsilon$ at 0, and probability $1-\epsilon$ uniformly distributed on $(-\pi,\pi)$. Due to the symmetry of the problem, $W$ and $W^{\prime}$ are uniformly distributed over $[0,2\pi)$.

The generalization error of the ERM algorithm $W$ can be computed exactly as (see Appendix B):

[TABLE]

For the second algorithm $W^{\prime}$, since the noise $\xi$ is independent from $S$, it follows that

[TABLE]

The bound via $I(W;S)$ in Lemma 1 is not applicable, since $W$ is deterministic given $S$ and $I(W;S)=\infty$. Moreover, for the second algorithm $W^{\prime}$,

[TABLE]

since $\xi$ has a singular component at 0, and $h(\xi)=-\infty$.

Applying the ISMI bound in Theorem 2 to the ERM algorithm $W$, we have that

[TABLE]

which we need to compute the conditional distribution $P_{W|Z_{i}=z_{i}}$. Note that given $Z_{i}=z_{i}$, the ERM solution

[TABLE]

depends on the other samples $Z_{j}$, $j\neq i$. Moreover, it can be shown that $P_{W|Z_{i}=z_{i}}$ is equivalent to the phase distribution of a Gaussian random vector $\mathcal{N}(\frac{z_{i}}{n},\frac{n-1}{n^{2}}I_{2})$ in polar coordinates.

Due to symmetry, we can always rotate the polar coordinates, such that $z_{i}=(r,0)$, where $r\in\mathbb{R}^{+}$ is the $\ell_{2}$ norm of $z_{i}$. Then, $P_{W|Z_{i}=z_{i}}$ is a function of $r$, and can be equivalently characterized and computed by the distribution f\big{(}\phi\big{|}\|Z_{i}\|=r\big{)} provided in Appendix B. Since the norm of $Z_{i}$ has a Rayleigh distribution with unit variance, it then follows that

[TABLE]

Applying Theorem 2, we obtain

[TABLE]

Similarly, we can compute the ISMI bound for $W^{\prime}$.

Numerical comparisons are presented in Fig. 1 and Fig. 2. In both figures, we plot the ISMI bound, the CMI bound in [19], and the true values of the generalization error, as functions of the number of samples $n$. In Fig. 1, we compare these bounds for the ERM solution $W$. Note that the CMI bound reduces to the classical chaining bound in this case. In Fig. 2, we evaluate these bounds for the noisy algorithm $W^{\prime}$ with $\epsilon=0.05$. Both figures demonstrate that the ISMI bound is closer to the true values of the generalization error, and outperforms the CMI bound significantly. More details about the computations of both bounds can be found in Appendix B.

V Noisy, Iterative Algorithms

In this section, we apply the ISMI bound in Theorem 2 to a class of noisy, iterative algorithms, specifically, stochastic gradient Langevin dynamics (SGLD).

V-A SGLD Algorithm

We begin by introducing some notation to be used in this section. Denote the parameter vector at iteration $t$ by $W_{(t)}\in\mathbb{R}^{d}$, and let $W_{(0)}\in\mathcal{W}$ denote an arbitrary initialization. At each iteration $t\geq 1$, we sample a training data point $Z_{U_{(t)}}\in S$, where $U_{(t)}\in\{1,...,n\}$ denotes the random index of the sample selected at iteration $t$, and compute the gradient $\nabla\ell(W_{(t-1)},Z_{U_{(t)}})$. We then scale the gradient by a step size $\eta_{(t)}$ and perturb it by isotropic Gaussian noise $\xi\sim\mathcal{N}(0,I_{d})$. The overall update rule is as follows [33]:

[TABLE]

where $\sigma_{(t)}$ controls the variance of the Gaussian noise.

For $t\geq 0$, let $W^{(t)}\triangleq\{W_{(1)},\ \cdots,\ W_{(t)}\}$ and $U^{(t)}\triangleq\{U_{(1)},\cdots,U_{(t)}\}$. We assume that the training process takes $K$ epochs, and the total number of iterations is $T=nK$. The output of the algorithm is $W=W_{(T)}$.

In the following, we use the same assumptions as in [29].

Assumption 1**.**

$\ell(w,Z)$* is $R$-sub-Gaussian with respect to $Z\sim\mu$, for every $w\in\mathcal{W}$.*

Assumption 2**.**

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

In [29], the following bound was obtained by upper bounding $I(W;S)$ in Lemma 1.

Lemma 3** ([29, Corollary 1]).**

The generalization error of the SGLD algorithm is bounded by

[TABLE]

V-B ISMI Bound for SGLD

We have the following proposition which characterizes the ISMI bound for the SGLD algorithm.

Proposition 3**.**

Suppose Assumption 1 and 2 hold, then we have the following ISMI bound on the generalization error for SGLD algorithm,

[TABLE]

where $U^{(T)}$ denotes the random sample path, and ${\mathcal{T}}_{i}(U^{(T)})$ denote the set of iterations for which samples $Z_{i}$ is selected for a given sample path $U^{(T)}$.

Proof.

To apply the ISMI bound for SGLD, we modify the result in Theorem 2 by conditioning on the random sample path $U^{(T)}$,

[TABLE]

where $\mathcal{U}$ denotes the set of all possible sample paths, and $I(W;Z_{i}|U^{(T)}=u^{(T)})$ is the mutual information 444Note that this mutual information is different from the conditional mutual information $I(W;Z_{i}|U^{(T)})=\mathbb{E}_{U^{(T)}}[I(W;Z_{i}|U^{(T)}=u^{(T)})]$. with conditional distribution $P(W,Z_{i}|U^{(T)}=u^{(T)})$.

Let ${\mathcal{T}}_{i}(u^{(T)})$ denote the set of iterations for which sample $Z_{i}$ is selected for a given sample path $u^{(T)}$. Using the chain rule of mutual information, we have

[TABLE]

where the last equality is due to the fact that given $u^{(T)}$ and $W_{(\tau-1)}$, $Z_{i}$ is independent of $W_{(\tau)}$, if $\tau\notin{\mathcal{T}}_{i}(u^{(T)})$. For $\tau\in{\mathcal{T}}_{i}(u^{(T)})$, i.e., if $Z_{i}$ is selected at iteration $\tau$, we have

[TABLE]

Since we assume that $\sup_{w\in\mathcal{W},Z\in\mathcal{Z}}\|\nabla\ell(W,Z)\|_{2}\leq L$, we have

[TABLE]

Due to the fact that $\xi$ is an independent Gaussian noise, h(\sigma_{(\tau)}\xi|W_{(\tau-1)})=\frac{d}{2}\log\big{(}2\pi e\sigma_{\tau}^{2}\big{)}, we have

[TABLE]

Combining with (V-B), it follows that

[TABLE]

where we remove the $\log$ term by using $\log(1+x)\leq x$. ∎

To compare the result of the ISMI bound in Proposition 3 and the bound in Lemma 3, we specify the parameters in the SGLD algorithms. As in [29], we set $\eta_{(t)}=\frac{c}{t}$, and $\sigma_{(t)}=\sqrt{\eta_{t}}$. We use the following “without replacement” sampling scheme for SGLD to further simplify the computation. Specifically, for the $k$-th training epoch, i.e., from the $((k-1)n+1)$-th to $kn$-th iterations, all training samples in $S$ are used exactly once.

Then, the ISMI bound can be further bounded as follows:

[TABLE]

where $(a)$ follows from the sampling scheme that all samples are used exactly once in each epoch; $(b)$ is due to the fact that $\sum_{k=1}^{K}\frac{1}{k}\leq\log(K)+1$; and $(c)$ follows by computing the integral $\int_{0}^{1}\sqrt{\frac{1}{x}+1+\log(K-1)}dx$.

Comparing with the bound in [29],

[TABLE]

it can be seen that our bound is tighter by a factor of $\sqrt{\log n}$ with the “without replacement” sampling scheme.

Remark 4**.**

*We note that for the typical use of SGLD, the standard deviation of the noise is $\sigma_{t}=\sqrt{2\eta_{t}/\beta_{t}}$, where $\beta_{t}$ denotes the inverse temperature at iteration $t$, and it is often set to be $\Theta(n)$. It is clear that $\beta=\Theta(n)$ will lead to a generalization bound that does not decay in $n$. Here, we choose $\beta_{t}=2$ for comparison with the bound in [29], while in practice $\beta$ may be a function of $n$ and grow with $t$. An analysis of the generalization error bound of SGLD for arbitrary choices of $\beta_{t}$ can be found in [32]. *

VI Empirical Evaluation of ISMI Bound for Logistic Regression

For some learning algorithms applied in practice, it is difficult to analytically characterize $P_{W|S}$, which makes the analytical evaluation of the ISMI bound challenging. In this section, we provide such an example, that of logistic regression, for which it is difficult to analytically characterize the learning algorithm via the conditional distribution $P_{W|S}$. We therefore empirically evaluate the ISMI bound via a mutual information estimator, and compare it to an empirical evaluation of the generalization error. We further note that the ISMI bound is much easier to estimate than the bound in Lemma 1 and the chaining bound in [19] due to the significant reduction in dimension that comes from estimating $I(Z_{i};W)$ instead of $I(S;W)$.

Consider the binary classification problem, where the samples $Z=(X,Y)$, consisting of features $X\in\mathbb{R}^{d}$ and labels $Y\in\{\pm 1\}$. We assume that training samples are generated from the following distribution,

[TABLE]

The marginal distribution of $X$ is the mixture of two Gaussian distributions, and we assume that $P(Y=-1)=P(Y=1)=1/2$.

A binary classifier is constructed as follows:

[TABLE]

We adopt classification error $\ell(w,Z)=\mathds{1}_{\{Y\neq\hat{Y}\}}$ to compute the generalization error, then the empirical risk with $n$ i.i.d. samples is

[TABLE]

Since the empirical risk function is not differentiable, we learn $W$ by minimizing the following loss function of logistic regression:

[TABLE]

In general, it is difficult to obtain a closed form solution for this optimization problem, and therefore, (56) is usually solved numerically. This makes it difficult to analytically characterize the conditional distribution $P_{W|Z_{i}}$, which in turn makes it challenging to compute the generalization error and the ISMI bound analytically.

Alternatively, we can empirically estimate the generalization error and the ISMI bound. Specifically, we train $W$ for $N$ times using $N$ sets of independent samples, and we use the K-nearest neighbor based mutual information estimator [34, 35] to estimate $I(W;Z_{i})$ with $N$ i.i.d. samples of $W$ and $Z_{i}$. Note that the K-nearest neighbor based mutual information estimator is consistent, and its mean squared estimation error can be upper bounded by $\mathcal{O}(N^{-\frac{2}{d_{W}+d_{Z}}})$, where $d_{W}=d$ is the dimension of the weights $W$, and $d_{Z}=d+1$ is the dimension of $Z$ [35]. Moreover, since we use classification error to compute generalization error, $\ell(W,Z)$ is bounded by 1. Then, by Hoeffding’s lemma, $\ell(W,Z)$ is $\frac{1}{2}$-sub-Gaussian. Thus, the ISMI bound can be estimated by

[TABLE]

where $\hat{I}(W;Z_{i})$ is the estimate of $I(W;Z_{i})$. If we apply an optimization algorithm that does not depend on the order of the samples, e.g., gradient descent and stochastic gradient descent with random shuffling, we then only need to estimate one $\sqrt{\frac{\hat{I}(W;Z_{1})}{2}}$ instead of estimating $\hat{I}(W;Z_{i})$ for all $1\leq i\leq n$.

We note that the bound in Lemma 1 is difficult to estimate due to the high dimension of $S$, which scales linearly with $n$. Specifically, the training dataset $S$ consists of $n$ samples, and therefore $I(W;S)$ is the mutual information between two random vectors with dimensions $d$ and $n(d+1)$. As shown in [35], due to the curse of dimensionality, it is impossible to construct a consistent mutual information estimator for large $n$. We also note that the exact computation of Wasserstein distances is costly in general, as it requires the solution of an optimal transport problem[36]. Moreover, similar high dimensional issue makes it even more difficult to directly estimate $\mathbb{W}(P_{W},P_{W}|S)$ in the Wasserstein distance based generalization bound in [24, 25].

In Fig. 3, we plot an empirical estimate of the ISMI bound using (57), and compare it to the generalization error. In the simulation, we chose the following model parameters: $d=2$ and $\mu_{1}=(1,1)$, $\mu_{-1}=(-1,-1)$ with $\Sigma=4I$. We used the K-nearest neighbor based mutual information estimator (revised KSG estimator) in [35] with $N=5000$ i.i.d. samples. It can be seen that the ISMI bound has a similar convergence behavior as the true generalization error as number of training samples $n$ increases.

VII Conclusions

In this paper, we proposed a tighter information-theoretic upper bound on the generalization error using the mutual information $I(Z_{i};W)$ between each individual training sample $Z_{i}$ and the output hypothesis $W$ of the learning algorithm. We showed that compared to existing studies, our bound is more broadly applicable, and is considerably tighter. More importantly, the individual sample mutual information is between two vectors whose dimensions do not scale with the sample size $n$. Therefore, unlike the existing bounds in [19, 25, 18], the ISMI bound can easily be evaluated empirically in practice. As suggested by recent works, the proposed ISMI bound could be further improved by combining with the chaining method [22], or data-dependent estimates [32]. The proposed information-theoretic framework can also be used to guide model compression in deep learning [37].

Appendix A Section IV-A Details

A-A Generalization Error

For this example, the generalization can be computed as

[TABLE]

Since ${\rm{Cov}}[Z]=\sigma^{2}I_{d}$, we have $\mathrm{gen}(\mu,P_{W|S})=\frac{2\sigma^{2}d}{n}$.

A-B Individual Mutual Information

We note that both $W$ and $n$ i.i.d. samples $Z_{i}$ are Gaussian, then the individual mutual information $I(W;Z_{i})$ is a function captured by the following covariance matrix,

[TABLE]

Then, we have

[TABLE]

for all $i=1,\dots,n$.

A-C Upper bound for CGF

Note that the CGF of $\ell(\widetilde{W},\widetilde{Z})$ is given by

[TABLE]

where $u\triangleq 2\sigma_{\ell}^{2}\lambda$. Further note that

[TABLE]

We therefore have the following upper bound on the CGF of $\ell(\widetilde{W},\widetilde{Z})$:

[TABLE]

Appendix B Section IV-B Details

B-A Generalization Error

Note that the expectation of the population risk is

[TABLE]

since $W$ and $Z$ are independent. Then, the generalization error can be computed as

[TABLE]

where the last step is due to the fact that the distribution of $\|\frac{1}{n}\sum_{i=1}^{n}Z_{i}\|_{2}$ is $\mathrm{Rayleigh}(\frac{1}{n})$.

B-B Individual Sample Mutual Information Bound

To compute the ISMI bound, we need the conditional distribution $P_{W|Z_{i}=z_{i}}$. Note that given $Z_{i}=z_{i}$, the ERM solution is

[TABLE]

Also note that since $\phi\in[0,2\pi)$, $P_{W|Z_{i}=z_{i}}$ is equivalent to the phase distribution of a Gaussian random vector $\mathcal{N}(\frac{z_{i}}{n},\frac{n-1}{n^{2}}I_{2})$ in polar coordinates. Since entropy is shift-invariant, we can always rotate the polar coordinates, such that $z_{i}=(r,0)$, where $r\in\mathbb{R}^{+}$ is the $\ell_{2}$ norm of $z_{i}$.

The joint distribution of radius and phase $(\rho,\phi)$ in polar coordinates can be obtained by applying the Jacobian method to the Gaussian distribution $\mathcal{N}(\frac{(r,0)}{n},\frac{n-1}{n^{2}}I_{2})$, and we have

[TABLE]

for $\rho\in[0,\infty),\ \phi\in[0,2\pi)$.

Then, the marginal distribution of $\phi$ can be computed by integrating out $\rho$ from the joint distribution f\big{(}\rho,\phi\big{|}\|Z_{i}\|=r\big{)}:

[TABLE]

where $Q(x)$ is the complementary cumulative distribution function of the standard normal distribution.

Thus, the ISMI bound for the ERM algorithm $W$ can be evaluated using the following expression via numerical integration:

[TABLE]

Similarly, the ISMI bound for algorithm $W^{\prime}$ can be computed via numerical integration using

[TABLE]

where the conditional distribution $P_{W^{\prime}|Z_{i}=z_{i}}$ can be characterized by the phase distribution

[TABLE]

B-C Chaining Mutual Information Bound

The CMI bound is computed based on the values provided in Table 1 in [19]. We note that the CMI bound in [19] is evaluated for the case $n=1$, i.e., there is only one training sample. To plot the CMI bound in [19] as a function of the number of samples $n$, we normalize the CMI bound by a $\sqrt{n}$ factor in Figures 1 and 2, since $\mathcal{O}(1/\sqrt{n})$ is the true convergence rate for the generalization error as shown in (B-A). For instance, Table 1 in [19] shows that the CMI bound for the ERM solution $W$ is 19.0352 when $n=1$, which is equivalent to the classical chaining bound. We therefore plot the curve $\frac{19.0352}{\sqrt{n}}$ as the CMI bound for comparison with the proposed ISMI bound in Figure 1.