M AX - SLICED B URES D ISTANCE FOR I NTERPRETING D ISCREPANCIES
Anonymous authors
Paper under double-blind review
A BSTRACT
We propose the max-sliced Bures distance, a lower bound on the max-sliced
Wasserstein-2 distance, to identify the instances associated with the maximum discrepancy between two samples.
The max-slicing can be decomposed into two asymmetric divergences each expressed in terms of an optimal slice or equivalently a ‘witness’ function that has large magnitude evaluations on a localized subset of instances in one distribution versus the other.
We show how witness functions can be used to detect and correct for covariate shift through reweighting and to evaluate generative adversarial networks.
Unlike heuristic algorithms for the max-sliced Wasserstein-2 distance that may fail to ﬁnd the optimal slice we detail a tractable algorithm that ﬁnds the global optimal slice and scales to large sample sizes.
As the Bures distance quantiﬁes differences in covariance, we generalize the max-sliced Bures distance by using non-linear mappings, enabling it to capture changes in higher-order statistics.
We explore two types of non-linear mappings: positive semideﬁnite kernels where the witness functions belong to a reproducing kernel Hilbert space, and task-relevant mappings corresponding to a neural network.
In the context of samples of natural images, our approach provides an interpretation of the Fr´echet Inception distance by identifying the synthetic and natural instances that are either over-represented or under-represented with respect to the other sample.
We apply the proposed measure to detect imbalances in class distributions in various data sets and to critique generative models.
I NTRODUCTION
Divergence measures quantify the dissimilarity between probability distributions.
They are fundamental to hypothesis testing and the estimation and criticism of statistical models, and serve as cost functions for optimizing generative adversarial neural networks (GANs).
Although a multitude of divergences exists, not all of them are interpretable.
A divergence is interpretable if can be expressed in terms of a real-valued witness function ω ( · ) whose level-sets identify the speciﬁc subsets that are not well matched between the distributions, speciﬁcally, subsets which have much higher or much lower probability under one distribution versus the other.
Localizing these discrepancies is useful for understanding and compensating for differences between two samples or distributions, to detect covariate shift (Shimodaira, 2000; Quionero-Candela et al., 2009; Lipton et al., 2018) or to evaluate generative models (Heusel et al., 2017).
While many divergences can be posed in terms of witness functions, not all witness functions are readily obtained or interpreted.
From an information-theoretic perspective, the most natural witness function is the logarithm of the ratio of the densities (Kullback & Leibler, 1951) as in the KullbackLeibler divergence.
Applying other convex functions to the density ratio constitutes the family of f -divergences (Ali & Silvey, 1966; R ´ enyi, 1961), which include the Hellinger, Jensen-Shannon, and others.
However, without a parametric model estimating the densities from samples is challenging (Vapnik, 2013).
Following Vapnik’s advice to “try to avoid solving a more general problem as an intermediate step,” previous work has sought to directly model the density ratio via kernel learning (Nguyen et al., 2008; Kanamori et al., 2009;
Yamada et al., 2011; 2013; Saito et al., 2018; Lee et al., 2019) or to estimate an f -divergence by optimizing a function from a suitable family (Nguyen et al., 2010) such as a neural network Nowozin et al.
Witness functions need not rely on the density ratio.
A wide class of divergences called integral probability metrics (IPMs) (M¨uller, 1997), which include total variation, the Wasserstein-1 distance, maximum mean discrepancy (MMD) (Gretton et al., 2007), and others (Mroueh et al., 2017), seek a witness function that maximizes the distance between the ﬁrst moments of the witness function evaluations.
In these cases the optimal witness function ω (cid:63) ( · ) has a greater expectation in one distribution compared to the other distribution.
An IPM between two measures µ and ν is expressed as sup ω ∈F |
E X ∼ µ [ ω ( X )] − E Y ∼ ν [ ω ( Y )] | for a given family of functions F .
A class of related divergences are the max-sliced Wasserstein-
p distances, which seek a linear (Deshpande et al., 2019) or non-linear slicing function (Kolouri et al., 2019) that maximizes the Wasserstein- p distance between the witness function evaluations for the two distributions.
However, there are two difﬁculties with computing the max-sliced Wasserstein distance for two samples.
The ﬁrst is that it is a saddlepoint optimization problem, whose objective evaluation requires sorting the samples.
Previous work has sought to approximate it using a ﬁrst moment approximation (Deshpande et al., 2019) or to use a ﬁnite number of steps of a local optimizer (Kolouri et al., 2019), without any guarantee of obtaining an optimal witness function.
Another difﬁculty is in the interpretation of the obtained witness function.
Unlike the density ratio, there is no notion of whether the witness function will take higher values for points associated to one distribution versus the other.
To address both of these issues we propose a max-sliced distance that replaces the Wassersteindistance with a second-moment approximation based on the Bures distance (Dowson & Landau, 1982; Gelbrich, 1990).
The Bures distance (Bures, 1969; Uhlmann, 1976) is a distance metric between positive semideﬁnite operators.
It is well-known in quantum information theory (Nielsen &
Chuang, 2000; Koltchinskii & Xia, 2015) and machine learning (Brockmeier et al., 2017; Muzellec & Cuturi, 2018; Zhang et al., 2020;
Oh et al., 2020; De Meulemeester et al., 2020).
1. C ONTRIBUTION
We propose a novel IPM-like divergence measure, the “max-sliced Bures distance”, to identify localized regions and instances associated with the maximum discrepancy between two samples.
The distance is expressed as the maximal difference between the root mean square (RMS) of the witness function evaluations sup ω ∈S (cid:12)(cid:12)(cid:12)(cid:112) E X ∼ µ [ ω 2 ( X )] − (cid:112) E Y ∼ ν
[ ω 2 ( Y )] (cid:12)(cid:12)(cid:12) , where S is an appropriate family of functions.
As | ∆ | = max { ∆ , − ∆ } , the max-sliced Bures can be expressed as the maximum of one-sided max-sliced divergences with optimal witness functions, ω µ>ν = arg max ω ∈S (cid:112) E µ [ ω 2 ( X )] − (cid:112) E ν
[ ω 2 ( Y )] , and ω µ<ν = arg max ω ∈S (cid:112) E ν [ ω 2 ( Y )] − (cid:112) E µ [ ω 2 ( X )] .
If the distributions are not well-matched, then ω µ>ν has large magnitude function evaluations under a ‘localized’ subset of µ and smaller magnitude values for ν , and the opposite for ω µ<ν .
The two samples { x i } mi =1 , { y i } ni =1 can be sorted by the magnitude of the witness function evaluations.
Crucially, we detail a tractable optimization procedure that is guaranteed to yield a global optimum witness function for the one-sided max-sliced Bures divergence.
When X = R d and the ﬁrst or second moments distinguish the distributions, linear witness functions can be used S = { ω ( · ) = (cid:104)· , w (cid:105) : w ∈ S d − 1 } , where S d − 1 denotes the unit sphere in R d .
The optimal witness function for the one-sided max-sliced Bures divergence ω µ>ν ( · ) = (cid:104)· , w µ>ν (cid:105) coincides with the subspace with the greatest difference in RMS, w µ>ν =
arg max w ∈ S
d − 1 (cid:112) w (cid:62) E [ XX (cid:62) ]
w − (cid:112) w (cid:62) E
[ Y Y (cid:62) ] w .
This optimization problem depends on the dimension d ; after computation of the covariance matrices, it is independent of the sample sizes m ≥ n .
In comparison, the optimal slice for the max-sliced Wasserstein may not be obtained, and even gradient ascent to a local optimum requires O ( m log m ) at each function/gradient evaluation.
Furthermore, the slice that maximizes the max-sliced Wasser-
[Equation]
[Table]
, (2) where ` π, ´ π, ` σ, ´ σ denote permutations and (cid:102) (cid:29) and (cid:102) (cid:28) denote expected inequalities with a large difference.
[Figure]
of stein lacks an intrinsic ordering, and it to the user to whether instances µ or ν have high or low values or magnitudes.
As second-order moments may be insufﬁcient for distinguishing the distributions, we explore non-linear mappings of the random variables.
Firstly, we consider a reproducing kernel Hilbert space (RKHS)
H with the family of witness functions S = { ω ( · ) = (cid:104) φ ( · ) , ω (cid:105)
H : ω, φ ( · ) ∈
H , (cid:104) ω, ω (cid:105) H = 1 } .
An example with Gaussian kernels is shown in Figure 1.
Secondly, we use a pre-trained neural network to create a task-relevant mapping, computing the second-order statistics of the hidden-layer activations, and apply this in the context of samples of natural images.
This enables interpretation of the Fr ´ echet Inception distance (FID) (Heusel et al., 2017) by identifying the subspace and images associated with discrepancies between synthetic and natural images.
We prove that the max-sliced Bures distance provides a lower bound on the max-sliced Fr´echet distance.
Because of their similarity, we develop the max-sliced Bures distance in the context of max-sliced versions of the total variation and Wasserstein-2 distances.
The kernel-based versions of these are novel contributions themselves.
The max-sliced total variation distance is a special case of the covariance feature matching proposed by Mroueh et al.
In experimental results, we show applications of the linear and kernel-based versions to detect imbalances in class distributions of natural images and to critique GANs.
We compare to other divergences expressed in terms of witness functions including MMD.
Finally, we propose algorithms to reweight an empirical distribution in order to minimize max-sliced divergences (with applications to generating conditional distributions and covariate shift correction).
M ETHODOLOGY
Consider a topological space X , a Borel σ -algebra B X , and the set Pr ( X ) of Borel probability measures on X .
Let µ, ν ∈ Pr( X ) denote two probability measures, and X ∼ µ and Y ∼
ν be two random variables X, Y ∈ X .
Let κ denote a positive-deﬁnite, bounded kernel function κ : X × X
→ B ⊂ R .
For any κ , there is an implicit mapping φ : X → H that maps any element x ∈ X to an element in the reproducing kernel Hilbert space (RKHS) φ ( x
) ∈ H such that κ ( x, y ) =
(cid:104) φ ( x ) , φ ( y ) (cid:105)
H = (cid:104) κ ( · , x ) , κ ( · , y ) (cid:105) H for x, y ∈ X , and ∀ ω ∈ H , ω ( x )
= (cid:104) ω, κ ( · , x ) (cid:105) H (Aronszajn, 1950).
∀ ω ∈ H , (cid:107) ω (cid:107) 2 = (cid:112) (cid:104) ω, ω (cid:105) H .
When clear, we drop the H subscript on the inner product.
A rank-1 RKHS operator is denoted as ω ⊗ ψ ∈ H × H with (cid:104) ( ω ⊗ ψ ) φ ( x ) , φ ( y ) (cid:105) =
(cid:104) ψ, φ ( x ) (cid:105)(cid:104) ω, φ ( y ) (cid:105)
= ψ ( x ) ω ( y ) for x, y ∈ X .
Denote by m X =
E X ∼ µ [ φ ( X )]
∈ H and m Y =
E Y ∼ ν
[ φ ( Y )] ∈ H the ﬁrst moments of the random variables in the RKHS.
The uncentered second moments are ρ X =
E
[ φ ( X ) ⊗ φ ( X )]
∈ H ×H and ρ Y = E
[ φ ( Y ) ⊗ φ ( Y )]
∈ H ×H .
The covariance operators are Σ X
= ρ X − m X ⊗ m X and Σ Y =
ρ Y − m Y ⊗ m Y .
2. D IVERGENCES AS D ISTANCE M ETRICS
Let D ( µ, ν ) denote a divergence D : Pr( X )
×
Pr
X ) → [0 , ∞ ) .
It is a distance metric between measures (a probability metric) if all of the following statements hold: (i) µ = ν = ⇒ D ( µ, ν ) = 0 , (ii) D ( µ, ν ) =
= ⇒ µ = ν , (iii) D ( µ, ν ) =
D ( ν, µ ) , (iv) D ( µ, ν ) ≤ D ( µ, ξ ) + D ( ν, ξ ) .
It is a semi-metric if all properties aside from (ii) hold.
M¨uller (1997) deﬁnes the class of integral probability metrics as the supremum of the absolute difference between expectations
[Equation]
With appropriate choice of the family of functions F , this form yields well-known divergences (Sriperumbudur et al., 2010), e.g., when F is the set of functions with Lipschitz constant less than 1, the resulting divergence is the Wasserstein-1 distance metric.
Another example of an
IPM is when F = { ω ∈ H : (cid:107) ω (cid:107) 2 ≤ 1 } , which yields MMD (Gretton et al., 2007), deﬁned as
D
[Equation] [Equation]
=
[Equation] [Equation]
For characteristic kernels such as the Laplacian and Gaussian kernels, the mean embedding E
X ∼ µ [ φ ( X )] : Pr ( X ) → H is an injective function (Sriperumbudur et al., 2008; Fukumizu et al., 2009; Sriperumbudur et al., 2010), capturing the full statistics of µ .
In these cases, MMD is a distance metric on Pr ( X ) ; likewise, distance metrics between the operators ρ X =
E [ φ ( X ) ⊗ φ ( X )] and ρ Y =
E [ φ ( Y ) ⊗ φ ( Y )] induce probability metrics for characteristic kernels (Zhang et al., 2020).
2. O PERATOR D ISTANCES FOR D EFINING D IVERGENCES
Total variation (TV) is a well-known probability metric and an integral probability metric (M ¨ uller, 1997), taking the form (1 / 2) (cid:80) i | p i − q i | for discrete measures, for which p i = µ ( x i ) and q i = ν
( x i ) where { x i } i = X .
The TV distance between operators in the RKHS is a divergence
[Equation]
[Equation]
where (cid:107)·(cid:107) 1 denotes the trace norm (Schatten 1-norm), which is the sum of the singular values.
The Bures distance generalizes the Hellinger distance (cid:112) (1 / 2) (cid:80) i ( √ p i − √ q i ) 2 to positive semideﬁnite operators (Fuchs & Van De Graaf, 1999; Bromley et al., 2014; Bhatia et al., 2019).
The kernel Bures divergence D H B ( µ, ν ) and the Bures distance d B ( ρ X , ρ Y ) are deﬁned as
[Equation]
The Bures distance is used to deﬁne the Wasserstein-2 (W2) distance between Gaussian measures, i.e., the Fr´echet distance (Fr´echet, 1957; Dowson & Landau, 1982).
The multivariate Fr´echet distance provides a lower bound for the W2 distance (Gelbrich, 1990).
The kernel Gauss-Wasserstein distance (Zhang et al., 2020; Oh et al., 2020) is deﬁned as
[Equation]
Zhang et al.
(2020) also proposed the kernel Wasserstein-
p distance between µ and ν ,
[Equation]
where Γ( µ, ν ) deﬁnes the set of all joint distributions coupling µ and ν , and d p κ ( X, Y )
= (cid:107) φ ( X ) − φ ( Y ) (cid:107) p 2 .
For p = 2 , d 2 κ ( X, Y ) =
κ ( X, X )+ κ ( Y, Y ) − 2 κ ( X, Y ) .
When p = 2 and φ ( x ) (cid:55)→ x ∈ R d such that H = R d , the standard W2 distance W R d 2 ( µ, ν ) is obtained.
2. D IVERGENCES BASED ON SLICING H ILBERT SPACES
The sliced Wasserstein distance (Wu et al., 2019; Deshpande et al., 2018; Kolouri et al., 2018), and max-sliced Wasserstein distance (Deshpande et al., 2019; Kolouri et al., 2019) evaluate discrepancies in linear or non-linear one-dimensional subspaces.
A motivation for this is the analytic solution of the Wasserstein- p distance in one dimension.
The max-sliced Wasserstein- p distance takes the form max- W R
[Equation]
µ, ν
[Equation] [Equation]
(cid:104)
[Equation] [Equation]
(cid:105)
[Equation]
] , p ≥ 1 .
Similarly, we propose the max-sliced Bures, the kernel TV, and the kernel Wasserstein-
p distances using the rank1 operator Ω = ω ⊗ ω ∈ H × H , which projects (slices) the RKHS along a one-dimensional subspace deﬁned by the ray ω ∈ H , with (cid:104) ω, φ ( X ) (cid:105) =
ω ( X ) , due to the reproducing property.
In this formulation, ω : X → R is the witness function from the set S = { ω ∈ H : (cid:107) ω (cid:107) 2 =1 } .
Notably, a linear slice in the RKHS is a possibly non-linear function in the input space.
For conciseness, we denote the mean square witness function evaluations
E [ ω 2 ( X )]
= (cid:104) ω, ρ X ω (cid:105) =
(cid:107)√ ρ X ω (cid:107) 22 as (cid:107) ω (cid:107) 2 µ , and E [ ω 2 ( Y )]
= (cid:104) ω, ρ Y ω (cid:105) =
(cid:107)√ ρ Y ω (cid:107) 22 as (cid:107) ω (cid:107) 2 ν .
The RMS (cid:107) ω (cid:107) µ is an Lsemi-norm on ω induced by the positive semideﬁnite operator √ ρ X .
The max-sliced kernel TV,
Bures, and W2 distances, derived in appendix A.1, are expressed as
[Equation]
respectively.
The inner supremums in equation 8 and equation 9 are the one-sided divergences.
The max-sliced TV distance is an IPM with F = { ω 2 ( x ) = (cid:104) φ ( x ) , ω (cid:105) 2 : ∀ ω ∈ S} and is a special case of the IPM Σ divergence proposed by Mroueh et al.
While not an IPM, the max-sliced kernel Bures distance can be directly related to max-sliced versions of the TV, Fr ´ echet, and W distance, as detailed by the following results.
Theorem 1.
The square of the max-sliced Bures distance in the RKHS H is less than or equal to twice the max-sliced TV distance, max- D H B ( µ, ν ) ≤ (cid:113) 2 (cid:0)
max- D H TV ( µ, ν ) (cid:1) .
Theorem 2.
The max-sliced Bures distance in the RKHS H is a lower bound on the kernel maxsliced Gauss-Wasserstein distance, max- D H B ( µ, ν ) ≤ max L - D H GW ( µ, ν )
≤ max U - D H GW ( µ, ν ) .
Theorem 3.
The max-sliced Bures distance in the RKHS is a lower bound on the kernel max-sliced
W2 distance, max- D H B ( µ, ν ) ≤ max- W H 2 ( µ, ν ) .
These results trivially translate to the linear kernel case
H = R d , S =
S
d − 1 , ω ( x ) = (cid:104) w , x (cid:105) , for w , x ∈ R d and X ⊆ R d .
The latter two show that the max-sliced Bures distance is a lower-bound on the max-sliced Fr´echet distance, which is a lower-bound on the max-sliced W2 distance.
The proofs and other relationships among the divergences are in Appendix A.2.
2. C OMPUTING THE M AX - SLICED D IVERGENCES
The max-sliced kernel Bures, TV, and W2 distances require solving optimization problems to ﬁnd the optimal witness function.
As noted by others investigating IPMs (Mroueh et al., 2017; Li et al., 2017; Kolouri et al., 2019), the witness function can be deﬁned using a family of functions implemented as a neural network.
In this context, Goodfellow et al.
(2014) use the divergence
D A ( µ, ν ) = max ω E [log ω
( X )]
+ E [log (1 − ω ( Y ))] , where ω : X → (0 , 1) .
In comparison, the Wasserstein-1 distance D W 1 ( µ, ν ) = sup ω ∈ Lip 1 E [ ω ( X )] − E [ ω ( Y )] requires a Lipschitz con- straint (Arjovsky et al., 2017; Gulrajani et al., 2017).
3 Table 1 compares the form and constraints.
[Equation]
[Table]
[Equation]
[Table]
[Equation]
2. S AMPLE -B ASED E STIMATORS FOR M AX -S LICED D IVERGENCES
We consider the case of ﬁnite samples expressed as empirical measures ˆ µ = (cid:80) mi =1 µ i δ x i and ˆ ν = (cid:80) ni
=
ν i
δ y
i for the samples { x i } mi =1 and { y i } ni =1 with discrete probability masses denoted as column vectors [ µ 1 , . . . , µ m ] (cid:62) =
µ ∈ [0 , 1] m , (cid:104) µ , 1 (cid:105) = 1 , and [ ν 1 , . .
ν n ] (cid:62)
=
ν ∈ [0 , 1] n , (cid:104) ν , 1 (cid:105) = 1 .
The kernel-based max-sliced divergences optimize the witness function ω ( · ) = (cid:80) li =
α i κ ( · , z i ) in terms of the dual variables α ∈ R l corresponding to a subset of the pooled sample { x i } mi =1 ∪{ y i } ni =1 .
The optimization problems and algorithms are detailed in appendix A.4.
For clarity, we proceed to the linear kernel case for a ﬁnite-dimensional embedding φ ( x ) = x ∈ R d .
After embedding, kernel evaluations correspond to vector inner-products κ ( x, y ) =
(cid:104) φ ( x ) , φ ( y ) (cid:105) =
x (cid:62) y .
Let X =
[ x 1 , . . .
, x m ] ∈ R d × m and Y =
[ y 1 , . . .
, y n ] ∈ R d × n denote the sample points with corresponding masses µ and ν , respectively.
The witness function is the inner product ω ( x ) = w (cid:62) x , where the variable w deﬁnes the slice with (cid:107) ω (cid:107) 2ˆ µ = w (cid:62) ρ X w with ρ X =
XD µ X (cid:62) and (cid:107) ω (cid:107) 2ˆ ν = w (cid:62) ρ Y w with ρ Y =
YD ν Y (cid:62) , where D v is diagonal with entries v .
For i.i.d.
samples,
D µ = 1 m I and D ν = 1 n I .
The max-sliced TV, Bures, and W2 divergences are
[Equation]
where λ 1 ( · ) denotes the largest magnitude eigenvalue of the argument and P ˆ µ, ˆ ν = { P ∈ [0 , 1] m × n | P1 n = µ , P (cid:62)
1 m = ν } is a transportation polytope.
These three optimizations differ in difﬁculty.
The ﬁrst two require only the sample means m X , m Y and covariance matrices Σ X , Σ Y , since ρ X
= m X m (cid:62) X
+ m − 1 m Σ X
and ρ Y =
m Y m (cid:62) Y + n − 1 n Σ Y (assuming unbiased covariance estimates).
The one-sided max-sliced TV divergences can be solved by ﬁnding the eigenvectors associated to the largest eigenvalues of ρ X − ρ Y and ρ Y − ρ X .
Likewise the optimal slice for each one-sided max-sliced Bures divergence requires solving a series of eigenvector problems.
Speciﬁcally, if ρ X and ρ Y are strictly positive deﬁnite, then the optimal witness function is ω µ>ν ( · ) =
(cid:104) w [Equation] (cid:63) , ·(cid:105) , where γ (cid:63) ∈ (0 , 1] solves the optimization problem
[Equation]
[Equation]
The general case involves checking the nullspace of ρ Y and is given in the Appendix A.5.
In comparison, the max-sliced W2 distance is a saddlepoint optimization problem (Deshpande et al., 2019).
Following Kolouri et al.
(2019), gradient ascent on w can be performed with ﬁrst-order solves.
Each gradient evaluation requires solving the transport map by sorting X (cid:62) w and Y (cid:62) w .
For this, we use
ADAM (Kingma & Ba, 2015) and quasi-Newton approaches, such as MIN F UNC (Schmidt, 2012).
The same approaches can be used to approximate max- D B after smoothing √· as √· + 0 .
E XPERIMENTS
We present various examples of using the proposed max-sliced divergences to identify the discrepancies between two samples.
We apply the proposed approach to detect mismatched distributions of natural and fake images using the internal representation of the Inception Network (Szegedy et al., 2016) as in the Fr´echet Inception distance (Heusel et al., 2017) and the Inception score (Salimans et al., 2016).
We investigate whether the witness functions detect covariate shift caused by class imbalances.
Then, we propose optimizing the weights ν to compensate for covariate shift.
Finally, we use the one-sided max-sliced Bures divergence to monitor mode dropping during GAN training.
3. I NTERPRETING THE F R ´ ECHET I NCEPTION D ISTANCE
We use a linear witness function to identify instances that are not well matched between two samples of real or fake images represented by internal activations of the Inception object classifying network (Szegedy et al., 2016).
Speciﬁcally, we search for witness functions with the form ω ( x ) =
(cid:104) w , φ ( x ) (cid:105) , where the vector φ ( x ) ∈ R 2048 is an Inception code—the internal activations of penultimate layer of the network after pooling (Heusel et al., 2017).
Figure 2 shows the performance of the proposed measure to identify instances associated with imbalanced representation of particular classes.
In particular, ˆ µ is a uniform sample from the training set and ˆ ν is a sample from the test set with less instances from one class.
Using the one-sided max-sliced Bures divergence we obtain the optimal slice ω µ>ν and apply it to the imbalanced sample ˆ ν , identifying the top-10 witness points with the largest magnitude witness function evaluations ω 2 µ>ν ( y ` σ (1) ) ≥ . . .
≥ ω 2 µ>ν ( y ` σ (10) ) , where ` σ is a permutation corresponding to sorting { y i } ni =1 by descending magnitude.
(While this may seem counterintuitive as it is expected that max 1 ≤ i ≤ m ω 2 µ>ν ( x i ) (cid:29) max 1 ≤ i ≤ n ω 2 µ>ν ( y i ) , since ω µ>ν corresponds to a one-dimensional subspace, { y ` σ ( i ) }
Ki =1 are the K instances from ˆ ν with the largest norm after projection to this subspace.)
The performance is quantiﬁed by the precision of the labels of these instances (ideally, these witness points should be from the underrepresented class).
Notably, the mean precision of the topinstances is 0.79 or better across the classes with a mean average precision (MAP) of 0.94 when the class probabilities differ by 2% (10.2% for majority and 8.2% for minority).
This is compared to a
MAP of 0.82 for the ﬁrst-moment based surrogate of the max-sliced W2 distance (Deshpande et al., 2019).
Computing the max-sliced W2 distance takes much longer to run on this sample size.
generate set of synthetic for an AutoGAN instance
-trained
(Gong et al.,
), which has an Inception score of
525 and a FID score of 12.41.
We applied both one-sided max-sliced Bures divergences to identify the two subspaces that maximize the difference in RMS between fake and real images.
Figure 3 details the top-10 images in each subspace and their realism scores R (Kynk¨a¨anniemi et al., 2019).
4 Applying the max-sliced
Wasserstein-2 distance yielded almost the same solution as w Real <
Fake (a linear correlation of 0.992).
[Figure] [Figure]
B ASELINE C OMPARISON ON C OVARIATE S HIFT D ETECTION
We compare the proposed max-sliced Bures distance and the resulting max-sliced Fr ´ echet distance to the max-sliced W2 distance for linear witness functions.
Figure 4 compares the divergence estimates instances associated with a simple case of covariate shift with the MNIST data set.
Notably, the precision of detecting class imbalances is higher for smaller samples using a kernel (Appendix A.8).
[Figure]
3. C OVARIATE S HIFT C ORRECTION BY R EWEIGHTING
We consider the task of reweighting the instances in one sample to minimize the max-sliced Bures distance.
This optimization problem can be expressed as min ν ∈ R n ≥ (cid:80) i ν i =
1 J ( ν ) , where J ( ν ) ∝ max- D R d B (ˆ µ, ˆ ν ) .
As shown in appendix A.7, this is a convex minimization problem with a simplex constraint on ν .
We apply the Frank-Wolfe algorithm
(Jaggi, 2013) to iteratively adjust the weight of one instance at each iteration.
The performance is quantiﬁed in terms of the Fr´echet Inception distance between the real-test images of the class and the reweighted sample of fake images.
For comparison, we also optimize reweightings that minimize the W2 distance and the max-sliced W distance (using 10 mini-batches of size n = m = 100 at each iteration).
The average FID distance across the classes is 49.15 for the max-sliced Bures reweighting compared to 68.1 and 72.5 for the mini-batch W2 and max-sliced W2.
(See Table 4 in the Appendix for full results).
Figure 5 shows results of a reweighting uniform distributions to match target distributions using either linear slices or random Fourier bases (Rahimi & Recht, 2008) for approximating a Gaussian kernel.
For the latter, using the max-sliced Bures as a loss achieves the lowest reweighted W distance, which is computed by solving a discrete transportation problem (Flamary & Courty, 2017)
[Figure]
Reweighting a uniform distribution to match various target distributions by minimizing the max-sliced W2 distance or the max-sliced Bures distance with either a linear kernel or random Fourier bases ( d =2,000, σ ∈ { 0 . 1 , 0 . 2 } ).
Examples follow Kolouri et al.
and uses
ADAM (Kingma & Ba, 2015) defaults and a learning rate of 10 − 2 .
A point’s size is proportional to their weights after 100 iterations.
Learning curves are the weighted W2 distance (log-scale).
3. D ETECTING M ODE D ROPPING
We create a 3-channel “Stacked MNIST” data set with 500,000 images from the MNIST training set to test mode dropping detection throughout GAN training (DCGAN architecture).
The training set ˆ µ has 1000 possible modes corresponding to all 3-digit combinations.
At the end of each training epoch (1000 iterations), a fake sample n = 10 4 is generated.
To verify mode coverage we use a 4-layer conv.
net trained on single-channel MNIST; any missing combination of 3-digit labels is considered a dropped mode.
Figure 6 details using the proposed approach to detect missing modes.
[Figure]
C ONCLUSION
We propose the max-sliced Bures distance, a lower-bound on the max-sliced W2 distance, which can be computed optimally with a tractable algorithm.
We show increased performance with kernelbased witness functions for covariate shift detection and correction, and also highlight its utility in the linear case when the feature space is the internal representation of a pre-trained network.
Importantly, the one-sided max-sliced Bures divergences enable direct interpretation of under- and over-representation between two samples, which can be used to identify systematic discrepancies.
A A PPENDIX
The appendix details the derivation of the proposed divergences, formal results relating them to existing divergences, algorithms and additional experimental results.
A. D ERIVATION OF THE MAX - SLICED KERNEL DIVERGENCES
Let U 1 = { ω ⊗ ω : | ω | 2 ≤ ω ∈ H} denote the set of rank-1 symmetric operators with bounded trace norm.
Let d denote either the TV distance ( d = d TV ) or the squared Bures distance ( d = d 2 B ) deﬁned in equation 4 and equation respectively.
In these cases the expression of the max-sliced distance can be simpliﬁed as
[Equation]
where (cid:104)· , ·(cid:105) HS :
( H × H ) × ( H × H ) → R denotes the inner-product deﬁning the Hilbert-Schmidt (Schatten-2 norm), (cid:104) Ω , ρ (cid:105) HS = (cid:104) ω ⊗ ω, ρ (cid:105) HS =
(cid:104) ω, ρω (cid:105) H , and δ ( p, q )
=
d ( p Ω , q Ω) denotes the distance between scaled versions of Ω , for p ∈ R ≥ 0 and q ∈ R ≥ 0 .
The equalities follow from the fact that Ω ρ Ω = ( ω ⊗ ω ) ρ ( ω ⊗ ω )
= (cid:104) ω, ρω (cid:105) Ω .
The max-sliced TV distance and squared max-sliced
Bures distance yield expressions for δ ( p, q ) that match the form of the underlying TV and Hellinger divergences: d TV ( p Ω , q Ω) yields δ ( p, q ) =
12 | p − q | , and d 2 B ( p Ω , q Ω) yields δ ( p, q ) =
( √ p −√ q ) 2 : max- d TV ( ρ X , ρ
[Equation]
(cid:44)
[Equation]
|
[Equation] [Equation]
(cid:105)
[Equation] [Equation]
(cid:105)| , and (15)
[Equation]
(16) where S = { ω ∈ H : (cid:107) ω (cid:107) 2 = 1 } .
The kernel-based divergences can be expressed in terms of the witness functions since (cid:104) ω, ρ X ω (cid:105) =
E [ (cid:104) φ ( X ) , ω (cid:105)(cid:104) φ ( X ) , ω (cid:105) ]
=
E [ ω 2 ( X )] and likewise for ρ Y .
The underlying divergences and distances are listed in Table 2.
[Table]
[Equation]
Slicing is natural for the operator-based distances, as it is inherent in their deﬁnition such that they coincide with the corresponding divergences between discrete probability laws (Fuchs & Van
De Graaf, 1999).
The scalar discrepancy measure δ ( · , · ) can be accumulated across a complete set of slices to obtain the original distances.
Consider a set (or countably inﬁnite sequence) of orthogonal trace-norm operators O = { Ω 1 = ω 1 ⊗ ω 1 , Ω 2 = ω 2 ⊗ ω 2 , . .
Ω k = ω
k ⊗ ω k } .
Since these operators are orthogonal and have unit trace-norm (cid:107) (cid:80)
ki
[Equation]
1 Ω i (cid:107) ∞
= 1 ,
[Equation]
= k (cid:88) i =
δ ( p i ( O ) , q i ( O )) , (17)
where p i ( O ) = (cid:104) Ω i , ρ X (cid:105) HS and q i ( O ) = (cid:104) Ω i , ρ Y (cid:105) HS .
To obtain the equality, one must optimize O over all possible sets of orthogonal trace-norm operators.
A.1. M AX - SLICED TOTAL VARIATION DISTANCE
The sliced kernel total variation distance is
[Equation]
ρ X − ρ Y
[Equation]
|
[Equation]
Maximizing over slices yields
[Equation]
where the ﬁrst equality follows from the fact that distance is maximized when (cid:107) Ω (cid:107) 1 = 1 .
This yields the expression
[Equation]
Notably, the penultimate expression in equation 18 can be related to the operator norm,
[Equation]
due to the dual norm deﬁnition.
Since ρ X − ρ Y is symmetric, the equality is achieved.
Thus, max- D H TV ( µ, ν ) = 12 (cid:107) ρ X − ρ Y
[Equation]
∞ .
For a linear kernel, this can be computed by ﬁnding the largest
magnitude
eigenvalue of ρ X −
[Equation]
Y
[Equation]
A.1. M AX - SLICED B URES DISTANCE
The sliced version of the Bures distance is
[Equation]
which can be simpliﬁed since (cid:107) Ω
[Equation]
X
[Equation]
(cid:107)
= (cid:107) ω (cid:107) 2 µ (cid:107) ω (cid:107) 22 , and (cid:107)
[Equation]
X
[Equation]
ω
[Equation]
ω
[Equation]
ν
[Equation]
Ω
[Equation]
1 = (cid:107) ω (cid:107) µ (cid:107) ω (cid:107) ν (cid:107) ω
[Equation]
Using these expressions, the max-sliced Bures distance is
[Equation]
The expression is monotonic with the norm of ω yielding
[Equation]
A.1. M AX - SLICED G AUSS -W ASSERSTEIN DISTANCES
Two max-sliced versions of the Gauss-Wasserstein or Fr´echet distance in the RKHS are
[Equation]
where σ ω ( X ) =
(cid:112) E [( ω ( X ) − (cid:104) m
X , ω (cid:105) ) 2 ] is the standard deviation of ω ( X )
[Equation]
and σ ω ( Y ) =
(cid:112)
E [( ω ( Y ) − (cid:104) m Y , ω (cid:105) ) 2 ] , and
[Equation]
Notably, max L - D H GW ( µ, ν ) is the supremum of the Fr´echet distance over all witness functions.
A.1. M AX - SLICED K ERNEL W ASSERSTEIN
A sliced version of the kernel Wasserstein distance between µ and ν relies on the the sliced distance
[Equation]
This distance is monotonic with the norm of ω and convex with respect to ω .
The max-sliced kernel
Wasserstein-
p distance, p ≥ 1 , is
[Equation]
which is a one-dimensional optimal transport problem.
Let ω
# µ and ω # ν denote the pushforward measures, then the divergence can be written as
[Equation]
[Equation]
where Π( ω # µ, ω # ν ) is the set of all joint distributions coupling the pushforward measures.
Assuming the measures are absolutely continuous and adopting the notation from Santambrogio (2015), let F ω,µ ( w )
=
(cid:82) w −∞
dω
# µ = ω
# µ (( −∞ , w ]) and F ω,ν ( w ) = (cid:82) w −∞ dω # ν = ω
# ν (( −∞ , w ]) denote the cumulative distribution functions of the pushforward measures with pseudo-inverses F − 1 ω,µ ( q ) = inf { w ∈ R : F ω,µ ( w ) ≥ q } and F − 1 ω,ν ( q )
= inf { w ∈ R : F ω,ν ( w ) ≥ q } .
As shown in Lemma 2.8 (Santambrogio, 2015), then the optimal transport plan π ∗ has cumulative distribution G ω ( w X , w Y )
=
min { F ω,µ ( w X ) , F ω,ν ( w Y ) } and the divergence is
[Equation] [Equation]
and for the case p = 1 (Santambrogio, 2015, Proposition 2.17),
[Equation]
The objective in the last quantity is an L 1 -norm version of the Cram ´ er–von Mises criterion (Schmid & Trede, 1995; Anderson, 1962).
That is, the max-sliced kernel Wasserstein-1 distance is equivalent to a max-sliced L 1 -norm version of the Cram ´ er–von Mises criterion, where the slicing corresponds to a function in the RKHS that witnesses the largest discrepancies between the measures.
The choice of p = 2 simpliﬁes, yielding the following optimization problem
[Equation]
A. R ELATIONSHIP BETWEEN THE KERNEL - BASED MAX - SLICED DIVERGENCES
The Bures distance can be used to lower bound the TV distance, d 2 B ( ρ X , ρ Y ) ≤ (cid:107) ρ X − ρ Y (cid:107) 1 = 2 d TV ( ρ X , ρ Y ) (Fuchs & Van De Graaf, 1999).
This inequality stems from the inequality between the squared Hellinger and total variation distances for discrete probability laws.
(cid:80) i
(cid:107)√ p i −√ q i (cid:107) 22 ≤ (cid:80) i
| p i − q i | , where the inequality holds for each summand, |√ p
i − √ q i | 2 ≤ |√ p i − √ q i | ( √ p i + √ q i ) =
| p i − q i | .
This inequality holds for the max-sliced divergences as stated in Theorem 1.
Proof of Theorem 1.
( (cid:107) ω (cid:107) µ − (cid:107) ω (cid:107) ν ) 2 ≤ (cid:12)(cid:12) (cid:107) ω (cid:107) 2 µ − (cid:107) ω
(cid:107) 2 ν (cid:12)(cid:12) , where the inequality is a consequence of the inequality between the arithmetic and geometric means, let a = (cid:107) ω (cid:107) 2 µ and b = (cid:107) ω (cid:107) 2 ν , then ( √ a − √ b ) 2 ≤ |√ a − √ b | ( √ a + √ b ) = | a − b | .
Taking the supremum over ω ∈ H with the constraint (cid:107) ω (cid:107) 2 ≤ 1 yields the desired the result (cid:0)
[Equation]
We now consider the relationship between the Gauss-Wasserstein or Fr´echet distance—which combines the distances between the ﬁrst and second-order moments—with the max-sliced Bures when it is applied directly to the uncentered covariance matrices.
For this we need the following lemma.
Lemma 4 (Reverse triangle inequality) .
For two vectors in R 2 , the difference between their Euclidean norms is less than or equal to the Euclidean norm of their differences.
For a, b, c, d ∈ R , (cid:12)(cid:12) √
[Equation]
(cid:12)(cid:12) ≤ (cid:112) ( a −
[Equation] [Equation]
+ ( b − d ) 2 .
Proof.
Let e = (cid:0) √ a 2 + b 2 − √ c 2 + d 2 (cid:1) 2 and f = ( a − c )
+ ( b − d ) 2 .
e =
[Equation]
the arithmetic and geometric mean inequality
[Equation]
Taking the square root of each side yields the inequality √ e ≤ √ f .
Proof of Theorem 2.
To relate the sliced Bures and Gauss-Wasserstein distances
[Equation]
we note that
[Equation] [Equation]
HS =
σ 2 ω ( X ) + (cid:104) m X , ω (cid:105) 2 and likewise for (cid:107) ω (cid:107) 2 ν .
Then,
[Equation]
(30) where the inequality relies on Lemma 4, with a = (cid:104) m
X , ω (cid:105) , b = σ ω ( X ) , c = (cid:104) m Y , ω (cid:105) , and d = σ ω ( Y ) .
Taking supremum over slices yields the desired inequality.
Theorem 2 shows that the max-sliced kernel Bures distance is a lower-bound on the max-sliced kernel Wasserstein-2 distance, since the latter is lower bounded by the max-sliced kernel GaussWasserstein distance.
Proof of Theorem 3.
Let ω ( X ) =
(cid:104) m
X , ω (cid:105) + ˜ ω
( X ) and ω ( Y ) = (cid:104) m
Y , ω (cid:105) + ˜ ω
( Y ) , where ˜ ω ( X ) = (cid:104) φ
( X ) − m X , ω (cid:105) and ˜ ω
( Y ) =
(cid:104) φ ( Y ) − m Y , ω (cid:105) are zero mean, and E [˜ ω
2 ( X )]
= σ 2 ω ( X ) and
E [˜ ω 2 ( Y )]
= σ 2 ω ( Y ) .
The squared Fr´echet distance between random variables ω ( X ) and ω ( Y ) is
[Equation]
By H¨older’s inequality, E [˜ ω ( X )˜ ω
( Y )]
≤ σ ω ( X ) σ ω
( Y ) =
(cid:112) E
[˜ ω 2 ( X )]
E [˜ ω 2 ( Y )] .
Consequently,
[Equation]
Taking the inﬁmum over all possible joint distributions ( X, Y ) ∼ γ that are within the coupling distribution γ ∈ Γ , yields the sliced Wasserstein-2 distance on the right hand side.
Maximizing over slices, yields max U - D H GW ( µ, ν ) ≤
max- W H 2 ( µ, ν ) and combining with Theorem 2 yields the desired result.
A. R ELATIONSHIP TO O THER D IVERGENCES
Finally, we note that the terms in the max-sliced Gauss-Wasserstein divergences are related to kernel
Fischer discriminant analysis (KFDA) (Mika et al., 1999).
KFDA objective is based on the ratio of the difference in means to the pooled variances:
[Equation]
KFDA seeks a witness function which has widely separated means for the two measures, and minimal variance.
A. C OMPUTING THE M AX -S LICED K ERNEL D IVERGENCES
We assume the witness function 6 ω ∈ H is of the form ω
= (cid:80) n + m i
=
α i φ ( z i ) with ω ( · ) = (cid:80) n + m i =
α i κ ( · , z i ) where z i = (cid:26) x i , 1 ≤ i ≤ m y i − m , m < i ≤ n + m and α ∈ R m + n .
In this case, (cid:107)
[Equation]
(cid:104)
[Equation] [Equation]
µ ) (cid:105)
= m (cid:88)
j
[Equation]
µ
[Equation] [Equation]
[Equation] φ ( x
[Equation]
) (cid:105) =
[Equation]
=
[Equation]
=
[Equation]
, where κ ( z i , z j ) =
K ij , K = (cid:20) K XX , K XY
K
[Equation]
= (cid:20) K XZ
K
[Equation]
Z (cid:21) , ( · ) ◦ 2 denotes the elementwise squaring of a matrix/vector, and D v = diag ( v ) denotes a diagonal matrix whose diagonal entries are the vector v .
Similarly, (cid:107) ω (cid:107) 2ˆ ν = (cid:104) ν , ( K Y Z α ) ◦ 2 (cid:105)
= (cid:107) D 12 ν K Y
Z α (cid:107) 22 .
In order for the constraint (cid:107) ω (cid:107) 22 ≤ 1 = ⇒ α (cid:62) K α ≤ 1 to ensure a bounded solution, we assume
K is strictly positive deﬁnite.
For this purpose, we add a small value to its diagonal K + 10 − 9 I when necessary in the optimization procedures.
For computational purposes when m or n are large, a subset (possibly random) of landmark points of size l < n + m can be used to form the witness function ω
= (cid:80) li =
α i φ ( z τ i ) , where { τ i } li =1 ⊂ { 1 , . . .
, m + n } .
In this case, K XZ ∈ R m × l and
K Y Z ∈ R n × l with [ K XZ ] i,j = κ ( x i , z τ j ) and [ K Y Z ] i,j = κ ( y i , z τ j ) .
In this case, the constraint also needs to be adjusted.
A.4. M AX -S LICED K ERNEL TV D ISTANCE
Using these expressions, the max-sliced kernel TV distance is
[Equation]
[Equation]
The solution is the generalized eigenvector corresponding to the largest magnitude eigenvalue of the generalized eigenvalue problem Av = λ
Kv , where A = K (cid:62) XZ D µ K XZ − K (cid:62) Y Z D
ν K Y Z .
α (cid:63) = arg max [Equation] α (cid:62)
A α α (cid:62) K α .
The witness function is ω (cid:63) ( · ) =
(cid:80) m
+ n i =
1 α (cid:63)i κ ( · , z i ) .
A.4. M AX -S LICED K ERNEL B URES D ISTANCE
The sample-based max-sliced kernel Bures distance is
[Equation]
The last expression shows that the max-sliced Bures distance can be expressed as a bilevel optimization problem, where the inner optimization problem—which we refer to as one-sided max-sliced
Bures divergence—is a difference of convex functions:
[Equation]
Without loss of generality
[Equation]
we will consider the ﬁrst case ( s = 1 ),
[Equation]
Inspired by the approach Landsman (2008)
[Equation]
we relate this problem to a quadratic program, speciﬁcally, the quadratically constrained quadratic program
[Equation]
(Q
[Equation]
for c 1 , c 2 ∈ R ≥ 0 .
The solution of which can be obtained as in the max-sliced kernel TV distance by solving a generalized eigenvalue problem.
For (P) , the Lagrangian function is
[Equation] [Equation]
Let G =
{ α : g ( α ) > 0 , h ( α ) > 0 } denote the set of points where g and h are differentiable.
Then for α ∈ G and K positive deﬁnite, L ( α , λ ) is differentiable, and
[Equation]
where the equality follows from ∇ α g 2 ( α )
= 2 g ( α ) ∇ α g
( α ) and likewise ∇ α h 2 ( α )
= 2 h ( α ) ∇ α h ( α ) .
For (Q) , the Lagrangian function and its gradient are
[Equation]
[Equation]
If c 1 = 1 2 g ( α ) and c 2 = 1 2 h ( α ) , then ∇ α L ( α , λ ) = ∇ α ¯ L ( α , λ ) .
Let α ∗ denote a global optimum of (Q) .
If c 1 = 1 2 g ( α ∗ ) and c 2 = 1 2 h ( α ∗ ) , then ∇ α L ( α , λ ) | α = α ∗ = ∇ α ¯ L ( α , λ ) | α = α ∗ .
Consequently, α ∗ is a local optimum of (P) .
By the Karush–Kuhn–Tucker conditions, it is a necessary condition for all optima of (P) in G to have this form.
Thus, any global optimum of (P) that lies in G corresponds to a global optimum of (Q) for particular values of c 1 , c 2 .
The family of solutions to (Q) that includes all local optima of (P) in G , is
[Equation]
where the bounds are due to the non-negative functions and h 2 ( α ∗ γ ) ≥ g 2 ( α ∗ γ )
=
⇒ 1 2 c 2 ≥ 12 c 1 = ⇒ c
1 ≥ c 2 = ⇒ c
2 c 1 ≤ 1 .
Notably, γ = 1 = ⇒
c
1 = c 2 corresponds to the one-sided max-sliced TV.
The global optimum of (P) within G is necessarily within { α ∗ γ } γ ∈ (0 , 1] and can be found as the solution to the bound scalar optimization problem
min γ ∈ (0 , 1] g ( α ∗ γ ) − h ( α ∗ γ ) .
A remaining case for a global optimum is a non-differentiable point α (cid:63) / ∈ G , specifically, g ( α (cid:63) )
=
(the case of h ( α (cid:63) )
= 0 is trivial), which corresponds to α (cid:63) ∈
Null( D ν K Y Z ) .
In this case, the generalized eigenvalue problem must be restricted to the nullspace α ∗ ∅ =
arg max α ∈ Null( D ν K YZ ): α (cid:62) K α ≤ 1 h 2 ( α ) .
Let V ∈ R
( m + n ) × p denote a matrix of p orthonormal columns that spans the nullspace, then α ∗ ∅ = V β ∗ , where β ∗ = arg max β ∈ R p : β (cid:62) V (cid:62) KV β ≤ 1 (cid:107) D 12 µ K XZ V β (cid:107) 22 .
Overall, the one-sided max-sliced kernel Bures distance can be computed using a combination of a line search for γ ∈ (0 , 1] and checking the solution in the nullspace as described in Algorithm 1.
Input: { ( x i , µ i ) } mi =1 , { ( y i , ν i ) } ni =1 , κ ( · , · ) , τ ⊆ { 1 , m + n } 1 z i = (cid:26) x i i ≤ m, y i − m i + 1 ≤ i ≤ m + n 2 D 12 µ K XZ =
[ √ µ i κ ( x i , z τ j )] m, | τ | i =1 ,j = 3 D 12 ν K Y Z =
[ √ ν i κ ( y i , z τ j )] n, | τ | i =1 ,j = 4 A = K (cid:62) XZ D µ K XZ 5 B
=
K
(cid:62) Y Z D ν K Y
Z 6 K
=
[ κ ( z τ i , z τ j )] m, | τ | i =1 ,j = 7 α (cid:63) =
O NE S IDED M AX S LICED K ERNEL B URES ( A , B , K ) 8 ω µ>ν ( z ) :
z (cid:55)→ (cid:80) | τ | i =
1 α (cid:63)i κ ( z, z τ i )
Output: ω
µ>ν 1 O NE S IDED M AX S LICED K ERNEL B URES ( A , B , K ) α γ : γ (cid:55)→ arg max α : α (cid:62) K α ≤ 1 α (cid:62) ( γ A − B )
α γ (cid:63) = arg max 0 <γ ≤ 1 (cid:113) α (cid:62) γ
A α γ − (cid:113) α (cid:62)
γ B α γ α ∅ = arg max α ∈ Null( B ): α (cid:62) K α ≤ 1 α (cid:62)
A α if (cid:113) α (cid:62) ∅
A α ∅ > (cid:113) α (cid:62) γ (cid:63)
A α γ (cid:63) − (cid:113) α (cid:62) γ (cid:63) B α γ (cid:63) then α (cid:63) =
α ∅ else α (cid:63) =
α γ (cid:63) end return α (cid:63)
A.4. M AX -S LICED K ERNEL W ASSERSTEIN
The empirical version of max-sliced kernel Wasserstein divergence can be expressed in terms of either equation 24 or equation 26,
[Equation]
where P ˆ µ, ˆ ν = { P ∈ [0 , 1] m × n | P1 n = µ , P (cid:62) 1 m = ν } is a transportation polytope.
The empirical distribution functions are F ω, ˆ µ ( w )
=
(cid:80) mi =1 µ
i I ω ( x i ) ≤ w and F ω, ˆ ν ( w ) =
(cid:80) ni
=
ν i
I ω
( y i ) ≤ w with inverses F − 1 ω, ˆ µ ( q ) =
min { ω ( x i ) , i ∈ { 1 , . . .
, m } : F ω, ˆ µ ( ω ( x i )) ≥ q } and F − 1 ω, ˆ ν ( q ) = min { ω ( y i ) , i ∈ { 1 , . . . , n } : F ω, ˆ ν ( ω ( y i )) ≥ q } .
For ﬁxed ω , the optimal transport plan ´ P is based on the sorted values ω ( x (1) ) ≤ ω ( x (2) ) ≤ · · · ≤ ω ( x ( m ) ) and ω ( y (1) ) ≤ ω ( y (2) ) ≤ · · · ≤ ω ( y ( m ) ) .
Denote the sorted values as vectors ´ ω X
and ´ ω Y , with [ ´ ω X ] i = ω
( x ( i ) ) and [ ´ ω Y ] i = ω
( y ( i ) ) , and denote by ´ µ and ´ ν the corresponding permuted versions of µ and ν .
In the case of equal samples sizes of uniform measure, m = n and µ = 1 m 1 m
and ν =n 1 n , elements in P ˆ µ, ˆ ν are scaled elements in the Birkhoff polytope (the set of doubly stochastic matrices), and the solution to the linear program is a permutation and max- W H p (ˆ µ, ˆ ν )
= sup ω ∈H : (cid:107) ω (cid:107) 2 ≤ 1 (cid:110)(cid:2)(cid:80) mi =1 (cid:12)(cid:12)
ω
[Equation]
x ( i ) )
[Equation]
ω
[Equation]
y ( i ) )
[Equation]
p (cid:3)
[Equation]
=
[Equation]
ω X − ´ ω Y (cid:107) p (cid:111) , where (cid:107)·(cid:107) p denotes the (cid:96) p norm in m dimensions.
As in the continuous case, the discrete transport plan between the sorted measures has the cumulative distribution ´ G ∈ [0 , 1] m × n with ´ G i,j = min { (cid:80) ik =1 ´ µ k , (cid:80) jk
=1 ´ ν k } .
The optimal transport plan between the sorted measures is given by taking the ﬁrst difference over both rows and columns of ´ G , ´ P 1 , 1 = ´ G 1 , 1 , ´ P 1 ,j = ´ G 1 ,j − ´ G 1 ,j − 1 , ´ P i, 1 = ´
G i, 1 − ´ G i − 1 , 1 , and ´
P i,j = ´ P i,j − ´ P i − 1 ,j − ´ P i,j − 1 .
Using the sorted witness function evaluations the distance can be written as
[Equation] [Equation]
We now turn our attention to the optimization of the function ω parametrized in terms of α , ω ( · ) = (cid:80) m + n i =
α i κ ( · , z i ) .
For arbitrary sample sizes with p = 2 , the max-sliced kernel W2 distance is
[Equation]
with unsorted values ω X
= [ ω ( x 1 ) , . . .
, ω ( x m )] (cid:62) =
K XZ α and ω Y =
[ ω ( y 1 ) , . .
ω ( y n )] (cid:62) =
K Y Z α .
The max-sliced kernel W2 distance can be expressed in terms of equation 45 or equation 48:
[Equation]
where Q P = K
[Equation]
XZ D µ K XZ
+ K
[Equation]
Y Z D ν K Y Z − K (cid:62) XZ PK Y Z − K
(cid:62) Y Z P (cid:62) K XZ is a symmetric matrix positive semideﬁnite matrix.
(This can be seen since the objective is greater than or equal to zero for all choices of α .)
For ﬁxed P , equation 49 is a convex maximization, that can be solved as a generalized eigenvalue problem.
For ﬁxed α , the optimization in terms of P is a linear program with the solution detailed above.
However, when maximizing with respect to α , P is a matrix-valued function of α .
To ﬁnd a local maximum for α we use the unconstrained optimization equationand perform gradient ascent with respect to α , wherein each iteration we compute the optimal transport plan.
This approach is also used in computing the generalized max-sliced Wasserstein distance Kolouri et al.
A. C OMPUTING THE M AX -S LICED D IVERGENCES FOR THE L INEAR C ASE
The linear case follows from the kernel case with some further simpliﬁcation.
The objectives of the one-sided max-sliced Bures divergences are each a difference of convex functions, whose stationary points are maximum eigenvalue problems, and correspond to reweighted versions of the one-sided max-sliced TV divergence.
Assuming the covariance matrices are strictly positive deﬁnite
[Equation]
w (cid:62) γ ρ X w γ − (cid:113) w (cid:62) γ ρ Y w γ (cid:12)(cid:12)(cid:12) , where w γ = arg max w : (cid:107) w (cid:107) 2 ≤ 1 w (cid:62) ( γ ρ X − ρ Y ) w for the one-sided case.
If the matrices are singular, then cases where w is in the nullspace must be checked.
Without loss of generality, the algorithm to obtain the optimal slice for the one-sided max-sliced Bures divergence is described in Algorithm 2.
Input: ρ X , ρ
Y ∈ R d × d 1 w γ : γ (cid:55)→ arg max w : (cid:107) w (cid:107) 2 ≤ 1 w (cid:62) ( γ ρ X − ρ Y ) w 2 γ (cid:63) = arg max 0 <γ ≤ 1 (cid:113) w (cid:62) γ ρ X w γ − (cid:113) w (cid:62) γ ρ Y w γ 3 if then w µ>ν = w γ
5 else v = arg max w ∈ Null( ρ Y ): (cid:107) w (cid:107) 2 ≤ 1 w (cid:62) ρ X w if (cid:112) v (cid:62) ρ X v > (cid:113) w (cid:62) γ (cid:63)
X w γ (cid:63) − (cid:113) w (cid:62) γ (cid:63) then w
= v else
>
end 12 end
Output: w µ>ν
As an alternative to Algorithm 2, ﬁrst-order algorithms can be applied, nevertheless, obtaining a global optimal cannot be guaranteed easily in this case.
To make the objective differentiable, the square root, which is non-differentiable at 0, should be smoothed √· ≈ √· + (cid:15) 2 (e.g., (cid:15) 2 = 0 . 01 ).
A. S AMPLE - BASED M AX -S LICED B URES D ISTANCE IS A R ELAXATION OF THE M AX -S LICED W ASSERSTEIN -2 D ISTANCE
We show that the max-sliced Bures distance is a relaxation of the max-sliced W2 distance.
Let ρ =
ZZ (cid:62) with Z ∈ R
d × p denote a strictly positive deﬁnite matrix, such that (cid:112) w (cid:62)
[Equation]
w > 0 ,
[Equation]
w
[Equation]
w
[Equation]
√ w
[Equation]
ZZ
[Equation]
w
[Equation]
max
[Equation]
S p
[Equation] [Equation]
θ, Z (cid:62) w (cid:105) , where θ ∗ = arg max θ ∈ S p
[Equation] [Equation]
θ, Z
[Equation]
w (cid:105) = 1 √
[Equation]
ZZ
[Equation]
Z
[Equation]
w .
Using this form, when ρ X and ρ Y are non-singular, the one-sided sliced Bures can be expressed as
[Equation]
since ρ X = XD 12 µ ( XD 12 µ ) (cid:62) and ρ Y =
YD 12 ν ( YD 12 ν ) (cid:62) .
Squaring this quantity and taking the square root yields the objective of the max-sliced Bures distance in a similar form to the max-sliced W distance, max-
[Equation]
[Equation]
(50) where Q ˆ µ, ˆ ν = { D 12 µ Θ D 12 ν : Θ ∈ R m × n , (cid:107) Θ (cid:107) 1 ≤ 1 } .
This holds since
[Equation]
A. C OVARIATE S HIFT C ORRECTION A LGORITHMS
For covariate shift correction the goal is to minimize the divergence by adjusting the weights ν of the instances in one sample ˆ ν .
This results in the following convex optimization problem:
[Equation]
where
[Equation]
For ﬁxed α the function ˜ J ( · , α ) is a sum of a linear function and a convex function f ( · ) = −√· is convex since √· is concave.
J ( ν ) = max α : α (cid:62) K α ≤ 1 ˜ J ( ν , α ) is convex since maximizing over α preserves convexity (Nesterov, 2018, Theorem 3.1.8).
In the linear kernel case, the cost is
[Equation]
Again, J ( ν ) is convex since ˜ J ( · , w ) is convex for ﬁxed w .
In this case, the gradient has the intuitive form of
[Equation]
To solve this convex minimization over a probability simplex we apply the Frank-Wolfe (conditional gradient) algorithm Jaggi (2013) to iteratively adjust the weight of one instance ν ← (1 − γ ) ν
+ γ e i , where e i is an indicator vector,
i = arg min 1 ≤ j ≤ n [ ∇ ν J ( ν )]
j , and γ ∈ [0 , 1] is the stepsize.
A beneﬁt of the Frank-Wolfe scheme is that it requires only rank-1 updates of YD ν Y (cid:62) , which are needed for updating w .
A. A DDITIONAL E XPERIMENTS R ESULTS
We start by comparing the proposed max-sliced Bures distance to the max-sliced W2 distance for two-dimensional data.
We compare the ﬁxed-point algorithm for solving each one-sided max-sliced Bures divergence with gradient-based approaches for the max-sliced W2 distance using ADAM with parameters lr = 1e-3, beta1 = 0.9, beta2 = 0.999, epsilon = 1e-08 , capping the number of iterations at 1000 or until the change in the slice is minimal (cid:107) w − w old (cid:107) ∞ < 10 − 6 .
In two-dimensions, a near optimal slice can be obtained by a ﬁne grid search of the sliced Bures and sliced W2 distance as shown in Figure 7 for two zero-mean Gaussian distributions.
Figure 8 shows the cases for success rate of the gradient-based optimizations across 10 trials at varying sample sizes for 2- and 1000-dimensional zero-mean Gaussians.
The effect of the number of gradient iterations is reported in Figure 9.
For kernel-based divergences, maximum mean discrepancy (MMD) detects differences in the ﬁrst moments of the distributions in the RKHS.
Using uncentered second-moments, the kernel-based max-sliced Bures distance (MSB) may detect some of the same differences.
Figure 10 details witness function evaluations of each for six data sets generated from two-dimensional distributions, where a Gaussian kernel function is used.
Notably, the one-sided MSB divergences correspond to localized regions, which are not distributed outliers.
This is beneﬁcial for the ‘precision’ of the witness function, but the ‘recall’ of MMD is better.
This beneﬁt of this localization depends on the task.
A. C OVARIATE S HIFT D ETECTION WITH M AX - SLICED K ERNEL D IVERGENCES
We proceed to generalize the comparisons in Figure 4 on imbalanced samples on MNIST to the kernel case.
We use a Gaussian kernel κ σ with the parameter σ set as the median Euclidean distance
[Figure] [Figure]
[Figure]
{+[Table] (sample size)
 D i v e r gen c e max-sliced Wasserstein-2 {ADAM} max-sliced Bures {Eig.} max-sliced Bures {ADAM}   12 log 2 (sample size) 2 
 D i v e r gen c e max-sliced Wasserstein-2 {ADAM} max-sliced Bures {Eig.} max-sliced Bures {ADAM}   12 log 2 (sample size) 2 
+}
{+D i v e r gen c e max-sliced Wasserstein-2 {ADAM} max-sliced Bures {Eig.} max-sliced Bures {ADAM}   12 log 2 (sample size) 2 
 D i v e r gen c e max-sliced Wasserstein-2 {ADAM} max-sliced Bures {Eig.} max-sliced Bures {ADAM}   12 log 2 (sample size) 2 
 D i v e r gen c e max-sliced Wasserstein-2 {ADAM} max-sliced Bures {Eig.} max-sliced Bures {ADAM}+}
0-   -2 0MMD
MSB
MMD
MSB
MMD
MSB
MMD
MSB
MMD
MSB
MMD
MSB in the pooled sample. To ease computation for large-sample sizes, we let τ be a random subset of the pooled samples { x i } mi =1 ∪{ y i } ni =1 of size l = min { 500 , m + n } . In this case, the max-sliced Fr´echet refers to max L - D H GW , which is equal to the square root of the sum of square of MMD and the square of the max-sliced Bures using centered kernels, as in equation 23.
The kernel-based max-sliced W distance should be an upper bound of the max-sliced Fr´echet.
However, in practice the optimal slice (witness function) may not be obtained.
[Figure]
We now compare the proposed kernel-based max-sliced divergences to existing baselines.
A primary baseline for this task is to train a logistic regression model with kernel basis functions to distinguish the two samples, and then use the probability estimates of the instances as the witness function evaluations ω ( x ) =
Pr ( H 0 | X = x ) = 1 − Pr
( H 1 | X = x ) , where H 0 :
X ∼ ˆ ν
and H 1 : X ∼ ˆ µ .
As additional baselines we also tested three methods for importance reweighting and density ratio estimation: kernel mean matching (KMM) (Huang et al., 2007), least-squares importance estimation (uLSIF) (Kanamori et al., 2009), and relative density-ratio estimation (RuLSIF)
(Yamada et al., 2011), but all methods were outperformed by logistic regression with kernel bases.
We also compare with kernel Fischer discriminant analysis (KFDA) (Mika et al., 1999), and the linear cases of max-sliced Wasserstein-2 distance, its ﬁrst moment approximation, max-sliced Bures, and logistic regression.
For all kernel methods, a Gaussian kernel κ σ is used with the parameter σ set as the median Euclidean distance in the pooled sample.
Using the MNIST data set again, we test three scenarios of covariate shift.
For each, one sample has a mismatched probability for one class l ∈ { 0 , . . .
, 9 } and the other sample has a balanced sample: (Scenario 1) ˆ µ is balanced and ˆ ν is missing l ; (Scenario 2) ˆ µ is imbalanced with l only appearing in 2% of the cases, compared to 10.¯8% for the other classes and ˆ ν is balanced; (Scenario 3) ˆ µ is balanced and ˆ ν consists of only images from l .
In each case, ˆ µ is a sample of 500 images from the training set and ˆ ν is a sample of 500 images from the test set.
A threshold-free way to assess covariate shift detection is to use the area-under-the-curve (AUC) of the receiver operator curve (ROC), where positive instances correspond to class l .
For some methods, the witness function (or its magnitude) may be ambiguous in sign, i.e., the values may be high (or large) for either the underor over-sampled instances (namely, max-sliced W2).
To be generous, on each run we choose the ordering with the highest AUC.
The results are reported in Table 3.
The other baselines KMM, uLSIF, and RuLSIF are not shown (their AUC scores across the scenarios are worse than the logistic regression with kernel baseline).
In a separate set of runs we also compute the realism scores (Kynk¨a¨anniemi et al., 2019) with k = 3 where ˆ ν is considered the real set, and ˆ µ are synthetic, to prioritize instances; results for the three scenarios are 0.75 ± 0.11, 0.67 ± 0.12, and 0.94 ± 0.04, which is better than linear logistic regression and KFDA, but far worse than the kernel logistic regression baseline.
[Table]
A. C OVARIATE S HIFT C ORRECTION FOR C LASS -C ONDITIONAL S UBSAMPLING
Figure 12 shows 20 synthetic images—generated by AutoGAN trained on CIFAR10 ( n =50,000)— with the largest weights after reweighting in order to minimize the max-sliced Bures distance to the subset of training images for each class separately ( m =5,000).
Computing the max-sliced Bures distance with the entire training set of 50,000 points is tractable, since it does not depend on the sample size.
The realism scores of the selected images have a median and range of 0.96 (0.63–1.34).
Figure 13 shows the same but based on the weights optimized by using the W2 distance with the mini-batch optimization.
The realism scores of the selected images have a median and range of 1.
Figure 14 shows the same but based on the weights optimized by using the max-sliced
Wasserstein-2 distance with the mini-batch optimization.
The realism scores of the selected images have a median and range of 0.97 (0.65–1.25).
Finally, Figure 15 shows the synthetic images selected for having the highest realism scores; notably this set lacks class correspondence.
The optimizations in all three cases use the Frank-Wolfe algorithm
(Jaggi, 2013) with simplex constraints.
The default step-size schedule γ = 2 k +2 and the same stopping criterion is used
[Equation]
| < 10 − 3 .
This yields roughly the same number of iterations for each method.
The optimization starts from a uniform weighting, which means only ∼ 2000 instances are actually reweighted.
The Fr ´ echet Inception distances after reweighting are detailed in Table 4.
Based on the quantitative and qualitative results it appears that the W2 distance with mini-batch approximation weights high-quality synthetic images the most, but the diversity of the highly weighted instances may not capture the full distribution for a class.
In this regard, the max-sliced Bures better captures the diversity of the class
[Table]