Sliced Kernelized Stein Discrepancy

1 INTRODUCTION . Discrepancy measures for quantifying differences between two probability distributions play key roles in statistics and machine learning . Among many existing discrepancy measures , Stein discrepancy ( SD ) is unique in that it only requires samples from one distribution and the score function ( i.e . the gradient up to a multiplicative constant ) from the other ( Gorham & Mackey , 2015 ) . SD , a special case of integral probability metric ( IPM ) ( Sriperumbudur et al. , 2009 ) , requires finding an optimal test function within a given function family . This optimum is analytic when a reproducing kernel Hilbert space ( RKHS ) is used as the test function family , and the corresponding SD is named kernelized Stein discrepancy ( KSD ) ( Liu et al. , 2016 ; Chwialkowski et al. , 2016 ) . Variants of SDs have been widely used in both Goodness-of-fit ( GOF ) tests ( Liu et al. , 2016 ; Chwialkowski et al. , 2016 ) and model learning ( Liu & Feng , 2016 ; Grathwohl et al. , 2020 ; Hu et al. , 2018 ; Liu & Wang , 2016 ) . Although theoretically elegant , KSD , especially with RBF kernel , suffers from the ” curseof-dimensionality ” issue , which leads to significant deterioration of test power in GOF tests ( Chwialkowski et al. , 2016 ; Huggins & Mackey , 2018 ) and mode collapse in particle inference ( Zhuo et al. , 2017 ; Wang et al. , 2018 ) . A few attempts have been made to address this problem , however , they either are limited to specific applications with strong assumptions ( Zhuo et al. , 2017 ; Chen & Ghattas , 2020 ; Wang et al. , 2018 ) or require significant approximations ( Singhal et al. , 2019 ) . As an alternative , in this work we present our solution to this issue by adopting the idea of “ slicing ” . Here the key idea is to project the score function and test inputs onto multiple one dimensional slicing directions , resulting in a variant of SD that only requires to work with one-dimensional inputs for the test functions . Specifically , our contributions are as follows . • We propose a novel theoretically validated family of discrepancies called sliced Stein discrepancy ( SSD ) , along with its scalable variant called max sliced kernelized Stein discrepancy ( maxSKSD ) using kernel tricks and the optimal test directions . • A GOF test is derived based on an unbiased estimator of maxSKSD with optimal test directions . MaxSKSD achieves superior performance on benchmark problems and restricted Boltzmann machine models ( Liu et al. , 2016 ; Huggins & Mackey , 2018 ) . ∗Work done at Microsoft Research Cambridge • We evaluate the maxSKSD in model learning by two schemes . First , we train an independent component analysis ( ICA ) model in high dimensions by directly minimising maxSKSD , which results in faster convergence compared to baselines ( Grathwohl et al. , 2020 ) . Further , we propose a particle inference algorithm based on maxSKSD called the sliced Stein variational gradient descent ( S-SVGD ) as a novel variant of the original SVGD ( Liu & Wang , 2016 ) . It alleviates the posterior collapse of SVGD when applied to training variational autoencoders ( Kingma & Welling , 2013 ; Rezende et al. , 2014 ) . 2 BACKGROUND . 2.1 KERNELIZED STEIN DISCREPANCY . For two probability distributions p and q supported on X ⊆ RD with continuous differentiable densities p ( x ) and q ( x ) , we define the score sp ( x ) = ∇x log p ( x ) and sq ( x ) accordingly . For a test function f : X → RD , the Stein operator is defined as Apf ( x ) = sp ( x ) T f ( x ) +∇Txf ( x ) . ( 1 ) For a function f0 : RD → R , the Stein class Fq of q is defined as the set of functions satisfying Stein ’ s identity ( Stein et al. , 1972 ) : Eq [ sq ( x ) f0 ( x ) +∇xf0 ( x ) ] = 0 . This can be generalized to a vector function f : RD → RD where f = [ f1 ( x ) , . . . , fD ( x ) ] T by letting fi belongs to the Stein class of q for each i ∈ D. Then the Stein discrepancy ( Liu et al. , 2016 ; Gorham & Mackey , 2015 ) is defined as D ( q , p ) = sup f∈Fq Eq [ Apf ( x ) ] = sup f∈Fq Eq [ ( sp ( x ) − sq ( x ) ) T f ( x ) ] . ( 2 ) When Fq is sufficiently rich , and q vanishes at the boundary of X , the supremum is obtained at f∗ ( x ) ∝ sp ( x ) − sq ( x ) with some mild regularity conditions on f ( Hu et al. , 2018 ) . Thus , the Stein discrepancy focuses on the score difference of p and q. Kernelized Stein discrepancy ( KSD ) ( Liu et al. , 2016 ; Chwialkowski et al. , 2016 ) restricts the test functions to be in a D-dimensional RKHS HD with kernel k to obtain an analytic form . By defining up ( x , x′ ) = sp ( x ) Tsp ( x′ ) k ( x , x′ ) + sp ( x ) T∇x′k ( x , x′ ) + sp ( x′ ) T∇xk ( x , x′ ) + Tr ( ∇x , x′k ( x , x′ ) ) the analytic form of KSD is : D2 ( q , p ) = ( sup f∈HD , ||f ||HD≤1 Eq [ Apf ( x ) ] ) 2 = Eq ( x ) q ( x′ ) [ up ( x , x′ ) ] . ( 3 ) 2.2 STEIN VARIATIONAL GRADIENT DESCENT . Although SD and KSD can be directly minimized for variational inference ( VI ) ( Ranganath et al. , 2016 ; Liu & Feng , 2016 ; Feng et al. , 2017 ) , Liu & Wang ( 2016 ) alternatively proposed a novel particle inference algorithm called Stein variational gradient descent ( SVGD ) . It applies a sequence of deterministic transformations to a set of points such that each of mappings maximally decreases the Kullback-Leibler ( KL ) divergence from the particles ’ underlying distribution q to the target p. To be specific , we define the mapping T ( x ) : RD → RD as T ( x ) = x+ φ ( x ) whereφ characterises the perturbations . The result from Liu & Wang ( 2016 ) shows that the optimal perturbation inside the RKHS is exactly the optimal test function in KSD . Lemma 1 . ( Liu & Wang , 2016 ) Let T ( x ) = x + φ ( x ) and q [ T ] ( z ) be the density of z = T ( x ) when x ∼ q ( x ) . If the perturbation φ is in the RKHSHD and ||φ||HD ≤ D ( q , p ) , then the steepest descent directions φ∗q , p is φ∗q , p ( · ) = Eq [ ∇x log p ( x ) k ( x , · ) +∇xk ( x , · ) ] ( 4 ) and ∇ KL [ q [ T ] ||p ] | =0 = −D2 ( q , p ) . The first term in Eq . ( 4 ) is called drift , which drives the particles towards a mode of p. The second term controls the repulsive force , which spreads the particles around the mode . When particles stop moving , the KL decrease magnitude D2 ( q , p ) is 0 , which means the KSD is zero and p = q a.e . 3 SLICED KERNELIZED STEIN DISCREPANCY . We propose the sliced Stein discrepancy ( SSD ) and kernelized version named maxSKSD . Theoretically , we prove their correctness as discrepancy measures . Methodology-wise , we apply maxSKSD to GOF tests , and develop two ways for model learning . 3.1 SLICED STEIN DISCREPANCY . Before moving to the details , we give a brief overview of the intuition on how to tackle the curse-offimensionality issue of SD ( The right figure of Figure 1 ) . For detailed explanation , refer to appendix B.1 . This issue of Stein discrepancy ( Eq.2 ) comes from two sources : the score function sp ( x ) and the test function f ( x ) defined on X ⊂ RD . First , we notice that comparing sp and sq is equivalent to comparing projected score srp = s T p r and s r q for all r ∈ SD−1 on an hyper-sphere ( Green square in Figure 1 ( Right ) ) . This operation reduces the test function ’ s output from RD to R ( Green circle in Figure 1 ( Right ) ) . However , its input dimension is not affected . Reducing the input dimension of test functions is non-trivial , as directly removing input dimensions results in the test power decrease . This is because less information is accessed by the test function ( see examples in appendix B.1 ) . Our solution to this problem uses Radon transform which is inspired by CT-scans . It projects the original test function f ( x ) in Stein discrepancy ( Eq . 2 ) ( as an RD → R mapping ) to a group of R → R functions along a set of directions ( g ∈ SD−1 ) . Then , this group of functions are used as the new test functions to define the proposed discrepancy . The invertibility of Radon transform ensures that testing with input in the original space RD is equivalent to the test using a group of low dimensional functions with input in R. Thus , the above two steps not only reduce the dimensions of the test function ’ s output and input , but also maintain the validity of the resulting discrepancy as each step is either equivalent or invertible . In detail , assume two distributions p and q supported on RD with differentiable densities p ( x ) and q ( x ) , and define the test functions f ( · ; r , g ) : RD → R such that f ( x ; r , g ) = frg ◦ hg ( x ) = frg ( x Tg ) , where hg ( · ) is the inner product with g and frg : R→ R. One should note that the r and g in f ( · ; r , g ) should not just be treated as parameters in a test function f . In fact , they are more like the index to indicate that for each pair of r , g , we need a new f ( · ; r , g ) , i.e . new frg , which is completely independent to other test functions . The proposed sliced Stein discrepancy ( SSD ) , defined using two uniform distributions pr ( r ) and pg ( g ) over the hypersphere SD−1 , is given by the following , with frg ∈ Fq meaning f ( · ; r , g ) ∈ Fq : S ( q , p ) = Epr , pg [ sup frg∈Fq Eq [ srp ( x ) frg ( xTg ) + rTg∇xT gfrg ( xTg ) ] ] . ( 5 ) We verify the proposed SSD is a valid discrepancy measure , namely , S ( q , p ) = 0 iff . q = p a.e . Theorem 1 . ( SSD Validity ) If assumptions 1-4 in appendix A are satisfied , then for two probability distributions p and q , S ( q , p ) ≥ 0 , and S ( q , p ) = 0 if and only if p = q a.e . Despite this attractive theoretical result , SSD is difficult to compute in practice . Specifically , the expectations over r and g can be approximated by Monte Carlo but this typically requires a very large number of samples in high dimensions ( Deshpande et al. , 2019 ) . We propose to relax such limitations by using only a finite number of slicing directions r from an orthogonal basis Or of RD , e.g . the standard basis of one-hot vectors , and the corresponding optimal test direction gr for each r. We call this variant maxSSD , which is defined as follows and validated in Corollary 1.1 : Smax ( q , p ) = ∑ r∈Or sup frgr∈Fq , gr∈SD−1 Eq [ srp ( x ) frgr ( xTgr ) + rTgr∇xT grfrgr ( xTgr ) ] . ( 6 ) Corollary 1.1 . ( maxSSD ) Assume the conditions in Theorem 1 , then Smax ( q , p ) = 0 iff . p = q a.e .

Summary: This paper tries to solve the curse-of-dimensionality problem of KSD and corresponding mode-collapse problem of SVGD by projecting both the input and output of test function onto 1D slices. By doing so, the paper proposes the new discrepancies called SSD and maxSKSD, and a new variant of SVGD called S-SVGD. Experiments on goodness-of-fit test (synthetic high-dim Gaussian & RBM) and model learning (ICA on synthetic data & amortized SVGD on MNIST) are reported in the main body of the paper.