Integral Probability Metrics PAC-Bayes Bounds
Abstract
We present a PAC-Bayes-style generalization bound which enables the replacement of the KL-divergence with a variety of Integral Probability Metrics (IPM).
We provide instances of this bound with the IPM being the total variation metric and the Wasserstein distance .
A notable feature of the obtained bounds is that they naturally interpolate between classical uniform convergence bounds in the worst case (when the prior and posterior are far away from each other), and improved bounds in favorable cases (when the posterior and prior are close).
This illustrates the possibility of reinforcing classical generalization bounds with algorithm- and data-dependent components, thus making them more suitable to analyze algorithms that use a large hypothesis space.
Introduction and Related Work
Classical statistical learning theory is based on a worst-case perspective which can be too pessimistic to model practical machine learning.
In reality, data is rarely worst-case and experiments demonstrate learning tasks that are solved with much less data than predicted by traditional theory.
A primary manifestation of the traditional worst-case perspective is demonstrated by uniform convergence (UC); a genre of generalization bounds which form the backbone of the classical theory (Vapnik, 1999).
These bounds guarantee that the generalization gap of all hypotheses in the output-space of the algorithm simultaneously vanish as the training-set size grows.
The key algorithmic insight these bounds provide is summarized by the Empirical Risk Minimization principle (ERM) which asserts that it suffices to output any hypothesis in the class which minimizes the empirical risk.
Consequently, UC arguments provide non-trivial guarantees only if the hypothesis class used by the algorithm is restricted (e.g. has low Rademacher complexity or bounded VC dimension).
In contrast, practical learning approaches such as deep learning algorithms use huge hypothesis classes whose VC dimensions rapidly increase with the size and depth of the underlying network.
Hence, the rate guaranteed by UC arguments is often much slower than the rate observed in practice (Bachmann, Moosavi-Dezfooli, & Hofmann, 2021; Nagarajan & Kolter, 2019b; Neyshabur, Bhojanapalli,
McAllester, & Srebro, 2017; C. Zhang, Bengio, Hardt, Recht, & Vinyals, 2017).
A further shortcoming of UC bounds, and the associated ERM principle, is that they are algorithm- and data-independent; 1 that is, they do not utilize beneficial properties of the data and/or the algorithm.
For example, in practice, regularized algorithms often perform better than Empirical Risk Minimization, but this cannot be expressed by UC bounds and the ERM principle.
The PAC-Bayes (PB) framework is a prominent example of a theoretical framework that does not require the UC property.
This framework was pioneered by Shawe-Taylor and Williamson (1997) and McAllester (1998) and developed in later papers, e.g. (Catoni, 2007; Lever, Laviolette, & ShaweTaylor, 2013; Maurer, 2004; McAllester, 2003; Seeger, 2002); see Guedj (2019) and Alquier (2021) for extensive surveys.
Begin, Germain, Laviolette, and Roy (2016) introduced a general strategy that produces PB bounds from change-of-measure inequalities leading to bounds based on the Rényi’s α -divergence, and Alquier and Guedj (2018); Ohnishi and Honorio (2021); Picard-Weibel and Guedj (2022) further extended PB bounds to other Csiszár’s f -divergences.
PB theorems consider the generalization performance of stochastic predictors.
These bounds are nonuniform 2 by nature, and are algorithm and data-dependent.
They are usually based on a complexity term that depends on the Kullback-Leibler (KL) divergence between a data-dependent posterior distribution and a data-independent prior distribution.
There are additional notable works on data and algorithm-dependent guarantees.
The classical work of Bousquet and Elisseeff (2002) and Xu and Mannor (2012) studied generalization guarantees that depend on data and algorithm-dependent stability measures.
A further line of recent papers tries to incorporate noise robustness/resilience.
In Miyaguchi (2019), a PAC-Bayes transportation bound is used to measure the contribution of randomization to PB.
This is done via optimal transport and
Lipschitzness, based on the usual KL-PB bound.
The work of Wei and Ma (2019) uses data-dependent
Lipschitz smoothness to improve margin bounds, and Nagarajan and Kolter (2019a) passes from standard PB to a deterministic bound by assuming noise-resilience on the training data.
This property translates to the test data, implying that good training smoothness leads to good test smoothness.
Finally, Yang, Sun, and Roy (2019) measure data-dependent smoothness around each hypothesis (for each sample) and merge Catoni’s bound (Catoni, 2007) with Rademacher theory, to obtain fast rates.
Recent work by Aminian, Bu, Wornell, and Rodrigues (2022); Lopez and Jog (2018); Rodríguez Gálvez, Bassi, Thobaben, and Skoglund (2021); Wang, Diaz, Santos Filho, and Calmon (2019); J. Zhang, Liu, and Tao (2021) and Neu and Lugosi (2022) proved information-theoretic bounds on the expected generalization gap using the Wasserstein and the total-variation (TV) distances.
Our work is within the PB framework, and therefore enjoys the following advantages: (i)
The bounds are “in high probability” over the sample rather than in expectation.
(ii) PB bounds are sample dependent, i.e., bound the generalization gap for a specific sample-dependent posterior, while information-theoretic bounds are formulated as expectation over all sample sets, thereby providing a basis for empirical algorithms, e.g., Alquier (2021); Dziugaite and Roy (2017).
(iii)
The reference measure in PB can be any sample-independent distribution, while information-theoretic bounds consider a specific reference.
Our work introduces uniform convergence assumptions, while the above-mentioned papers each used different assumptions.
Recently, Chee and Loustau (2021) proposed PB bounds with the entropy regularized optimal transport distance for an online-learning setting with a finite class.
The optimal transport interpretation of the Wasserstein distance has been used recently in other contexts to derive generalization bounds.
Chuang, Mroueh, Greenewald, Torralba, and Jegelka (2021) proposed a bound that uses a data-dependent complexity measure, evaluated via the Wasserstein distance of independently sampled subsets of the training data in the feature space.
Hou, Kassraie,
Kratsios, Rothfuss, and Krause (2022) used the principles of optimal transport to derive an instancebased bound based on the local Lipschitz regularity of the learned prediction function in the data space.
In the modern deep learning regime, measures of the hypothesis class complexity used in UC bounds, such as the VC dimension or Rademacher complexity, are enormous, making the bounds extremely loose for any reasonable number of samples, as opposed to PB bounds (Dziugaite & Roy, 2017;
Jiang, Neyshabur, Mobahi, Krishnan, & Bengio, 2019).
However, these complexity measures often have closed-form formulas for models such as neural networks, which show explicitly the effect of the model architecture (number of layer, activation functions etc.).
This in contrast to PB bounds, in which the dependence on the hypothesis class is less explicit (but see Anthony and Bartlett (1999); Neyshabur, Tomioka, and Srebro (2015) for exceptions for neural networks).
Therefore, we believe that extension of UC bounds to incorporate data- and algorithm-dependence can facilitate the design of better performing architectures.
A further advantage of PB bounds is their non-uniformity (the generalization gap bound depends on the learning output), hence we can use the bound as a minimization objective for a structural minimization algorithm, where the complexity term acts as a regularizer.
In cases where the hypothesis class is very large, but we have some prior knowledge on which hypotheses are more likely to have low population loss (e.g. prefer simpler hypothesis as suggested by Occam’s Razor), then in PB one can inject this knowledge as the prior distribution, effectively lowering the generalization bound.
Can the rich theory of UC bounds be extended to help explain generalization with modern large scale models?
Can this theory be used to prove data and algorithm dependent guarantees?
In this paper, we take a step in the direction of answering these questions positively.
To achieve this goal, we show a new technique to incorporate UC bounds within the PAC-Bayes framework.
We prove new PB bounds with Integral Probability Metric (IPM) (Müller, 1997) to measure distances between distributions, rather than the standard KL or f -divergences used so far.
Specifically, we focus on utilizing two specific IPMs: the total variation and Wasserstein metrics This IPM framework allows greater flexibility, as it does not require the support of the posterior to be a subset of the prior’s support (absolute continuity) as in standard KL-PB bounds, and it applies to deterministic as well as stochastic prior and posterior distributions.
In fact, the IPM-based PB bounds we introduce match, at worst, the rate of the UC bound used.
Recently, Livni and Moran (2020) showed that the classical
KL-PB theorem cannot imply meaningful distribution-free generalization bounds for 1-dimensional linear classification.
In contrast, our derived IPM-PB bounds do imply such bounds, because linear classifiers satisfy uniform convergence.
We note that the work of Audibert and Bousquet (2003, 2007) showed a different approach to utilizing the UC assumption to derive PAC-Bayes bounds.
Their work assumed a UC property to utilize the generic chaining technique, resulting in more refined, variance-sensitive, bounds.
In contrast to our work, their bound is not fully empirical, and the assumed UC bound is not used by the resulting bound.
The Total Variation PAC-Bayes (TVPB) bound (Thm. 6) applies in any setting where uniform convergence (UC)
(Def. 1) holds, an assumption that is satisfied by many natural learning problems.
For example, in binary problems, UC holds with a rate of O ( (cid:112) VC ( H ) / m ) , where VC ( H ) is the VC dimension (Vapnik, 1999).
As observed in practice, for large models of deep neural networks with very large VC dimensions the learning rate on natural datasets is often much faster than predicted
UC bounds.
To explain this gap, we must turn to data and algorithm-dependent bounds.
The TVPB improves the gap bound to be O (cid:16)(cid:112) VC ( H ) D TV ( Q,P ) / m (cid:17) , effectively multiplying the VC dimension by the total variation distance of the posterior from the prior.
Intuitively, simpler posteriors (closer to the prior assumed before observing the data) lead to a better generalization gap.
Compared to the vanilla KL-PB, the TV distance can be small even in cases where the KL distance can be very large, and in any case, the bound only improves over the original UC bound.
In addition, the TV-PB bounds incorporate properties of H via the VC dim.
We also explore settings where the generalization gap function exhibits a certain smoothness property, and show a PB bound with the Wasserstein metric (Thm. 11), We analyze this smoothness property and show an explicit Wasserstein-based bound in the finite hypothesis class setting and in a linear regression setting.
In the latter setting, we show that a standard UC bound can be improved by a factor of O ( (cid:112) W 1 ( Q, P )) , where W 1 ( Q, P ) is the 1 st order Wasserstein distance between posterior and prior over the unit-sphere.
We conduct a numerical simulation to demonstrate the improvement of the Wasserstein PB bound over the UC bound, and, in cases of narrow prior distributions, over the KL-PB bound.
The experiment also investigates the case of non-randomized predictors by setting the prior and posterior as Dirac delta measures, which f -divergence based PB bounds are unable to use.
Figure 1 illustrates graphically the organization of the claims in the paper. [Figure]
Preliminaries 2.1 The Learning Problem
We begin with a short description a standard supervised learning task.
Consider a domain Z , 4 a distribution D over Z , a hypothesis set H , and finally, a loss function ℓ : H × Z → [0 , 1] .
The tuple ( D , H , ℓ ) defines a learning problem: The learning algorithm receives as input a training set
S =
{ z i } mi =1 ∈
Z m sampled i.i.d from D and selects an hypothesis h ∈ H .
The performance of h is measured by the expected risk , L D ( h ) def = E z ∼D ℓ ( h, z ) .
While the expected risk is unavailable to algorithm, the empirical risk , ˆ L S ( h ) def = 1 m (cid:80) mi =1 ℓ ( h, z i ) , can be evaluated using the training data.
The generalization gap is defined by ∆ S ( h ) def =
L D ( h ) − ˆ L S ( h ) .
Several of our results will assume that the learning problem ( D , H , ℓ ) satisfies the uniform convergence property ; i.e. the existence of a uniform upper bound on the generalization gap which applies simultaneously to all hypotheses in H .
Definition 1 (Uniform convergence, (Vapnik & Chervonenkis, 2015)) .
The learning problem ( D , H , ℓ ) satisfies the uniform convergence property, if there exists a bound function u ( m, δ ′ ) >
0 s.t.
for any m ∈ N + , δ ∈ (0 , 1) we have
[Equation]
UC type bounds are a major part of the foundations of theoretical machine learning.
Unfortunately, they suffers from a few drawbacks.
First, currently known bounds tend to be extremely loose in many cases, most notably for deep networks.
Second, the setup does not provide a natural way to encode prior knowledge into the bounds, particularly when dealing with deep networks - the hypothesis set is usually rich enough to express all relevant functions, and the training algorithms that might utilize some prior knowledge are not themselves a part of UC based bounds.
Finally, UC bounds are usually not data-dependent, a property which is critical to explain the generalization of DNNs on real-world data (Nagarajan & Kolter, 2019b; C. Zhang et al., 2017).
2.2 PAC-Bayes Bounds
Let M ( H ) denote the set of all probability measures over H .
For any probability measure Q ∈
M ( H ) , we define the expected loss, empirical loss and generalization gap by
[Equation]
PAC-Bayes theory bounds the expected loss, simultaneously for all “posterior" (sample-dependent) probability measures Q ∈ M ( H ) , with high probability over the samples S ∼ D m , given any “prior" (sample-independent) probability measure P ∈ M ( H ) .
A key feature of most PAC-Bayes bounds is their dependence on the KL divergence between the two distributions P, Q , KL ( Q ∥ P ) def = (cid:82)
H ln
( d Q d P )d Q , where d Q d P is the Radon–Nikodym derivative of Q w.r.t.
P .
While KL is a natural measure of divergence between probability distributions, it restricts the applicability of the resulting bounds to cases where the support of Q is contained in the support of P .
The following bound was introduced by McAllester (1998).
Proposition 2 (Classical KL-PB Bound) .
For any prior P ∈ M ( H ) and δ ∈ (0 , 1) , with probability at least 1 − δ over the samples S ∼ D m , for all Q ∈ M ( H ) , we have
[Equation]
A Template for IPM PAC-Bayes Bounds 3.1 General IPM-PB Bound
Definition 3 (Integral Probability Metric) .
(Müller, 1997; Sriperumbudur, Fukumizu, Gretton,
Schölkopf, & Lanckriet, 2009, 2012)
The Integral Probability Metric (IPM) between two probability measures P and Q over H is defined as
[Equation]
where F is a set of real-valued bounded functions H → R .
By definition, IPM distance measures are symmetric and non-negative.
Note that the KL-divergence is not a special case of IPM, rather it belongs to the family of f -divergences, that intersect with IPM only at the Total-Variation (Sriperumbudur et al., 2009, 2012).
The following proposition assumes that for any fixed sample S ′ ∈ Z m , the function f S ′ ( h ) def = 2( m − 1)∆ 2 S
( h ) is a member of a family of functions that depend on the sample, denoted F S ′ .
Thus, the IPM-PB bound allows us to ‘convert’ some knowledge we have about the properties of the generalization gap function ∆ S ( h ) to a generalization bound.
Since we do not specify yet the collection of function families {F S } S ∈Z m , the bound does not convey an explicit rate, and it should rather be seen as a template .
In the next sections, we will derive explicit bounds with specific IPMs divergences.
Namely, we will derive a total-variation distance based bound by selecting a collection of bounded function sets, and a Wasserstein distance based bound, by selecting a collection of Lipschitz function sets.
Proposition 4 (Template for IPM PB Bound) .
For any fixed dataset S ′ ∈
Z m , m ∈ N + , let F S ′ be a family of bounded and measurable functions H
→ R .
Assume that for any number of samples, m , and
sample S
∈
Z m , the function 2( m − 1)∆ 2 S ′ ( · ) is in F S ′ .
Then for any prior P ∈ M ( H ) and δ ∈ (0 , 1) , with probability at least 1 − δ over the samples S ∼ D m , for all Q ∈ M ( H ) , we have
[Equation]
The proof is in Appendix A.1.
The main idea is to use the IPM definition and the assumption as a change-of-measure inequality, instead of the variational formula by Donsker and Varadhan (1975), which is used in the classical KL-PB bound proof.
The rest of the proof is similar to the classical derivation (McAllester, 2003; Shalev-Shwartz & Ben-David, 2014).
Note that, similarly to the classical KL-PB bound, Proposition 4 does not require the UC property to hold.
However, in the next sections we will see that assuming an existence of a UC bound u ( m, δ ) , and selecting a particular collection of function families {F S } S ∈Z m , will result in explicit bounds that can improve upon the worst-case nature of the original u ( m, δ ) bound.
3.2 Template for Seeger Type IPM PAC-Bayes Bound
The work of Seeger (2002) and Maurer (2004) presented a different form of the PAC-Bayes theorem with fast O ( 1 / m ) rate as the dominant term, if the empirical risk is low.
We denote by kl ( p ∥ q ) the KL divergence between two Bernoulli distributions B ( p ) and B ( q ) , p, q ∈ [0 , 1] , that is,
[Equation]
For any h ∈ H , define the relative entropy of the empirical risk with respect to the expected one as ∆ kl S ( h ) def =
kl (ˆ L S ( h ) ∥
L D ( h )) .
Similarly, and overloading notations, we define for any distribution
Q ∈ M ( H ) , ∆ kl S ( Q ) def =
kl
(ˆ L S ( Q ) ∥
L D ( Q )) .
By replacing the KL-based change-of-measure inequality step in the proof of (Maurer, 2004, Thm.
5) with the IPM property (Def. 3) we get a similar bound for IPM measures (see a detailed proof in
Appendix A.2).
Proposition 5 (Template for Seeger Type IPM PAC-Bayes Bound) .
Assume f S ( h ) def = m · ∆ kl S ( h ) ∈
F S .
Then for any prior P ∈ M ( H ) and δ ∈ (0 , 1) , with probability at least 1 − δ over the samples
S ∼ D m , for all Q ∈ M ( H ) , we have [Equation]
[Equation]
By applying the Refined Pinsker’s relaxation (McAllester (2003), Eq. 6) we can immediately derive the following, looser, but easier to interpret bound
[Equation]
When ˆ L S ( Q ) is small (as is typical with modern deep networks), the final term determines the convergence rate.
We defer the investigation of PB bounds derived from Prop. 5 to Appendix B.
In the following sections we focus on investigating the implication of the Template IPM PB Bound of
Prop. 4.
Total-Variation PAC-Bayes Bounds
In this section, we investigate a PB bound with the total-variation (TV) distance, D TV ( Q, P ) def = sup A ∈ Σ H | P ( A ) − Q ( A ) | , where Σ H is the standard Borel sigma-algebra associated with H .
The
TV distance can be described as an IPM with the family of functions
[Equation]
for any M ≥ 0 .
To see this, note that
[Equation]
where equality ( i) holds since in the supremum it suffices to take the class of indicator functions { M · 1 A ( h ) , A ∈ Σ H } , since the functions in F ∞ M are bounded in [0 , M ] .
Theorem 6 (Template for Total-Variation PB Bound) .
Assume that there exists some uniform convergence bound u ( m, δ ′ )
(Definition 1), then, for any prior P ∈ M ( H ) and δ ∈ (0 , 1) , with probability of at least 1 − δ over samples S ∼ D m , for all Q ∈ M ( H ) , we have ∆ S ( Q ) ≤ (cid:115)
[Equation]
The proof (Appendix A.3) follows directly from the general IPM PB bound (Prop. 4) and the uniform convergence assumption, using a union bound argument.
This bound can be seen as a template to be used to derive explicit PB bounds, by plugging in existing UC bounds.
Note that while we require the existence of UC bound, the resulting bound is nonuniform (since it depends on the data-dependent posterior).
Compared to the original UC bound, u ( m, δ ) , the bound in (11) is roughly multiplied by a factor of (cid:112)
D TV ( Q, P ) ∈
[0 , 1] , ensuring tighter guarantees, especially if the posterior is close to the prior.
For example, consider a binary classification case, with the zero-one loss function and VC ( H ) class
H .
The well-known UC theorem states that the generalization gap converges uniformly at a rate
O (cid:16)(cid:112) VC ( H ) / m (cid:17) .
Proposition 7 (VC Bound, Boucheron, Bousquet, and Lugosi (2005)) .
There exists some universal constant 5 c > 0 s.t.
for any δ ∈ (0 , 1) we have
[Equation]
Using Thm.
6, we derive the following algorithm and data-dependent bound.
Corollary 8 (Total-Variation PB Bound for VC Classes) .
Consider a binary classification problem, with the zero-one loss, and hypothesis class H , with finite VC dimension, VC ( H ) .
There exists some universal constant c > 0 s.t. for any prior P ∈ M ( H ) and δ ∈ (0 , 1) , with probability of at least 1 − δ over samples S ∼ D m , for all Q ∈ M ( H ) , we have
[Equation]
Compared to the UC bound of Prop. 7, Cor. 8 multiplies the dominant term of the bound by a nonuniform (data and algorithm-dependent) factor of (cid:112) D TV ( Q, P ) , which is guaranteed to tighten the bound.
Note that the total-variation based bound of Aminian et al.
and Rodríguez Gálvez et al.
(2021) assume Lipschitz loss function, while our TV bound allows non continuous loss functions such as the zero-one loss.
The TV based bound of Wang et al.
(Thm. 1) is not directly comparable, since the empirical risk term is multiplied by a factor that goes to infinity for TV distance that goes to 1.
Wasserstein PAC-Bayes Bounds 5.1 Template for Wasserstein-PB Bound
In this section, we provide a PAC-Bayes generalization bound with the Wasserstein metric between posterior and prior and a certain smoothness parameter of the generalization gap function.
We explore learning settings for which H can be paired with a distance metric ρ : H × H
→ R ≥ 0 s.t.
( H , ρ ) is a
Polish metric space (complete, separable metric space).
Given the distance metric ρ , we can define the Wasserstein distance between any two probability measures on H .
Definition 9 (Wasserstein Distance) .
For any two probability measures P, Q on H with finite first moment, the 1 st order Wasserstein distance is
[Equation]
(14) where Γ( Q, P ) denotes the set of all couplings of Q and P , that is, the set of all joint measure on
H × H whose marginals are Q and P .
The following proposition gives a dual representation for the first-order Wasserstein distance.
Proposition 10 (Kantorovich-Rubinstein Duality (Villani, 2006)) .
For any 0 ≤ K , and any two probability measures P, Q ∈ M ( H ) ,
[Equation]
where F Lip K is the set of K -Lipschitz functions w.r.t.
ρ ( h, h ′ ) , i.e. functions that satisfy
[Equation]
We can write the Kantorovich-Rubinstein duality (15) using IPM formulation (Def. 3),
[Equation]
Using this duality we will prove the following bound.
Theorem 11 (Template for Wasserstein-PB Bound) .
Assume that for any δ ′ ∈ (0 , 1] , w.p. at least 1 − δ ′ over the sampling S ∼ D m , the squared generalization gap function, ∆ 2 S ( · ) , is K -Lipschitz w.r.t.
the metric ρ with some K = K ( m, δ ′ ) .
Then, for any prior P ∈ M ( H ) and δ ∈ (0 , 1) , with probability of at least 1 − δ over samples S ∼ D m , for all Q ∈ M ( H ) , we have
[Equation]
The proof (Appendix A.4) follows directly from Proposition 4 and the assumption, via the union bound.
Theorem 11 can be seen as a template to be used for deriving Wasserstein-PB bounds in various learning settings where the ∆ 2 S ( · ) is K -Lipschitz with high probability (over samples), where the rate of K = K ( m, δ / 2 ) should be O
( 1 / m ) to ensure a factor O ( 1 / √ m ) multiplying the divergence between posterior and prior, as in the KL-PB bound.
Such a result can be challenging to prove since it requires uniform convergence of the slope between any two hypotheses in H .
Next, we will show specific learning settings where this property holds and the resulting generalization bounds.
We note that recent work Neu and Lugosi (2022) also establishes a Wasserstein-based informationtheoretic generalization bound.
This work assumed infinitely smooth loss functions and established bounds on the expected generalization gap, rather than high-probability bounds.
In fact, Neu and
Lugosi (2022) noted that obtaining such bounds is an open problem.
Observe, though, that our bound depends on the Lipschitz constant K = K ( m, δ ) which needs to assessed; see sections 5.2 and 5. for specific examples.
The general problem remains open.
At a more pragmatic level, we note that learning algorithms derived from minimizing Wasserstein based PB bounds have an added benefit of more stable optimization compared to KL based approaches, due to lower gradient variance, as noted in Arjovsky, Chintala, and Bottou (2017), and this distance measure can be approximated efficiently from finite samples (Cuturi, 2013; Weed & Bach, 2017)
5.2 Wasserstein-PB Bound for Finite Classes
We first investigate the simple case of a finite hypothesis class with a loss function ℓ which is GLipschitz w.r.t.
the metric ρ .
Note that for finite classes, UC always holds.
We derive the following bound from Thm.
Theorem 12 (Wasserstein-PB Bound for Finite Classes) .
Let H be a finite hypothesis class.
Assume that for any fixed z ∈ Z , ℓ ( h, z ) is a G -Lipschitz function in h w.r.t the metric ρ .
Then for any prior P ∈ M ( H ) and δ ∈ (0 , 1) , with probability of at least 1 − δ over samples S ∼ D m , for all
Q ∈ M ( H ) , we have
[Equation]
The proof (Appendix A.5) makes use of standard union bound arguments, Hoeffding’s concentration inequality and the template Wasserstein-PB bound (Thm. 11).
Notice that compared to the standard
UC bound for finite classes, Thm.
12 multiplies the bound by a nonuniform factor of (cid:112) GW 1 ( Q, P ) .
5.3 Wasserstein-PB Bound for Loss-Gradient UC Classes
In this section we show a Wasserstein-PB bound for learning problems ( D , H , ℓ ) , with H ⊂ R d , for some dimension d ∈ N + , that satisfy the standard UC property, and, additionally, satisfy UC property for the loss gradient, as defined below.
Definition 13 (Loss-Gradient UC Property) .
A learning problem ( D , H , ℓ ) , with H ⊂ R d , is said to satisfy the loss-gradient UC property , if: (i) the loss function ℓ ( h, z ) is differentiable w.r.t.
h on
Int ( H ) ×
Z and continuous w.r.t.
h on H × Z .
(ii)
The problem satisfies the uniform convergence property (Def. 1).
(iii)
The empirical average of the loss gradient converges uniformly in L 2 norm sense to its mean.
I.e., there exists a bound function u grad ( m, δ ) > 0 , s.t. for any δ ∈ (0 , 1) we have u grad ( m, δ )
→ m →∞ 0 and for any m ∈ N + ,
[Equation]
where ∇ h ℓ ( h, z ) denotes the gradient of ∇ h ℓ ( h, z ) w.r.t.
h , for a fixed z ∈ Z , and Int ( H ) is the interior of H .
We call u grad , the UC bound of the loss gradient.
Theorem 14 (Wasserstein-PB Bound for Loss-Gradient UC Classes) .
Let ( H , ρ ) be a metric space such that H ⊂ R
d is a closed and convex set, and ρ is the L 2 distance.
Assume the learning problem satisfies the loss-gradient UC property (Def. 13), with UC bound u , and a UC bound of the loss gradient, u grad .
Then for any prior P ∈ M ( H ) and δ ∈ (0 , 1) , with probability of at least 1 − δ over samples S ∼ D m , for all Q ∈ M ( H ) , we have
[Equation]
The proof (Appendix A.6) is derived from the assumptions and the template Wasserstein-PB bound (Thm. 11).
In learning problems that satisfy the loss-gradient UC property, we often have u ( m, δ / 4 ) , u grad ( m, δ / 4 ) ∈
O ( 1 / √ m ) , and then the resulting bound is O ( 1 / √ m ) .
We provide full analysis that shows such a rate for the following linear regression example.
5.4 Linear Regression Example
Based on the Wasserstein-PB Bound for Loss-Gradient UC Classes (Thm. 14), we derive the following corollary.
Corollary 15 (Wasserstein-PB Bound for Linear Regression) .
Consider a data distribution of pairs
z
=
( x, y ) , where x is sampled from an unknown distribution supported on a d -dimensional ball of radius r > 0 , B dr def = { x ∈ R d : ∥ x ∥ 2 ≤ r } , and the target, y = f ( x ) , is set by an unknown, possibly random, target function f : B dr → [ − 1 , 1] .
The hypothesis space H is B d 1 / r , and the loss function is ℓ ( x, y, h )
= 1 4 ( h ⊤ x − y ) 2 .
Then, for any prior P ∈ M ( H ) and δ ∈ (0 , 1) , with probability at least 1 − δ over samples S ∼ D m , for all Q ∈ M ( H ) ,
[Equation]
where W 1 ( Q, P ) denotes the 1 st order Wasserstein distance with the L 2 metric,
[Equation]
The full expression of the bound appears in the theorem’s proof (Appendix A.7).
Ignoring logarithmic factors we obtained a UC bound of ˜ O (cid:18)(cid:113) dm (cid:19) , and a Wasserstein-PB bound of ˜ O (cid:18)(cid:113) rW 1 ( Q, P ) dm (cid:19) .
Note that from Thm.
6 we can also deduce a TV-PB bound of order ˜
O
(cid:18)(cid:113) D TV ( Q, P ) dm (cid:19) .
In comparison, the standard KL-PB bound is of order ˜ O (cid:18)(cid:113) KL ( Q ∥ P ) m (cid:19) .
The TV-PB bound is, at worst, roughly the same as the UC bound, since D TV ( Q, P ) ≤ 1 .
Note that since Q and P are distributions over a sphere of radius 1 / r , then rW 1 ( Q, P ) ≤ 2 .
Hence, the Wasserstein-PB is also, at worst, roughly the same as the UC bound.
However, the KL-PB bound can be either tighter or looser, depending on P and Q .
In cases where the mass of the posterior Q is concentrated in a region of the hypothesis set where the prior P is arbitrarily small, then the KL divergence can be arbitrarily large, making the KL-PB bound extremely loose compared to the UC, TV-PB, and
Wasserstein-PB bounds.
The numerical experiment described in Appendix C demonstrates this by investigating different prior distributions with different widths.
In particular, for posteriors and priors that are Dirac delta distributions (i.e., deterministic predictors), we show that the Wasserstein-PB considerably improves over the UC bound, while the KL-PB bound is undefined.
We therefore demonstrated non-vacuous guarantees for deterministic models within the PB framework without requiring additional derandomization steps.
We can compare the bound of Thm.
14 to the bound of Corollary 8, in Neu and Lugosi (2022), which is also dependent on the Wasserstein distance between a data-dependent output (posterior) and a base measure (prior).
The bound of Thm.
14 is different in the sense that (i) It holds with high probability instead of in expectation.
(ii) Instead of assuming infinitely-smooth loss function with β -bounded directional derivatives, Thm.
14 assumes a loss-gradient UC class.
(iii)
The bound scales as ˜ O (cid:16)(cid:112) W 1 · u · u grad (cid:17) instead of ˜ O
(cid:18)(cid:113) W 2 · dβm (cid:19) .
In our linear regression example, Cor. 15 scales as ˜ O
(cid:16)(cid:112) W 1 · u · u grad (cid:17) =
˜ O (cid:32)(cid:114) r (cid:113) dm r (cid:113) dm (cid:33) =
˜ O (cid:18) r (cid:113) dm (cid:19) .
The loss is infinitely-smooth with β ∈ O (1 + r + r 2 ) , and therefore Cor.
8 of Neu and Lugosi (2022) scales as ˜ O (cid:18)(cid:113) W 2 · dβm (cid:19) =
˜ O (cid:18) r (cid:113) dm (1 + r + r 2 ) (cid:19) , i.e., looser by a factor of ˜ O (cid:0) √ 1 + r + r 2 (cid:1) compared to Cor. 15.
Discussion
We have presented high-probability PB bounds based on integral probability metrics, that extend standard PB bounds based on KL divergence and
more recent f -divergence and α -divergence based bounds, to a new class of distances.
Our bounds interpolate between classic UC bounds and PB bounds, by allowing data- and algorithm-dependent complexity terms.
As in all PB results, our bounds suggest improved rates when the PB posterior is close to the prior.
While we have extended high-probability PB bounds for IPMs to novel distance measures, it is still an open question to do so without the UC assumption.
Possible directions for future research include: (i) Deriving high probability IPM-PB bounds (e.g.
Wasserstein or TV based), without global UC assumptions (possibly using localization based approaches, e.g. , Local Rademacher complexities (Bartlett, Bousquet, & Mendelson, 2005; Koltchinskii & Panchenko, 2004) are computed only on a subset of hypotheses with small empirical risk).
This may allow non-vacuous bounds for large-scale models where global UC-based bounds are extremely vacuous.
(ii) Derivation of algorithms that utilize the bounds as minimization objectives.
Based on the optimization advantages of Wasserstein based costs mentioned in Sec. 5.1, these could lead to enhanced practical utility.
Such an advantage could play an important role in meta-learning schemes where PB methods have been widely used in recent years (Amit & Meir, 2018).
Acknowledgments
We thank Nadav Merlis and Daniel Soudry for helpful discussions of this work, and the anonymous reviewers for their helpful comments.
SM is a Robert J. Shillman Fellow; he acknowledges support by
ERC grant 802599, by ISF grant 1225/20, by BSF grant 2018385, by an Azrieli Faculty Fellowship, by Israel PBC-VATAT, and by the Technion Center for Machine Learning and Intelligent Systems (MLIS).
The work of RM is partially supported by ISF grant 1693/22, by the Ollendorff Center of the
Viterbi ECE Faculty at the Technion and by the Skillman chair in biomedical sciences.
Checklist
For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?
[Yes] (b) Did you describe the limitations of your work?
[Yes] (c) Did you discuss any potential negative societal impacts of your work?
[N/A] (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
[Yes] 2.
If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?
[Yes] (b) Did you include complete proofs of all theoretical results?
[Yes] 3.
If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[Yes] (b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
[Yes] (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[Yes] (d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
[N/A] 4.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
[N/A] (b) Did you mention the license of the assets?
[N/A] (c) Did you include any new assets either in the supplemental material or as a URL?
[N/A] (d) Did you discuss whether
and how consent was obtained from people whose data you’re using/curating?
[N/A] (e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content?
[N/A] 5.
If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?
[N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
[N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
[-In the extreme case of σ P = 
-]
[N/A]
[-In particular, for the function f ( z ) def = (cid:13)(cid:13)(cid:13) µ 2 + Σ 1 / 2 2 z (cid:13)(cid:13)(cid:13) 
-]