On Leave-One-Out Conditional Mutual Information For Generalization
Abstract
We derive information theoretic generalization bounds for supervised learning algorithms based on a new measure of leave-one-out conditional mutual information (loo-CMI).
In contrast to other CMI bounds, which may be hard to evaluate in practice, our loo-CMI bounds are easier to compute and can be interpreted in connection to other notions such as classical leave-one-out cross-validation, stability of the optimization algorithm, and the geometry of the loss-landscape.
It applies both to the output of training algorithms as well as their predictions.
We empirically validate the quality of the bound by evaluating its predicted generalization gap in scenarios for deep learning.
In particular, our bounds are non-vacuous on image-classification tasks.
Introduction
Generalization in classical machine learning models is often understood in terms of the bias-variance trade-off: over-parameterized models fit the data better but are more likely to overfit.
However, the push for ever larger deep network models, driven by their remarkable empirical generalization, has spurred a race for the development of new theories of generalization with two main objectives: (1) guiding the practitioners in designing over-parameterized architectures and training schemes that would generalize better and (2) providing provable bounds on the performance of the model after deployment.
In this context information theoretic bounds on generalization [27, 31, 5, 21, 29, 12, 16] are particularly appealing.
They capture the intuitive idea that models that memorize less about the training set will generalize better.
In particular, since they deal with the amount of information stored in the weights rather than the number of weights, they are especially adaptable to over-parameterized models [17, 2, 1].
However, the bounds are often vacuous in realistic settings, require modified training algorithms, or are difficult to compute and/or to relate to relevant properties of the network/training algorithm.
To remedy this problem, we ask whether it is possible to develop a bound that is: (a) realistically tight for standard training algorithms used in deep learning; (b) interpretable, meaning that they can be rewritten in terms of quantities the practitioners can control during the training process and that may guide toward the design of better training algorithms, (c) that applies out-of-the-box to standard training algorithms.
Contributions.
As a step in this direction, we introduce a new information theoretic bound based on leave-one-out conditional mutual information (loo-CMI) 1 .
The main intuition behind the bound is to ask the counterfactual question “If we were to remove one sample at random from the training set and retrain, how well would we be able to infer which sample was removed?”.
As we shall see, bounding this quantity is enough to bound sharply the amount of memorization in a deep network, and hence its generalization (Theorems 3.1 and 3.2).
At the same time, the bound is easy to interpret and connects to several empirical quantities and other generalization theories, such as: (1) stability of the optimization algorithm
(Section 4.4), (2) flatness of the loss landscape at convergence (Section 4.5).
Moreover, we show that our bound can also be interpreted in the function space (that is, looking at the activations rather than the weights, (Theorem 3.2) leading to tighter bounds for over-parameterized models through the data processing inequality.
Our loo-CMI bound falls in the general class of conditional mutual information bounds, introduced by [29].
However, while standard CMI bounds iterate over all the subset of samples of size N ( 2 N subsets for a dataset of 2 N samples) in order to compute the bound, our loo-CMI bounds only need to iterate over N possible samples to remove, leading to an exponential reduction in the cost to estimate the bound.
Moreover, our strategy of removing a single sample is easy to relate to the stability of the training path to slight perturbation of the training set, thus giving a clear connection between generalization and the geometry of the loss landscape (Figure 1).
Empirically, we show that our bound can be computed and is non-vacuous for state-of-the-art deep networks fine-tuned on standard image classification tasks (Table 1).
We also study the dependency of the bound on the size of the dataset (Figure 2), and the hyper-parameters (Figure 3).
Preliminaries
We denote with [ n ] the set { 1 , . . . , n } and with R d the set of real-valued and complex-valued d − dimensional vectors.
A T and A − 1 denote the transpose and inverse of a square matrix A .
If x is a vector with n elements, then diag ( x ) denotes the n × n square matrix with the elements of x on its main diagonal; ∥·∥ is the standard Euclidean norm on vectors.
We use uppercase letters to denote random variables, and lowercase letters to denote a specific value the random variable can take.
If X is a random variable, then p X is the probability distribution of X , and E x ∼ p X
[ x ] denotes the expected value of X .
For a pair of random variables X and Y , p ( X,Y ) and p X p
Y are the joint and product distributions, respectively.
If p and q are two probability densities over R d such that p is absolutely continuous with respect to q ( p ≪ q ), the Kullback-Leibler Divergence from p to q is defined as:
[Equation]
The Mutual Information between two random variables X and Y , and the conditional mutual information of X and Y given a third random variable Z are respectively given by:
[Equation]
2.1 Problem Formulation
Let Z n = { Z 1 , . . .
Z n } ∈ Z n be a collection of n independent and identically distributed samples drawn from some unknown distribution P .
We consider the setting of supervised learning where
Z = X
× Y ; i.e., each sample is composed of a feature vector and a label.
Let A : Z n → W ⊆ R K be a possibly stochastic training algorithm, where W is a set of weights (e.g., weights in a neural network).
Given a loss function ℓ : W × Z → R + , the empirical risk for weights w on Z n is given by (cid:98) L ( w, Z n ) =
1 /n (cid:80) ni =1 ℓ ( w, Z i ) , and the "true" loss is given by L ( w )
=
E Z ′ ∼P
[ ℓ ( w, Z ′ )] ; i.e., the average loss w incurs on random samples drawn from P .
Our goal is to find a w ∗ ∈ W such that w ∗ =
arg min w ∈W L ( w ) .
Since we do not have access to the data-generating distribution P , computing the true loss is infeasible.
However, we do have access to a set of samples Z n , and one approach is empirical risk minimization where we find the w ∈ W that minimizes (cid:98) L ( w, Z n ) .
If the difference L ( w ) − (cid:98) L ( w, Z n ) is small, then the true loss of an empirical risk minimizer would be nearly optimal within the considered hypothesis class.
Therefore, it is of great interest to study the generalization gap of A (given n − 1 samples for ease of notation):
[Equation]
2.2 Conditional Mutual Information Bounds
We start by recalling the main results on Conditional Mutual Information (CMI) bounds that we return to later.
Let ˜ Z 2 n ∈ Z n × 2 be a dataset with 2 n samples drawn from a distribution P , grouped into n pairs.
Let S ∼ S
=
Uniform ( { 0 , 1 } n ) a uniform random binary vector of size n .
S selects one sample from each pair in ˜ Z 2 n to form ˜ Z nS ∈ Z n given by ( ˜ Z nS ) i =
˜
Z 2 nS i +1 .
In this setting, the bound of Steinke and Zakynthinou [29] arises from asking the following question : “if a model is trained with a subset of samples chosen through the random binary indexing vector S , how much information does the output of the algorithm provide about S ?”
Here, the Shannon Mutual Information is useful because it is a measure of the amount of ‘information’ obtained about one random variable after observing another random variable.
Theorem 2.1 (Theorem 2, [29]) .
Let ℓ : W × Z → [0 , 1] be a bounded loss function.
Let A be a possibly stochastic training algorithm.
Let W
=
A ( ˜ Z nS ) be the output of A given the dataset ˜ Z nS , then
[Equation]
Haghifam et al. [12] improved Theorem 2.1 by moving the expectation over ˜ Z 2 n outside the square root, and by measuring the information the output of the algorithm provides on random subsets of the indexing vector S .
Harutyunyan et al.
[16] further improved the bound by moving the expectation over the random subsets outside the square root.
Theorem 2.2 (Theorem 2.6, [16]) .
Let ℓ : W × Z → [0 , 1] be a bounded loss function.
Let m ∈ [ n ] and let V ∼ V be a random subset of [ n ] of size m .
Let W
=
A ( ˜ Z S ) be the output of A given the dataset ˜ Z S , then
[Equation]
One appealing aspect of the CMI bounds is that they work for any training algorithm and any model, including over-parameterized networks trained with Stochastic Gradient Descent (SGD).
However, computing the bounds requires iterating over all possible 2 N values of S .
Moreover, the bound is difficult to interpret since the value of W =
A ( ˜ Z nS ) for different values of S can vary significantly and in unpredictable ways for large non-convex models.
Taking inspiration from leave-one-out cross validation, we propose a different CMI bound which avoids both problems by removing a single sample from the dataset.
As we shall see, the bound allows for interpretable expressions, faster computation, and a connection to notions of stability.
Leave-One-Out Conditional Mutual Information
Let Z n = { Z 1 , Z 2 , . . .
Z n } ∈ Z n be a dataset of n i.i.d.
samples drawn from P .
Let U ∼ U =
Uniform([ n ]) be uniform random variable taking values over the indices of the samples in Z .
U removes a single sample from Z n to form Z n − 1 − U =
Z n \ Z U , the dataset without the U th sample.
Let W =
A ( Z n − 1 − U ) be the output of the algorithm, then we define the pointwise leave-one-out conditional mutual information as:
[Equation]
In this setting, loo - CMI( A , z n ) measures the amount of information that the output of the algorithm
W reveals about U , the index (identity) of the left-out sample.
As is the case with the conditional mutual information terms in Theorems 2.1 and 2.2, loo - CMI is bounded.
Specifically, loo - CMI is upper bounded by the entropy of U , H ( U ) = log( n ) .
Hence, it does not suffer from the same issues as information stability bounds [27, 31].
Moreover, the output of an algorithm is not significantly affected by the inclusion or exclusion of one sample when the input consists of thousands of other samples, so we expect loo - CMI to be small for large values of n .
We use (7) to derive a bound on the generalization of error in expectation (Theorem 3.1).
3.1 Generalization Bounds Based on loo - CMI
loo - CMI is deeply intertwined with leave-one-out cross validation, so before we derive bounds on the generalization through loo - CMI , it is helpful to first obtain a probabilistic upper bound on the leave-one-out cross validation error ( loo - cv ).
loo - cv measures the difference in loss between the samples which the algorithm trained on, and the sample that was left out.
We begin with a precise definition of loo - cv and a probabilistic bound on it.
[Equation]
Lemma 3.1.
Let ℓ : W × Z → [0 , 1] be a bounded loss function, then for all t > 0 , w ∈ W , and z n ∈ Z ,
[Equation]
To derive the generalization bounds included in this paper, we make heavy use of Lemma 3.1.
We begin with a generalization bound in expectation.
Theorem 3.1.
Let A : Z n − 1 → W be a training algorithm.
Let ℓ : W × Z → [0 , 1] be a bounded loss function, then
[Equation]
where c n is given in (9).
The bound in (10) guarantees a good generalization error when loo - CMI( A , z n ) is small, i.e., when the output of the algorithm W tells us little about the identity of the sample that was removed from the training dataset.
If a parametric learning algorithm ends up memorizing the samples in the training set, the generalization gap could be large, and so would loo - CMI as we can determine U by checking which sample was not memorized by the algorithm.
On the other extreme, a learning algorithm which outputs a constant parameter regardless of the training set does not have a generalization gap, and in this case loo - CMI would equal 0 as W provides no information about U .
We note that the bound of Theorem 3.1 most closely resembles the bound of
Theorem 2.2 when m = 1 .
The latter bound is given by:
[Equation]
Recalling the setting of Section 2.2, (11) is computed by first going through every possible value of
S ∈ { 0 , 1 } n .
For every value of S and for each i ∈ [ n ] , one fixes S =
( S 1 , S 2 , . . .
, S n ) excluding
S i which varies uniformly over { 0 , 1 } .
In other words, the i th component of the training set ˜ Z nS is equally likely to be either ˜
Z i, 1 or ˜ Z i, 2 .
The right-hand-side of (11) then bounds the generalization using the information that the weights contain about S i .
In other words, the question asked is: “can the output of the algorithm help us determine which one of ˜ z
i, 1 or ˜ z i, 2 was present in the training set?".
On the other hand, we ask a different question: "can the output of the algorithm help us determine the index of the sample that was removed?”.
This subtle change leads to an exponential reduction in the cost of computing the since we only need to iterate over the values of U ∈ [ n ] , as opposed to S which can take 2 n possible values.
3.2 Extension To The Function Space
The previous bounds used the output of the algorithm, the weights, to bound the generalization error.
A different approach studied in [16] is to assume that given a training dataset and a test sample, the algorithm outputs a prediction on the test sample.
In other words, we assume the algorithm is a possibly stochastic function h : Z n − 1 × X −→ R which takes a training set z , a new unlabeled sample x new , and outputs a possibly stochastic prediction h ( z, x new ) on the new sample.
The set of predictions R can be different than the set of labels
Y (e.g., class probabilities in multi-class classification tasks).
Unlike bounds with respect to the weights, this approach is applicable to both parametric algorithms (e.g. neural networks) and non-parametric algorithms (e.g. k -nearest neighbors).
In the former case, the output of the algorithm A specifies a prediction function g ∈ { g w : X → R | w ∈ W} from a class of functions parameterized by the weights i.e. h ( z, x new ) =
g A ( z ) ( x new ) .
We redefine the loss to be a non-negative function ℓ : R × Y → R + that measures the distance between the predictions and the ground-truth labels.
Given ℓ and a training set Z n , the true loss of an algorithm h is given by L ( h, Z n ) =
E ( x ′ ,y ′ ) ∼P
[ ℓ ( h ( Z n , x ′ ) , y ′ )] , and the empirical loss is given by (cid:98) L ( h, Z n ) = 1 /n (cid:80)
[Equation]
=
[Equation]
ℓ ( h ( Z n , X i ) , Y i ) .
The generalization gap of h given is then defined as:
[Equation]
Recalling the setup of Section 3, we define the pointwise functional leave-one-out conditional mutual information floo - CMI as:
[Equation]
While loo - CMI measures the reduction in uncertainty about U after having known the weights, floo - CMI measures the reduction in uncertainty after having known the predictions made on the whole dataset (including the prediction on the removed sample).
We derive a bound with respect to floo - CMI( h, z n ) using a proof that closely resembles the proof of Theorem 3.1.
Theorem 3.2.
Let ℓ : R × Y → [0 , 1] be a bounded loss function.
Let U be a uniform random variable over [ n ] , then
[Equation]
14) where c n is given in (9).
For a parametric algorithm A with prediction function h ,
U — A ( z n − 1 − U ) — h ( z n − 1 − U , X ) is a Markov chain.
By the data-processing inequality, we have that floo - CMI( h, z n ) ≤ loo - CMI( A , z n ) .
Since the bounds of Theorems 3.1 and 3.2 differ only through loo - CMI and floo - CMI , the bound of Theorem 3.2 is sharper.
Computing The Bounds 4.1 Computable Upper Bounds to loo - CMI and floo - CMI
To evaluate and interpret the bounds for a parametric algorithm A with prediction function h , we require a computable closed-form expressions for loo - CMI and floo - CMI .
However, obtaining a closed form expression is a difficult task in most cases.
To see the reason behind this difficulty, note that loo - CMI can be written as:
[Equation]
(15) where p W | z n = E u ∼U (cid:2) p W | z n ,u (cid:3) is a mixture distribution.
Since one of the terms in the KullbackLeibler divergence of (15) is a mixture distribution, it is unlikely that one could find a closed-form expression for the mutual information except for simple cases (e.g., we cannot obtain a closedexpression for Gaussian mixtures[8]).
To find a computable and interpretable bound on the generalization error, we propose upper bounds on the conditional mutual information which can be interpreted and evaluated.
Theorem 4.1.
Let A be a parametric algorithm with prediction function h .
Let U, Z n , W be defined as before.
Let U ′ be an identical independent version of U and W ′ =
A ( z n − 1 U ′ ) , then
[Equation]
One can alternatively upper bound loo - CMI through the convexity of the Kullback-Leibler divergence (similarly for floo - CMI ).
However, since the function ( − ln) is strictly convex, it is easy to show that the bound of Theorem 4.1 is strictly tighter through Jensen’s inequality.
Moreover, we opt to use this bound as it allows for more interpretable expressions (Corollary 4.1).
4.2 Geometry Aware Synthetic Randomization
Theorems 3.1 and 3.2 are valid in theory, but the outputs of algorithms used in practice are deterministic given a fixed input, and so we do not have a distribution of weights or predictions.
This means that (7) and (13) are degenerate.
Even if we vary the random seed for neural networks trained with SGD, and run the algorithm for each dataset and random seed, we obtain a set of discrete distributions with unequal supports.
We alleviate this issue by taking measures similar to the ones made in [12, 22, 15], and also previously in related contexts in [2, 9, 1, 23].
In particular, we add noise to the output of a deterministic algorithm, and use the now stochastic algorithm to bound the information in the weights and predictions.
We avoid adding isotropic Gaussian noise for the weight-based bounds, as changing the values of some weights may have little effect on the loss compared to others.
Therefore, we add noise while taking into account the geometry of the loss landscape.
Specifically, let A be deterministic algorithm, and let A Σ ( z n − 1 − U )
=
A ( z n − 1 − U ) + N , where N ∼
N ( 0 , Σ) .
A choice of Σ that would incorporate the notion that weights have a varying effect on the loss is Σ = diag( α 1 , . . .
, α K ) with α k >
for all k
and α i not necessarily equal to α j for i ̸
=
j .
Similarly, we can add noise to predictions made by a deterministic algorithm as h σ ( z n − 1 − U , x )
=
h ( z n − 1 − U , x ) + M , where M ∼ N ( 0 , σ 2 I ) .
The generalization bounds derived for
A Σ or h σ are not directly applicable to A and h , but if certain Lipschitz continuity assumptions hold for the loss function ℓ , it is possible to derive bounds for deterministic algorithms from the bounds of their noisy counterparts (Theorems 4.2 and 4.3).
We begin by applying Theorem 4.1 to get interpretable expressions for the conditional mutual information terms.
Corollary 4.1.
Let A be a deterministic algorithm with prediction function h .
Let w −
i
=
A ( z n − 1 − i ) and h − i =
h ( z n − 1 − i , x n ) be the output and predictions respectively when i th sample is removed from the training set.
Using Theorem 4.1, we obtain
[Equation]
4.3 Connections to Stability and Bounds For Deterministic Algorithms
It is easy to see the connection between classical definitions of stability and the right-hand side of (18), (19).
After all, to compute the right-hand side of (18) and (19), one only needs to know how the weights and predictions of the algorithm change when two datasets differ by one sample.
Based on the observation that components of the weight might have differing degrees of effect on the loss, and the right-hand side of (18), we introduce the idea of "relative" weight stability.
Moreover, we also use classical definitions of functional stability, and here we find it useful and intuitive to define two notions of stability: one with respect to the samples that the two datasets share, and one with respect to any other sample.
The definitions are as follows.
Definition 4.1.
Let z n − 1 , ˆ z n − 1 ∈
Z n be datasets such that z n − 1 and ˆ z n − 1 differ by at most one sample.
Letting a
=
A ( z n − 1 ) − A (ˆ z n − 1 ) , then we say that a deterministic algorithm A has ϵ weight stability relative to positive semi-definite matrix Σ if a T Σ − 1 a ≤ ϵ 2 .
Definition 4.2.
Let z n − 1 , ˆ z n − 1 ∈
Z n be datasets such that z n − 1 and ˆ z n − 1 differ by at most one sample, and without loss of generality, assume they differ at the first sample i.e. z 1 ̸
= ˆ z 1 and z k = ˆ z k for all k ̸ = 1 .
Let h be prediction function h : Z n − 1 × X → R d , then we say h has β -train stability if (cid:13)(cid:13) h
( z n − 1 , z k ) − h
(ˆ z n − 1 , z k ) (cid:13)(cid:13) 2 ≤ β 2 for all k ̸ = 1 .
Moreover, we say that h has β 1 -test stability if (cid:13)(cid:13) h ( z n − 1 , x ′ )
− h (ˆ z n − 1 , x ′ )
(cid:13)(cid:13) 2 ≤ β 21 for all x ′ ∈ X .
Given bounds on the noisy version of a deterministic algorithm, one can derive bounds on the deterministic algorithm itself by making certain Lipschitz continuity assumptions on the loss function.
Moreover, we show that given relative weight stability and functional stability of deterministic algorithm, we can add an optimal amount of noise to find a bound on the deterministic algorithm.
The following technique is also used by [16, 22] for the case of isotropic noise.
We generalize this result for arbitrary values of the noise covariance.
Theorem 4.2.
Let A : Z n − 1 → W be a deterministic algorithm, and let ℓ : W × Z
→ R + be the loss that a set of weights incur on a sample.
If ℓ ( w, · ) is L -Lipschitz in the weights, then if A has ϵ -weight stability relative to a positive semi-definite Σ , then gen( A ) ≤ (cid:113) 4 c n ϵL (cid:112) tr(Σ) .
Theorem 4.3.
Let h : Z n × X
→
R
d be a deterministic prediction function.
Assume that h has β -train stability and β 1 -test stability.
Assume the loss function ℓ : R ×
Y is L -Lipschitz continuity in the first coordinate, then gen( A ) ≤ (cid:113) 4 c n L (cid:112) nd ( nβ 2 + 2 β 21 ) .
4.4 Bounds For Stochastic Gradient Descent
As we have seen, the quality of our information bound depends on the stability of the optimization algorithm.
Stochastic Gradient Descent is the most commonly used method in deep learning, and it is therefore interesting to study the stability of SGD and how it translates into generalization from an information theoretic point of view.
Assuming the gradient updates are γ -bounded, ∥ w t +1 − w t ∥ ≤ √ γ for γ > 0 and all t > 0 , where w t are the weights at iteration t , then one can bound ∥ w − i − w − j ∥ with respect to the number of iterations of SGD [14], and obtain a bound on the generalization through the right-hand-side of (18).
In particular, we derive the following bound which show how training for a shorter time T and having more bounded updates γ both contribute to generalization.
The bound relates to the observation that early stopping (training for a smaller number of steps T) may improve the generalization of the network.
Lemma 4.1.
Let the gradient update rule be γ -bounded.
Suppose we run SGD for T steps, and we use the same starting value for the weights, then gen( A σ
[Equation]
4.5 Local interpretation of the bound
The bound in Corollary 4.1 depends on non-local quantities as it requires retraining the model a linear number of times on subsets of the data.
Ideally, one wants to bound the generalization without having to retrain.
We now provide a qualitative approximation of the bound using local quantities, in order to connect the bound with the geometry of the loss function.
Assume that the training algorithm minimizes the loss A ( z n ) =
arg min w ∈W ℓ ( z n , w ) , and let w ∗ =
A ( z n ) be the minimum to which the algorithm converges, and similarly let w ∗ −
i =
A ( z n − 1 − i ) .
Using influence functions [19], we can approximate:
[Equation]
where H w is the Hessian of the loss, ∇ w ℓ ( z i , w ∗ ) is the gradient of the loss on z i for w ∗ , and we have introduced the per-sample gradient at convergence rescaled by the Hessian g i =
1 n H − 1 w ∗ ∇ w ℓ ( z i , w ∗ ) .
Using this notation we can approximate:
[Equation]
where in ( ∗ ) we have assumed Σ −
= H w ∗ , which is the optimal variance of the noise (Section 4.2).
From (23) we see that converging to a flat minimum, that is, having a small norm of the Hessian H w ∗ , is expected to correlate to better generalization.
This is in accordance with several empirical results connecting flat minima to better generalization [18, 9, 6].
It should be noted that flatness by itself cannot explain generalization, since we can reparameterize the network to have identical predictions (and hence generalization) but arbitrarily large Hessian [7].
Indeed, in (23) we see that the role of the Hessian is mediated by the similarity of the (re-scaled) per-sample gradients at convergence, which can change under reparameterization.
In particular, studying flatness by itself is not enough.
Related Work
Over the past several years, there has been a significant line of research in using information theory both to interpret the behavior of DNNs [2, 28, 1] and to derive generalization bounds [27, 31, 5, 21, 29, 12, 16], with the earliest works on this line from [3, 27, 31].
Many of them develop information stability bounds which involve the mutual information between the output of the algorithm and the samples [27, 31, 5, 21] which could degenerate, e.g., if the data is continuous.
This was observed in [29], which then proposed a new bound based on the conditional mutual information (see Theorem 2.1) which is non-degenerate.
This work was further extended by [12, 16].
In particular [16] developed bounds based on prediction outputs rather than the algorithm outputs, and its combination with the conditional mutual information framework made it the state-of-the-art in terms of informationtheoretic generalization bounds.
The two issues identified in our work related to this is one of computability (as the [16] bound might in principle require computation exponential in the size of the dataset) and interpretability; this is the main focus of our work.
In particular, our loo-CMI framework can not only be computed more efficiently, but also connects to classical leave-one-out cross validation measures used extensively in practice (see Theorem 3.1).
Another aspect introduced in [16, 22] is the application of CMI bounds to deterministic algorithm outputs (as one can only run a finite number of runs of even a randomized algorithm on a dataset).
We use this viewpoint in our work, but using a geometry-aware method (see Theorem 4.2), which takes into account the loss landscape in the algorithmic output space.
There has also been a line of work connecting SGD to generalization including earlier works in [13] and more recently its connection with information theoretic bounds [25, 22].
We apply these ideas to the framework of loo-CMI (see Lemma 4.1).
The connection to classical notions of algorithmic stability can also be made using the loo-CMI framework (see Theorem 4.3).
Experiments
Model and datasets.
We now study the behavior of our loo - CMI and floo - CMI bounds on realworld image classification tasks.
In particular, we fine-tune an off-the-shelf ResNet-18 model pretrained on ImageNet on a set of standard image-classification tasks (see also Table 1): MIT- [26] and Oxford Pets [24] with a few thousand examples each.
On both datasets we fine-tune forepochs using stochastic gradient descent with learning rate η = 0 .
05 , momentum m = 0 .
99 , batch size B = 256 , and weight decay λ = 0 . 0005 .
We compute w − i /h
− i by removing sample i from the training set and re-training from scratch.
For all the experiments we remove 10 samples one at a time across 3 random seeds and use corollary 4.1 to compute the information bounds. We used 2 NVIDIA 1080Ti GPUS and the experiments take approximately 1-2 days.
[Figure]
[Figure] [Figure] [Figure]
[Figure]
the value of σ in (14) improves the floo - CMI bound (blue line) but also makes h σ increasingly different from h (as shown by the increased test error, orange line), thus making bounding the generalization of the original algorithm h given h σ more difficult.
Non-vacuous bounds.
In Table 1, we compute the loo - CMI and floo - CMI bounds using (10) and (14).
We show that while the former bound is vacuous, the floo - CMI bound provides non-vacuous generalization bounds on all datasets.
This is significant, as obtaining non-vacuous bound for largescale models remains a challenging problem.
The failure of the bound based on loo - CMI is expected here, since loo - CMI looks directly at the information contained in the weights of the network without considering how this information is used.
In large models with millions of parameters, most of the information in the weights do not significantly affect the predictions and should ideally be discarded.
This is indeed done by the floo - CMI bound.
Effect of the dataset size.
To show how the quality of the bound changes for different sizes of the dataset, we subsample randomly and without replacement n samples from each dataset.
In
Figure 2, we plot the resulting generalization gap alongside our generalization bounds as the size of the subsample varies (in terms of fractions of the dataset).
We observe that the generalization bound becomes tighter as the size of the training set grows.
This is not surprising as we expect the model to become more stable for larger datasets.
Synthetic randomization.
In Section 4.2, we introduced artificial Gaussian noise N (0 , Σ) in order to obtain better bounds.
Increasing the noise variance Σ in the loo - CMI bound, or σ in floo - CMI , improves the bound on the CMI, but also increases the test error of the model.
Hence, we need to select a value of Σ and σ that provides a good trade-off between the two.
In Figure 3, we show how changing the noise variance affects each term.
Conclusion
In this work, we introduced a bound based on a new information-theoretic measure, i.e., leave-one-out conditional mutual information and analyzed its properties.
loo-CMI bounds reduce the computational costs compared to prior CMI bounds, connects to several empirical quantities and generalization theories.
It leads to non-vacuous bounds for several deep learning applications in image classification.
There are several open questions including, whether it is possible to (i) obtain tighter bounds using the leave-one-out framework for general learning algorithms; (ii) further reduce the computational cost of computing loo-CMI 2 and (iii) obtain deeper understanding the effect of loss landscape and dynamics of learning algorithms on learning performance using the information-theoretic lens.
Acknowledgments and Disclosure of Funding
This work was supported in part by NSF grants #1955632, #2007714 and #2139304, ONR N0001422-1-2252, and ARO W911NF-17-1-0304.
The views and conclusions contained in this document are those of the authors and should not be interpreted as representing the official policies, either expressed or implied, of the Army Research Office or the U.S. Government.
The U.S. Government is authorized to reproduce and distribute reprints for Government purposes notwithstanding any copyright notation herein.
A Proofs A.1 Proof of Lemma
We begin by restating the result of Lemma 3.1.
Lemma.
Let ℓ : W × Z → [0 , 1] be a bounded loss function, then for all t > 0 , w ∈ W , and z n ∈ Z ,
[Equation]
Proof : To show this, we use tools from [30].
Let ℓ ∈ [0 , 1] n and t > 0 .
For u ∈ [ n ] , we define:
[Equation]
If we obtain a bound on E u ∼U [exp( t · f ( ℓ, u ))] for every ℓ ∈ [0 , 1] n , then we also obtain a bound on E u
∼U [exp ( t · (loo - cv( w, z n , u )))] for all z n ∈ Z n , w ∈ W .
Now,
[Equation]
where c := tc n = t (1 + 1 n − 1 ) .
Since the maximum of g must be at a stationary point, then
[Equation]
Setting (27) to 0 , then the right-hand side expression of (27) must be 0 .
This implies that ℓ j is either at the extremities or we must have a value of ln( b ) /c , where
[Equation]
Let m be the number of ℓ j that are 0 , and k the number of ℓ j that are 1 .
Solving for b using (28), we get that
[Equation]
We find that b depends on ℓ only through the number of 0 s and 1 s in ℓ .
Writing g ( ℓ ) with respect to m, k and b , we obtain
[Equation]
Hence, finding the maximum of g ( ℓ ) is equivalent to finding the solution for the following optimization problem.
[Equation]
Since m and k are bounded non-negative integers, we can find the solution of the above maximization through a brute-force search.
However, we can relax the constraints of (31), and allow m, k to take non-integer values over [0 , n ] .
Using the method of Lagrange multipliers, we have that:
[Equation]
The partial derivative of (32) with respect to m and k are given by:
[Equation]
ln
[Equation]
b ) n
[Equation]
Setting the partial derivatives of (33) and (34) to 0 yields
[Equation]
Substituting these expressions in (32), we see that we now need to find
[Equation]
Since 1 +
[Equation]
1 −
[Equation]
+
[Equation]
−
[Equation]
(cid:17) > 0 for all c >
[Equation]
this expression is maximized by the largest value k can take.
Hence, the we obtain our maximum for
[Equation]
Plugging into our expression, we have
[Equation]
Recalling that E u ∼U [exp( t · f ( ℓ, u ))]
= 1 /ng ( ℓ ) , we obtain
[Equation] [Equation]
Using the Taylor expansion of the exponent in the right-hand side of (38), we have that:
[Equation]
Finally, we obtain
[Equation] [Equation]
Recalling that c = tc n , we obtain the desired result.
A.2 Proof of Theorem 3.
We begin restating the Theorem.
Theorem.
Let A : Z n − 1 → W be a training algorithm.
Let ℓ : W × Z → [0 , 1] be a bounded loss function, then
[Equation]
We state a series of lemmas and definitions that are essential for the proof.
Lemma A.1.
Let z n ∈ Z n , and let w ∈ W , then
[Equation]
Proof.
We have that:
[Equation]
Lemma A.2.
Let A : Z n − 1 → W be a training algorithm
[Equation]
then
[Equation]
Proof.
[Equation]
Definition A.1.
We say that a random variable X ∼ P X is σ 2 − subgaussian if:
[Equation]
Lemma A.3 (Lemma 1, [31]) .
Let A, B be arbitrary random variables.
Let A ′ , B ′ be independent copies of A, B such that P A ′ ,B ′ =
P A ′ P B ′ .
Suppose f ( A ′ , B ′ ) is σ 2 -subgaussian, then
[Equation]
Now, fix Z n = z n
[Equation]
and let A
=
A ( z n − 1 − U ) and B = U
[Equation]
and let f ( A, B )
= loo - cv( z n , A, B ) .
[Equation]
Using Lemma A
[Equation]
1, E A ′ ,B ′ f
( A ′ , B ′ )
= 0 .
Moreover, we use f ( A ′ , B ′ ) is σ 2 -subgaussian with σ 2 = c 2 n / 4 using Lemma 3.1.
This gives us that:
[Equation]
where W =
A ( z n − 1 − U ) .
Taking the expectation over z n on both sides
[Equation]
we get that:
[Equation]
Since g ( · ) =
|·| is a convex function
[Equation]
then
[Equation]
A.3 Proof of Theorem 3.
We begin by restating the Theorem.
Theorem.
Let ℓ : R × Y → [0 , 1] be a bounded loss function.
Let U be a uniform random variable over [ n ] , then
[Equation]
This proof is nearly identical to the proof of Theorem 3.1.
Here, we use the prediction function instead of the weights.
We use the alternate definition of the loss ℓ : R × Y → [0 , 1] .
Let g ∈ R n be a set of prediction, then we also use an alternate definition of the leave-one-out cross validation where
[Equation]
Lemma A.4.
Let z n ∈ Z n , and let g ∈ R n , then
[Equation]
Lemma A.5.
Let h : Z n − 1 × X
→ R be a training algorithm, then
[Equation]
Now, we use Lemma A.3 with A = h ( z n − 1 − u , x n ) , B = U , and f ( A, B ) = loo - cv( z n , A, B ) .
[Equation] [Equation]
The proof then follows in the exact same way as it did for Theorem 3.1, and we get the desired result:
[Equation]
A.4 Proof of Theorem 4.
We begin by restating the Theorem.
Theorem.
Let A be a parametric algorithm with prediction function h .
Let U, Z n , W be defined as before.
Let U ′ be an identical independent copy of U , then
[Equation]
We begin by proving a more general result.
Let W ∼ P W , U ∼ P U be arbitrary random variables.
Suppose w ∈ W and u ∈ U .
Let q ( u ′ ) be any probability distribution over U , then we define
[Equation]
We begin with the following lemma.
Lemma A.6.
[Equation]
Proof.
We note that
[Equation]
where the last inequality follows from the concavity of ln and Jensen’s inequality.
Continuing from (74), we obtain
[Equation]
Since I ( W ; U ) =
E u ∼ P
U
[KL ( p ( w | u ) ∥∥ p ( w ))] , we obtain the result of the lemma.
To find a tight bound, we minimize Φ( q ( u ′ ) , u ) with respect to the probability distribution q ( u ′ ) .
Since ϕ ( q ( u ′ , u ) is an upper bound to KL ( p ( w | u ) ∥∥ p ( w )) for any q ( u ′ ) , the infimum of Φ with respect to q ( u ′ ) is still an upper bound.
Since we are trying to find the infimum with respect to a probability distribution, the Lagrangian is given by
[Equation]
We "differentiate" with respect to q ( u ′ ) , and obtain
[Equation]
Setting (81) to 0 , and solving for q ( u ′ ) for every u ′ ∈ U , we obtain
[Equation]
To find the expression for Φ( q ∗ ( u ′ )
[Equation]
u )
[Equation]
note that:
[Equation]
Plugging in the above into Φ( q ∗ ( u ′ ) , u ) , we obtain
[Equation]
Finally, taking the expectation with respect to U and using Lemma A.6, we obtain that
[Equation]
We apply the above to I ( W ; U | z n ) and I ( h ( z n − 1 − i , x n ) , U | z n ) to obtain the expressions of Theorem 4.1.
A.5 Proof of Theorem 4.
We begin a restatement of the theorem and the following lemma:
Theorem.
Let A : Z n − 1 → W be a deterministic algorithm, and let ℓ : W × Z
→ R + be the loss that a set of weights incur on a sample.
If ℓ ( w, · ) is L -Lipschitz in the weights, then if A has ϵ -weight stability relative to a positive semi-definite Σ , then gen( A ) ≤ (cid:113) 4 c n ϵL (cid:112) tr(Σ) .
Lemma A.7.
Let A : Z n − 1 → W be a deterministic algorithm Let the loss ℓ : W × Z
→ R + be
L-Lipschitz in the weights.
Let Σ be a positive definite matrix, then
[Equation]
Proof.
We have that:
[Equation]
where δ ′ =
ℓ ( A ( z n ) , z ′ ) − ℓ ( A Σ ( z n ) , z ′ ) and δ i = ℓ
( A ( z n ) , z i ) − ℓ ( A Σ ( z n ) , z i ) .
Recalling that
A Σ ( z n ) =
A ( z n ) + N where N ∼ N (0 , Σ) , we have that | δ ′
| ≤ L ∥ N ∥ (also for | δ i | ).
This gives us that:
[Equation]
where the last inequality follows from the fact that for N ∼ N (0 , Σ) , E [ ∥ N ∥ ] ≤ tr(Σ) .
Now, suppose A has ϵ -weight stability relative to Σ , and let α > 0 , then using (10) and (18), we have
[Equation]
Hence, for all α > 0 , we have that:
[Equation] [Equation]
Finding the minimum of the above with respect to α , we obtain:
[Equation] [Equation]
[Equation]
A.6 Proof of Theorem 4.
We begin with a restatement of the theorem, and the following lemma where proof is identical to the lemma used in Section A.5.
Theorem.
Let h : Z n × X
→
R
d be a deterministic prediction function.
Assume that h has β -train stability and β 1 -test stability.
Assume the loss function ℓ : R ×
Y is L -Lipschitz continuity in the first coordinate, then gen( A ) ≤ (cid:113) 4 c n L (cid:112) nd ( nβ 2 + 2 β 2 1 ) .
Lemma A.8.
Let h : Z n − 1 × X
→ R ⊂ R d be a deterministic prediction function.
Let the loss ℓ : R × Y
→ R + be L-Lipschitz in the predictions.
Let σ 2 be a positive definite matrix, then
[Equation]
Now, for a prediction function h with β -train stability and β 1 -test stability, we have that
[Equation]
Hence, we have that
[Equation]
Optimizing the right-hand-side of the above over σ , we obtain the desired result:
[Equation]
A.7 Proof of Lemma 4.
We restate the lemma.
Lemma.
Let the gradient update rule be γ -bounded.
Suppose we run SGD for T steps, and we use the same starting value for the weights, then gen( A σ 2 I ) ≤ (cid:113) c
2 n T 2 γ σ 2 .
Definition A.2 (Definition 2.4, [14]) .
An update step G is said to γ -bounded if for all w ∈ W ,
[Equation] [Equation]
Lemma A.9 (Lemma 2.5, [14]) .
Let G 1 , . .
, G T and G ′ 1 , . . .
, G ′ T be two arbitrary sequences of updates.
Let w 0 = w ′ 0 (same initial weights), and let δ t =
∥
w t − w ′ t ∥ where w t and w ′ t are defined recursively through w t
+1 =
G t ( w t ) and w ′ t +1 =
G ′ t ( w ′ t ) .
If G t and G ′ t are γ -bounded, then δ t +1 ≤ δ t +
[Equation]
√ γ.
Let G 1 , . .
, G T be the sequence of updates made when SGD is used with the dataset z n − 1 − i , and let G ′ 1 , . . .
, G ′ T be the sequence of updates made when SGD is used with the dataset z n − 1 − j .
If the previous updates are γ -bounded, then Lemma A.9 gives us a bound on ∥ w − i − w − j ∥ if we use T iteration of SGD, and we obtain the desired result.
B Experimental Details and Additional Experiments
We fine-tune an ImageNet pre-trained ResNet-18 model on MIT67 and Oxford Pets datasets.
In order to compute the information bounds we use the results in corollary 4.1.
We compute w − i /h
− i by removing sample i from the training set and re-training from scratch.
For all the experiments we remove 10 samples one at a time across 3 random seeds and use corollary 4.1 to compute the information bounds.
We used 2 NVIDIA 1080Ti GPUs and the experiments take 1-2 days.
We also include the results of fine-tuning on the Caltech256 dataset [10] (with ≈ 30000 samples).
C Societal Impact and Limitations C.1 Societal Impact
Our work is of a theoretical nature, and so there is no direct social impact.
However applications of deep neural networks, if not done properly, could have several consequence in society, including in terms of fairness, ability for misinformation etc.
However, theoretical work that attempts to understand & predict the performance of deep neural networks could give tools to ameliorate these issues.
C.2 Limitations and Future Work.
There are several open questions that remain from this work, pointing to some of its limitations.
Natural ones include tighter bounds on the generalization gap, obtaining impossibility results in terms of error performance among several others.