Max-sliced Bures Distance for Interpreting Discrepancies

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 find the optimal slice , we detail a tractable algorithm that finds the global optimal slice and scales to large sample sizes . As the Bures distance quantifies 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 semidefinite 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 Fréchet 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 . 1 INTRODUCTION . 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 specific subsets that are not well matched between the distributions , specifically , 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 ; Rényi , 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 . ( 2016 ) . Witness functions need not rely on the density ratio . A wide class of divergences called integral probability metrics ( IPMs ) ( Müller , 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 first moments of the witness function evaluations . In these cases the optimal witness function ω ? ( · ) has a greater expectation in one distribution compared to the other distribution . An IPM between two measures µ and ν is expressed as supω∈F |EX∼µ [ ω ( X ) ] − EY∼ν [ ω ( 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 difficulties with computing the max-sliced Wasserstein distance for two samples . The first is that it is a saddlepoint optimization problem , whose objective evaluation requires sorting the samples . Previous work has sought to approximate it using a first moment approximation ( Deshpande et al. , 2019 ) or to use a finite number of steps of a local optimizer ( Kolouri et al. , 2019 ) , without any guarantee of obtaining an optimal witness function . Another difficulty 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 Wasserstein-2 distance 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 semidefinite 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.1 CONTRIBUTION . 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 ∣∣∣√EX∼µ [ ω2 ( X ) ] −√EY∼ν [ ω2 ( Y ) ] ∣∣∣ , 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 √ Eµ [ ω2 ( X ) ] − √ Eν [ ω2 ( Y ) ] , and ωµ < ν = arg maxω∈S √ Eν [ ω2 ( Y ) ] − √ 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 { xi } mi=1 , { yi } ni=1 can be sorted by the magnitude of the witness function evaluations.1 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 = Rd and the first or second moments distinguish the distributions , linear witness functions can be used S = { ω ( · ) = 〈· , w〉 : w ∈ Sd−1 } , where Sd−1 denotes the unit sphere in Rd . The optimal witness function for the one-sided max-sliced Bures divergence ωµ > ν ( · ) = 〈· , wµ > ν〉 coincides with the subspace with the greatest difference in RMS , wµ > ν = arg maxw∈Sd−1 √ w > E [ XX > ] w − √ w > E [ Y Y > ] 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 requiresO ( m logm ) at each function/gradient evaluation . Furthermore , the slice that maximizes the max-sliced Wasser- 1Four groups of ‘ witness points ’ ( top-K instances ) can be inspected to identify any discrepancies : ω2µ > ν ( xπ̀ ( 1 ) ) ≥ · · · ≥ ω2µ > ν ( xπ̀ ( K ) ) ︸ ︷︷ ︸ π̀ sorts { xi } mi=1 to reveal examples from µ̂ with large ω 2 µ > ν ̃ ω2µ > ν ( yσ̀ ( 1 ) ) ≥ · · · ≥ ω2µ > ν ( yσ̀ ( K ) ) ︸ ︷︷ ︸ σ̀ sorts { yi } ni=1 to find the examples from ν̂ with large ω 2 µ > ν , ( 1 ) ω2µ < ν ( xπ́ ( 1 ) ) ≥ · · · ≥ ω2µ < ν ( xπ́ ( K ) ) ︸ ︷︷ ︸ π́ sorts { xi } mi=1 to find examples from µ̂ with large ω 2 µ < ν ̃ ω2µ < ν ( yσ́ ( 1 ) ) ≥ · · · ≥ ω2µ < ν ( yσ́ ( K ) ) ︸ ︷︷ ︸ σ́ sorts { yi } ni=1 to find the examples from ν̂ with large ω 2 µ < ν , ( 2 ) where π̀ , π́ , σ̀ , σ́ denote permutations and ̃ and ̃ denote expected inequalities with a large difference . stein lacks an intrinsic ordering , and it is left to the user to determine whether instances from µ or ν have high or low values or magnitudes . As second-order moments may be insufficient 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 = { ω ( · ) = 〈φ ( · ) , ω〉H : ω , φ ( · ) ∈ H , 〈ω , ω〉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 Fréchet 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 Fréchet 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 . ( 2017 ) . 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 ) .

Summary: This paper studies a family of integral probability metric (IPM) divergence on Hilbert spaces. This family can be characterized by the choice of the witness function, and specific witness function may give rise to the Bures distance, the MMD, Wasserstein, as well as many sliced variants. While this family has been well understood for distributions on finite dimensional space, this paper extends this insight to distributions on (possibly infinite dimensional) Hilbert space. By leveraging the representer theorem, the paper provides the finite dimensional optimization problems that can be solved to estimate the divergence from samples. The power of the method is demonstrated on the covariate shift experiments.