Understanding the Variance Collapse of SVGD in High Dimensions

1 INTRODUCTION . A typical challenge in Bayesian learning is to efficiently and accurately learn a complex posterior distribution . Markov Chain Monte Carlo ( Robert and Casella , 2013 ) provides asymptotically accurate samples , but simulating the chain till convergence can be time-consuming . On the other hand , variational inference ( Wainwright et al. , 2008 ) approximates the intractable target distribution with parametric families by minimizing the KL divergence , which can be more computationally efficient . Among a plethora of variational approximations , particle-based inference algorithms are particularly compelling as they bypass the inherent distributional assumptions on the variational family . As a noticeable example , Stein Variational Gradient Descent ( SVGD ) ( Liu and Wang , 2016 ) is a deterministic particle-based inference algorithm that iteratively transports the particles by the functional gradient of KL divergence in the reproducing kernel Hilbert space ( RKHS ) . The functional gradient takes the form of a kernelized Stein ’ s operator and only requires access to the unnormalized target density . Despite the empirical successes of SVGD ( Liu , 2017 ; Haarnoja et al. , 2017 ; Kim et al. , 2018 ) , very few convergence guarantees have been established except in the mean-field limit , i.e. , the number of particles n → ∞ under fixed dimensionality d , for which it has been shown that the distribution of particles converges to the true invariant solution ( Liu , 2017 ; Lu et al. , 2019 ) . Moreover , it has been observed that as the problem dimensionality d becomes larger , the variance estimated by SVGD can be much smaller than variance of the target distribution ( Zhuo et al. , 2017 ) . This observation , which we refer to as the variance collapse phenomenon ( see Figure 1 ) , is highly undesirable for practitioners due to two reasons : ( i ) underestimating the variance leads to failure in explaining the uncertainty of model predictions , which is a key benefit of being Bayesian ; ( ii ) modern Bayesian inference problems are usually high-dimensional ; for instance , training Bayesian neural networks ( BNNs ) ( MacKay , 1992 ) requires inferring the posterior distribution of network weights , which could be more than millions of dimensions in real-world problems ( Krizhevsky et al. , 2012 ) . Our Contribution . We study the algorithmic bias of SVGD that leads to variance collapse in high dimensions . We focus on the commonly-used Euclidean distance kernel ( e.g. , Gaussian RBF kernel ) , and provide : ( i ) qualitative understanding on the pitfall of SVGD by comparing it to another interacting particle algorithm termed MMD-descent ; ( ii ) quantitative understanding on the variance underestimation in learning simple target distributions . Our findings can be summarized as follows . • In Section 3 we connect SVGD with MMD-descent , a kernel-based particle inference algorithm that performs maximum mean discrepancy ( MMD ) minimization , and empirically show that despite their similar updates , MMD-descent does not collapse the variance in high dimensions . • In Section 4 we identify the log derivative driving force as the problematic term in SVGD , and design experiments to illustrate that the combination of this driving force and the deterministic bias , i.e . the absence of particle resampling , prevents accurate variance estimation in high dimensions . • In Section 5 we argue that the proportional asymptotic limit : n , d → ∞ , d/n → γ ∈ ( 0 , ∞ ) , is more relevant to understanding the variance collapse phenomenon . As an example , we derive the exact dimension-averaged variance of SVGD and MMD-descent in learning Gaussian distribution in the proportional limit , under certain concentration assumption which we empirically verify . Our analysis confirms that variance estimated by SVGD scales inversely with the dimensionality . 2 BACKGROUND . 2.1 INTEGRAL PROBABILITY METRIC . To measure the “ closeness ” between distributions on X ⊆ Rd , one may consider the maximum discrepancy between the target p and sample distribution q over test functions F : DF ( p , q ) : = suph∈F Eq [ h ( x ) ] − Ep [ h ( y ) ] , known as the integral probability metric ( IPM ) ( Müller , 1997 ) . If F is a unit ball in the reproducing kernel Hilbert space ( RKHS ) H , the resulting DF is termed the maximum mean discrepancy ( MMD ) ( Gretton et al. , 2012 ) , and its squared value can be evaluated as : MMD2 ( p , q ) = Ex , x′ [ k ( x , x′ ) ] + Ey , y′k ( y , y′ ) − 2Ex , y [ k ( x , y ) ] ( 1 ) where x , x′ ∼ p , y , y′ ∼ q , and the kernel k : Rd × Rd → R satisfies E √ k ( x , x ) < ∞ . Stein ’ s Discrepancy . When integration under the target p is intractable , Stein ’ s method ( Stein et al. , 1972 ) can be used to construct zero-mean test functions w.r.t . p. For differentiable h in the Stein Class of p , i.e. , ∫ x ∇x ( h ( x ) p ( x ) ) dx = 0 , the Stein ’ s discrepancy ( Gorham and Mackey , 2015 ) is given as DSteinF ( p , q ) = suph∈F Eq [ h ( x ) > ∇x log p ( x ) +∇ > xh ( x ) ] . Note that the Stein ’ s discrepancy only involves the score of p and thus the normalization constant is not required . When h is restricted in the product RKHS Hd with inner product 〈f , g〉Hd =∑d i=1〈fi , gi〉H , the corresponding maximum discrepancy , known as kernel Stein discrepancy ( KSD ) , can be estimated efficiently from samples ( Liu et al. , 2016 ; Chwialkowski et al. , 2016 ) . 2.2 ( DETERMINISTIC ) PARTICLE INFERENCE ALGORITHMS . We now consider the approximation of p using a set of particles constructed as follows : starting from initial particles X = { xi } ni=1 , we iteratively optimize their positions using the deterministic update : xi = xi + η∆ ( xi ) , where η is the step size , and ∆ ( · ) : Rd → Rd represents the update direction . Motivated by the SVGD algorithm , we focus on the setting where the update ∆ is deterministic , in which case particles may converge to a deterministic fixed point . Note that there are other particlebased inference algorithms that do not fit into this description , such as sequential Monte Carlo ( Doucet et al. , 2001 ) , and stochastic variants of particle gradient method ( Chen et al. , 2018a ; Gallego and Insua , 2018 ) , which our current analysis does not cover ( see Appendix A.4 for more discussion ) . Stein Variational Gradient Descent . SVGD constructs the update direction ∆ as the optimal perturbation in the RKHS that decreases Kullback-Leibler divergence . In particular , constrain ∆ in RKHS unit ball , and take q = 1n ∑n i=1 δxi , then the update for each particle x can be computed as : ∆SVGD ( x ) = Ex′∼q [ k ( x′ , x ) ∇x′ log p ( x′ ) ︸ ︷︷ ︸ driving force +∇x′k ( x′ , x ) ︸ ︷︷ ︸ repulsive force ] : = 1 n n∑ i=1 [ S1 ( xi , x ) + S2 ( xi , x ) ] . ( 2 ) Intuitively , the log derivative term in the update rule S1 ( xi , x ) : = k ( xi , x ) ∇xi log p ( xi ) corresponds to a driving force that guides particles towards high likelihood regions , whereas the kernel derivative term S2 ( xi , x ) : = ∇xik ( xi , x ) provides a repulsive force to prevent the particles from collapsing into the mode . Typically-used kernels in SVGD include the Gaussian RBF kernel k ( x , x′ ) = exp ( −‖x− x′‖22 /2σ2 ) ( Liu and Wang , 2016 ; Zhuo et al. , 2017 ) , the linear kernel k ( x , x′ ) = x > x′ + 1 ( Liu and Wang , 2018 ) , and the inverse multi-quadratic ( IMQ ) kernel k ( x , x′ ) = 1/ √ 1 + ‖x− x′‖22/ ( 2σ2 ) ( Gorham and Mackey , 2017 ) . 3 CONNECTING SVGD WITH MMD MINIMIZATION . In this section , we introduce another particle inference algorithm termed MMD-descent , whose update rule closely relates to that of SVGD . Despite the similarity , we empirically observe MMD-descent accurately estimates the target variance independent of the dimensionality . By analyzing the similarity of the update rules and contrast in the algorithmic performances , we identify factors that lead to the variance collapse of SVGD . In this work we focus on the following class of kernel functions . Definition 1 ( Euclidean Distance Kernel ) . A positive semi-definite kernel is referred to as a Euclidean distance kernel if it can be represented as : k ( x , y ) = f ( ‖x− y‖22/σ2 ) . For example , the Gaussian RBF kernel and IMQ kernel both satisfy this definition . Note that σ is the tunable bandwidth that usually scales with the distance between particles . We now construct a kernel-based particle inference algorithm termed MMD-descent , which draws inspiration from the kernel herding algorithm ( Welling , 2009 ) and particle gradient descent ( Chizat and Bach , 2018 ) . Instead of greedily reducing the MMD by adding one particle a time , we consider an update rule similar to SVGD that transport all particles together to approximate the target distribution , by minimizing the MMD between the particles and the target via gradient descent : ∆MMD ( x ) = Ey∼p [ ∇xk ( x , y ) ︸ ︷︷ ︸ driving force ] + Ex′∼q [ −∇xk ( x , x′ ) ︸ ︷︷ ︸ repulsive force ] . This update can be seen as the finite-particle discretization of the MMD gradient flow ( Arbel et al. , 2019 ) . For Euclidean distance kernels we have −∇xk ( x , x′ ) = ∇x′k ( x , x′ ) = S2 ( x′ , x ) ; hence , MMD-descent and SVGD share the same repulsive force . Furthermore , when k is in the Stein class of p , the driving force of MMD-descent can be connected SVGD via a simple integration by parts : Ey∼p [ ∇xk ( x , y ) ] = −Ey∼p [ ∇yk ( x , y ) ] = Ey∼p [ k ( x , y ) ∇y log p ( y ) ] = Ey∼p [ S1 ( y , x ) ] . ( 3 ) Therefore , the MMD-descent update for a set of particles { xi } ni=1 can be equivalently written as : ∆MMD ( x ) = Ey∼p [ S1 ( y , x ) ] + 1 n n∑ i=1 S2 ( xi , x ) . ( 4 ) We remark that MMD-descent is not a practical algorithm due to the required integration under p , which is typically intractable . Instead , the purpose of introducing MMD-descent is to compare the update with SVGD and understand the cause of the variance collapse phenomenon . SVGD vs. MMD-descent . By comparing SVGD ( 2 ) and MMD-descent ( 4 ) , we observe that : • SVGD and MMD-descent have identical repulsive force . • In MMD-descent , the driving force is integrated under the target distribution p , whereas in SVGD the expectation is under the current particle distribution q . At the infinite particle limit , p = q is a fixed point for both update rules . In fact , under certain conditions , uniqueness of the fixed point has been established for both updates in this mean-field limit ( Lu et al. , 2019 ; Arbel et al. , 2019 ) . However , in the practical setting where the particle size is not significantly greater than the dimensionality , it is not clear if the two algorithms : ( i ) provide a reasonable “ approximation ” to the target distribution ; ( ii ) converge to similar fixed points . Given the analogous updates , one natural question to ask is : do SVGD and MMD-descent reliably estimate the target variance in high dimensions , and do they give similar solutions ? The empirical answer is in the negative : SVGD and MMD-descent converge to completely different solutions : SVGD is known to underestimate the marginal variance in high dimensions , even for simple Gaussian targets ( Zhuo et al. , 2017 ) . In contrast , MMD-descent does not exhibit the same issue empirically . Figure 1 illustrates this discrepancy in a Bayesian neural network ( BNN ) experiment . We consider a two-hidden-layer BNN with 100 hidden units each , so that inference in the weight space is a high-dimensional problem . We obtain ( approximate ) target samples via Hamiltonian Monte Carlo ( HMC ) ( Neal et al. , 2011 ) , and then initialize 100 particles from the target to be optimized by SVGD or MMD-descent with the Gaussian RBF kernel ( for detailed setup see Appendix D ) . Observe that although we used the same number of particles for SVGD and MMD-descent , SVGD severely underestimates the variance , whereas MMD-descent generates diverse predictions similar to HMC .

Summary: This paper provides an understanding of the variance collapse phenomenon of SVGD. The paper first (1) introduces the reader to the most important concepts and phenomena, then (2) gives an explanation for why this problem occurs, thanks to a comparison with an accurate (yet computationally intensive) algorithm they call MMD-descent. Finally (3) the paper shows how to fix SVGD with damping. The paper provides experiments and theory, nicely combined.