Tighter Expected Generalization Error Bounds via Wasserstein Distance

1 Introduction . A learning algorithm is a mechanism that takes a dataset s = ( z1 , . . . , zn ) of n samples zi ∈ Z taken i.i.d . from a distribution PZ as an input , and produces a hypothesis w ∈ W by means of the conditional probability distribution PW |S . The ability of a hypothesis w to characterize a sample z is described by the loss function ` ( w , z ) ∈ R. More precisely , a hypothesis w describes well the samples from a population PZ when its population risk , i.e. , LPZ ( w ) , E [ ` ( w , Z ) ] , is low . However , the distribution PZ is often not available and the empirical risk on the dataset s , i.e. , Ls ( w ) , 1n ∑n i=1 ` ( w , zi ) , is considered as a proxy . Therefore , it is of interest to study the discrepancy between the population and empirical risks , which is defined as the generalization error : gen ( w , s ) , LPZ ( w ) −Ls ( w ) . Classical approaches bound the generalization error in expectation and in probability ( PAC Bayes ) either by studying the complexity and the geometry of the hypothesis ’ space W or by exploring properties of the learning algorithm itself ; see , e.g. , [ 2 , 3 ] for an overview . More recently , the relationship ( or amount of information ) between the generated hypothesis and the training dataset has been used as an indicator of the generalization performance . In [ 4 ] , based on 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . [ 5 ] , it is shown that the expected generalization error , i.e. , gen ( W , S ) , E [ gen ( W , S ) ] , is bounded from above by a function that depends on the mutual information between the hypothesis W and the dataset S with which it is trained , i.e. , I ( W ; S ) . However , this bound becomes vacuous when I ( W ; S ) → ∞ , which occurs for example when W and S are separately continuous and W is a deterministic function of S. To address this issue , it is shown in [ 6 ] that the generalization error is also bounded by a function on the dependency between the hypothesis and individual samples , I ( W ; Zi ) , which is usually finite due to the smoothing effect of marginalization . Following this line of work , in [ 7 ] , the authors present data-dependent bounds based on the relationship between the hypothesis and random subsets of the data , i.e. , DKL ( PW |s ‖ PW |sjc ) where j ⊆ [ n ] . After that , a more structured setting is introduced in [ 1 ] , studying instead the relationship between the hypothesis and the identity of the samples . The authors consider a super-sample of 2n i.i.d . instances z̃i from PZ , i.e. , s̃ = ( z̃1 , . . . , z̃2n ) . This super sample is used to construct the dataset s by choosing between the samples z̃i and z̃i+n using a Bernoulli random variable Ui with probability 12 , i.e. , zi = z̃i+uin . In this paper , the two settings are referred to as the standard and randomized-subsample settings.1 In the randomized-subsample setting , the empirical generalization error is defined as the difference between the empirical risk on the samples from s̃ not used to obtain the hypothesis , i.e. , s̄ = s̃ \ s , and the empirical risk on the dataset s , i.e. , ĝen ( w , s̃ , u ) , L̄s ( w ) −Ls ( w ) = 1 n ∑n i=1 ( ` ( w , z̃i+ ( 1−ui ) n ) − ` ( w , z̃i+uin ) ) , where u is the sequence of n i.i.d . Bernoulli trial outcomes ui . The expected value of the empirical and the ( standard ) generalization errors coincide , i.e. , E [ ĝen ( W , S̃ , U ) ] = gen ( W , S ) . Also , the expected generalization error is controlled by the conditional mutual information between the hypothesis W and the Bernoulli trials U , given the super sample S̃ [ 1 ] , i.e. , I ( W ; U |S̃ ) , by the individual conditional mutual information [ 9 ] , i.e. , I ( W ; Ui|Z̃i , Z̃i+n ) , and by the “ disintegrated ” mutual information with a subset UJ of the Bernoulli trials [ 10 ] , i.e. , DKL ( PW |S̃ , U ‖PW |S̃ , UJc ) . A highlight of this setting is that these conditional notions of information are always finite [ 1 ] and smaller than their “ unconditional ” counterparts [ 10 ] , e.g. , I ( W ; U |S̃ ) ≤ I ( W ; S ) and I ( W ; U |S̃ ) ≤ n log ( 2 ) . Some steps towards unifying these results are taken in [ 8 ] , where the authors develop a framework that makes it possible to recover the expected generalization error bounds based on the mutual information I ( W ; S ) and the conditional mutual information I ( W ; U |S̃ ) . Then , the aforementioned framework is further exploited in [ 9 ] to recover the single-sample and the random-subsets bounds , which are based on I ( W ; Zi ) , DKL ( PW |S ‖PW |SJc ) , andDKL ( PW |S̃ , U ‖PW |S̃ , UJc ) , and to generate new individual conditional mutual information bounds , i.e. , I ( W ; Ui|Z̃i , Z̃i+n ) . Finally , in [ 11 , 12 ] , other systematic ways to recover some of the said bounds and obtain similar new ones are studied . In parallel , there were some attempts to bridge the gap between employing the geometry and complexity of the hypothesis space and the relationship between the hypothesis and the training samples . In [ 13 ] , the authors bound gen ( W , S ) with a function of weighted dependencies between the dataset and increasingly finer quantizations of the hypothesis , i.e. , { 2−k/2I ( [ W ] k ; S ) } k , which can be finite even if I ( W ; S ) →∞ . This result stems from a clever usage of the chaining technique [ 14 , Theorem 5.24 ] , and a comparison with this kind of approaches is given in Appendix A . Later , in [ 15 ] and [ 16 ] , it is shown that the expected generalization error is bounded from above by a function of the Wasserstein distance between the hypothesis distribution after observing the dataset PW |S and its prior PW , i.e. , Wp ( PW |S , PW ) , and by a function of the Wasserstein distance between the hypothesis distribution after observing a single sample PW |Zi and its prior PW , i.e. , Wp ( PW |Zi , PW ) , which are finite when a suitable metric is chosen but are difficult to evaluate . Concurrently , in [ 17 ] it is shown that a similar result holds , if the metric is the Minkowski distance , for the distribution of the data PS and the backward channel PS|W , i.e. , Wpp , ‖·‖ ( PS|W , PS ) . The main contributions of this paper are the following : • It introduces new , tighter single letter and random-subset Wasserstein distance bounds for the standard and randomized-subsample settings ( Theorems 1 , 2 , 3 , and 4 ) . • It shows that when the loss is bounded and the geometry of the space is ignored , these bounds recover from below ( and thus are tighter than ) the current relative entropy and mutual 1In [ 8 ] the latter is called the random-subset setting . However , this may cause confusion with the randomsubset bounds in the present work . information bounds on both the standard and randomized-subsample settings . In fact , they are also tighter when the loss is additionally subgaussian or under certain milder conditions on the geometry . However , these results are deferred to Appendix B to expose the main ideas more clearly . Moreover , Corollaries 1 and 2 overcome the issue of potentially vacuous relative entropy bounds on the standard setting . • It introduces new bounds based on the backward channel , which are analogous to those based on the forward channel and more general than previous results in [ 17 ] . • It shows how to generate new bounds based on a variety of information measures , e.g. , the lautum information or several f -divergences like the Hellinger distance or the χ2-divergence , thus making the characterization of the generalization more flexible . 2 Preliminaries . 2.1 Notation . Random variables X are written in capital letters , their realizations x in lower-case letters , their set of outcomes X in calligraphic letters , and their Borel σ-algebras X in script-style letters . Moreover , the probability distribution of a random variable X is written as PX : X→ [ 0 , 1 ] . Hence , the random variable X or the probability distribution PX induce the probability space ( X , X , PX ) . When more than one random variable is considered , e.g. , X and Y , their joint distribution is written as PX , Y : X⊗ Y→ [ 0 , 1 ] and their product distribution as PX ⊗ PY : X⊗ Y→ [ 0 , 1 ] . Moreover , the conditional probability distribution of Y given X is written as PY |X : Y⊗X → [ 0 , 1 ] and defines a probability distribution PY |X=x ( or PY |x for brevity ) over Y for each element x ∈ X . Finally , there is an abuse of notation writing PX , Y = PY |X × PX since PX , Y ( B ) = ∫ ( ∫ χB ( ( x , y ) ) dPY |X=x ( y ) ) dPX ( x ) for all B ∈ X⊗Y , where χB is the characteristic function of the set B . The natural logarithm is log . 2.2 Necessary definitions , remarks , claims , and lemmas . Definition 1 . Let ρ : X × X → R+ be a metric . A space ( X , ρ ) is Polish if it is complete and separable . Throughout it is assumed that all Polish spaces ( X , ρ ) are equipped with the Borel σ-algebra X generated by ρ . When there is no ambiguity , both the metric space ( X , ρ ) and the generated measurable space ( X , X ) are written as X . Definition 2 . Let ( X , ρ ) be a Polish metric space and let p ∈ [ 1 , ∞ ) . Then , the Wasserstein distance of order p between two probability distributions P and Q on X is Wp ( P , Q ) , ( inf R∈Π ( P , Q ) ∫ X×X ρ ( x , y ) pdR ( x , y ) ) 1/p , where Π ( P , Q ) is the set of all couplings R of P and Q , i.e. , all joint distributions on X × X with marginals P and Q , that is , P ( B ) = R ( B , X ) and Q ( B ) = R ( X , B ) for all B ∈ X . Remark 1 . Hölder ’ s inequality implies that Wp ≤Wq for all p ≤ q [ 18 , Remark 6.6 ] . Hence , since this work is centered on upper bounds the focus is on W , W1 . Definition 3 . A function f : X → R is said to be L-Lipschitz under the metric ρ , or simply f ∈ L-Lip ( ρ ) , if |f ( x ) − f ( y ) | ≤ Lρ ( x , y ) for all x , y ∈ X . Lemma 1 ( Kantorovich-Rubinstein duality [ 18 , Remark 6.5 ] ) . Let P1 ( X ) be the space of probability distributions on X with a finite first moment . Then , for any two distributions P and Q in P1 ( X ) W ( P , Q ) = sup f∈1-Lip ( ρ ) { ∫ X f ( x ) dP ( x ) − ∫ X f ( x ) dQ ( x ) } . ( KR duality ) Definition 4 . The total variation between two probability distributions P and Q on X is TV ( P , Q ) , supA∈X { P ( A ) −Q ( A ) } . Definition 5 . The discrete metric is ρH ( x , y ) , 1 [ x 6= y ] , where 1 is the indicator function . Remark 2 . A bounded function f : X → [ a , b ] is ( b− a ) -Lipschitz under the discrete metric ρH . Remark 3 . The Wasserstein distance of order 1 is dominated by the total variation . For instance , if P and Q are two distributions on X then W ( P , Q ) ≤ dρ ( X ) TV ( P , Q ) , where dρ ( X ) is the diameter of X . In particular , when the discrete metric is considered W ( P , Q ) = TV ( P , Q ) [ 18 , Theorem 6.15 ] . Lemma 2 ( Pinsker ’ s and Bretagnolle–Huber ’ s ( BH ) inequalities ) . Let P and Q be two probability distributions on X and define Ψ ( x ) , √ min { x/2 , 1− exp ( −x ) } , then [ 19 , Theorem 6.5 ] and [ 20 , Proof of Lemma 2.1 ] state that TV ( P , Q ) ≤ Ψ ( DKL ( P ‖Q ) ) .

Summary: This paper considers the problem of obtaining generalization bounds using the statistical measure of dependence between the input and output of a learning algorithm. Tha main contribution of this paper is to provide a generalization bound using the Wasserstein distance between the posterior distribution and P(W). This result is further improved by considering the single sample case where the generalization bounds are based on the Wasserstein distance P(W|Z_i) and P(W). Also the authors extend their results to the setup of Steinke and Zakynthinou ‘20. The authors specifically follow the literature of information-theoretic generalization bounds, and compare their bounds to the existing bounds in the literature. In all cases, the authors show that the bounds using Wasserstein distance are tighter than the bound using mutual information (or KL).  Also, they show for a toy example of estimating the mean of gaussian distribution their bound provides “order-wise” improvement over the mutual information bounds where the mutual information based bounds gives upper bounds on the generalization that decays as 1/sqrt(n) but the Wasserstein distance based bound gives a bound that decays as 1/n.  The authors also use the joint-range technique to provide generalization bounds based on f-divergences.