Challenges and Opportunities in High Dimensional Variational Inference

1 Introduction . A great deal of progress has been made in black-box variational inference ( BBVI ) methods for Bayesian posterior approximation , but the interplay between the approximating family , divergence measure , gradient estimators and stochastic optimizer is non-trivial – and even more so for highdimensional posteriors [ 1 , 10 , 29 , 31 ] . While the main focus in the machine learning literature has ∗ Equal contribution . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . been on improving predictive accuracy , the choice of method components becomes even more critical when the goal is to obtain accurate summaries of the posterior itself . In this paper , we show that , while the choice of approximating family and divergence is often motivated by low-dimensional illustrations , the intuition from these examples do not necessarily carry over to higher-dimensional settings . By drawing a connection between importance sampling and the estimation of common divergences used in BBVI , we are able to develop a comprehensive framework for understanding the reliability of BBVI in terms of the pre-asymptotic behavior of the density ratio between the target and the approximate distribution . When this density ratio is heavy-tailed , even unbiased estimators exhibit a large bias with high probability , in addition to high variance . Such heavy tails occur when there is a mismatch between the typical sets of the approximating and target distributions . In higher dimensions , even over-dispersed distributions miss the typical set of the target [ 18 , 27 ] . Thus , as illustrated in Fig . 1 , the benefits of heavy-tailed approximate families and divergences favoring mass-covering diminish as dimensionality of the target distribution increases . Building on these insights , we make the following main contributions : 1 . We develop a conceptual and experimental framework for predicting and empirically evaluating the reliability of BBVI based on the choice of variational objective , approximating family , and target distribution . Our framework also incorporates the Pareto k diagnostic [ 27 ] as a simple and practical approach for obtaining empirical and conceptual insights into the pre-asymptotic convergence rates of estimators of common divergences and their gradients . 2 . We validate our framework through an extensive empirical study using simulated data and many commonly used real datasets with both Gaussian and non-Gaussian target distributions . We consider the exclusive and inclusive Kullback-Leibler ( KL ) divergences [ 4 , 21 , 24 ] , tail-adaptive f -divergence [ 29 ] , χ2 divergence [ 7 ] , and α-divergences [ 12 ] , and the resulting variational approximation for isotropic Gaussian and Student-t and normalising flow families . 3 . Based on our framework and numerical results , we provide justified recommendations on design choices for different scenarios , including low- to moderate-dimensional and high-dimensional posteriors . 2 Preliminaries and Background . Let p ( θ , Y ) be a joint distribution of a probabilistic model , where θ ∈ RD is a vector of model parameters and Y is the observed data . In Bayesian analysis , the posterior p ( θ ) : = p ( θ | Y ) = p ( θ , Y ) /p ( Y ) ( where p ( Y ) : = ∫ p ( θ , Y ) dθ ) is typically the object of interest , but most posterior summaries of interest are not accessible because the normalizing integral , in general , is intractable . Variational inference approximates the exact posterior p ( θ | Y ) using a distribution q ∈ Q from a family of tractable distributionsQ . The best approximation is determined by minimizing a divergence D ( p ‖ q ) , which measures the discrepancy between p and q : qλ∗ = arg min qλ∈Q D ( p ‖ q ) , ( 1 ) where λ ∈ RK is a vector parameterizing the variational family Q . Thus , the properties of the resulting approximation q are determined by the choice of variational family Q as well as the choice of divergence D. The family Q is often chosen such that quantities of interest ( e.g. , moments of q ) can be computed efficiently . For example , q can be used to compute Monte Carlo or importance sampling estimates of the quantities of interest . Let w ( θ ) : = p ( θ , Y ) /q ( θ ) denote the density ratio between the joint and approximate distributions . For a function φ : RD → R , the biased self-normalized importance sampling estimator for the posterior expectation Eθ∼p [ φ ( θ ) ] is given by Î ( φ ) : = S∑ s=1 w ( θs ) ∑S s′=1 w ( θs′ ) φ ( θs ) , where θ1 , . . . , θS ∼ q are independent . Using importance sampling can allow for computation of more accurate posterior summaries and to go beyond the limitations of the variational family . For example , it is possible to estimate the posterior covariance even when using a mean-field variational family . Pareto Smoothed Importance Sampling . Since importance sampling estimates can have very high variance , Pareto smoothed importance sampling ( PSIS ) can be used to substantially reduce the variance with small additional bias [ 27 ] . This procedure modifies and stabilises extreme importance ratios using a generalized Pareto distribution fit to the upper tail of the distribution of the ratios . Variational families . Let qλ ( θ ) be an approximating family parameterised by a K-dimensional vector λ ∈ RK for D-dimensional inputs θ ∈ RD . Typical choices of q include mean-field Gaussian and Student ’ s t families [ 3 , 14 ] , full and low rank Gaussians [ 15 , 22 ] , mixtures of exponential families [ 17 , 19 ] , and normalising flows [ 25 ] . We focus on the most popular mean-field and normalizing flow families . Mean-field families assume independence across the D dimensions : q ( θ ) = ∏D i=1 qi ( θi ) , where each qi typically belongs to some exponential family or other simple class of distributions . Normalising flows [ 1 ] provide more flexible families that can capture correlation and non-linear dependencies . A normalizing flow is defined via the transformation of a probability density through a sequence of invertible mappings . By composing several maps , a simple distribution such as a mean-field Gaussian can be transformed into a more complex distribution [ 25 ] . f -divergences . The most commonly used divergences are examples of f -divergences [ 28 ] . For a convex function f satisfying f ( 1 ) = 0 , the f -divergence is given by Df ( p ‖ q ) : = Eθ∼q [ f ( p ( θ | Y ) q ( θ ) ) ] . The exclusive Kullback-Leibler ( KL ) divergence corresponds to f ( w ) = − log ( w ) , the inclusive KL divergence corresponds to f ( w ) = w log ( w ) , the χ2 divergence corresponds to f ( w ) = ( w − 1 ) 2 , and the general α-divergences correspond to ( wα − w ) / { α ( α− 1 ) } . We also consider the adaptive f -divergence proposed by Wang et al . [ 29 ] . Loss estimation and stochastic optimization . In all the cases we consider , minimizing the f - divergence is equivalent to minimizing the loss function Lf ( p ‖ q ) : = Eθ∼q [ f ( w ( θ ) ) ] ( although , see Wan et al . [ 28 ] for a different approach ) . Let L ( λ ) : = Lf ( p ‖ qλ ) denote the loss as a function of the variational parameters λ . The loss and its gradient G ( λ ) : = ∇λL ( λ ) can both be approximated using , respectively , the Monte Carlo estimates L̂ ( λ ) = 1S ∑S s=1 f ( w ( θs ) ) and Ĝ ( λ ) = 1 S ∑S s=1 g ( θs ) , ( 2 ) where θ1 , . . . , θS are independent draws from qλ and g : RK → RK is an appropriate gradient-like function that depends on f and w. The two most popular gradient estimators in the literature are the score function and the reparameterization gradient estimator [ 20 , 30 ] . The score function gradient corresponds to g ( θ ) = { f ( w ( θ ) ) − w ( θ ) f ′ ( w ( θ ) ) } ∇λ log qλ ( θ ) . It is a general-purpose estimator that applies to both discrete and continuous distributions q , but it is known to suffer from high variance . When this estimator is used for the mass-covering divergences such as the inclusive KL and general α-divergences with α > 1 , the importance weights are usually replaced with selfnormalized importance weights w ( θs ) / ∑S i=1 w ( θi ) . The reparameterization gradient [ 20 ] requires expressing the distribution qλ as a deterministic transformation of a simpler base distribution r such that Tλ ( z ) ∼ qλ with z ∼ r. This allows writing an expectation with respect to qλ as an expectation over the simpler distribution r. The reparameterization estimator corresponds to using g ( zs ) = ∇λf ( w ( Tλ ( zs ) ) ) ( for zs ∼ r ) in place of g ( θ ) , where w implicitly depends on λ as well . In the case of the adaptive f -divergence , the importance weights w ( θ1 ) , . . . , w ( θS ) are sorted , and the gradients corresponding to each sample are then weighed by the empirical rank . The gradient estimates can be used in a stochastic gradient optimization scheme such that λt+1 ← λt + ηtĜ ( λt ) , ( 3 ) where ηt is the step size . In practice , more stable adaptive stochastic gradient optimisation methods such as RMSProp or Adam [ 9 , 13 ] , which smooth or normalize the noisy gradients , are often used . Numerous prior work have studied some of the challenges tied to optimizing these divergence measures under the presence of noisy gradient estimates [ 1 , 10 , 29 , 31 ] . Particularly , when dealing with mass-covering divergences , the gradient estimates can become so noisy that convergence is not possible in practice , as we will illustrate later on . 3 Assessing the Reliability of Black-box Variational Inference . 3.1 Conceptual framework . How can we determine – both conceptually and experimentally – what is required to obtain reliable estimates of the variational divergence and optimal variational approximation ? As we have seen , the most common variational divergences and their Monte Carlo gradient estimators can be expressed in terms of the density ratio w ( θ ) . Reliable black-box variational inference ultimately depends on the behavior of w ( θ ) since ( 1 ) accurate optimization requires low-variance and ( nearly ) unbiased gradient estimates Ĝ ( λ ) , and ( 2 ) determining convergence and validating the quality of variational approximations can require accurate estimates L̂ ( λ ) of variational divergences [ 14 , 15 ] . While asymptotically ( in the number of iterations and Monte Carlo sample size S ) there may be no issues with stochastic optimization or divergence estimation , in practice black-box variational inference operates in the pre-asymptotic regime . Therefore , the reliability of black-box variational inference depends on the pre-asymptotic behavior of the w ( θ ) , and how it interacts with the choice of variational objective and gradient estimator . Before accounting for the effects of the objective and gradient estimator , first consider the behavior of the density ratio w ( θ ) , which can also be interpreted as an importance sampling weight with qλ ( θ ) as the proposal distribution [ cf . 2 , 16 , 29 ] . Pickands [ 23 ] proved , under commonly satisfied conditions , that for u tending to infinity , the distribution of w ( θ ) | w ( θ ) > u is well-approximated by the three-parameter generalized Pareto distribution GPD ( u , σ , k ) , which for k > 0 has density p ( w | u , σ , k ) = σ−1 { 1 + k ( w − u ) /σ } −1−1/k where w is restricted to ( u , ∞ ) . Since w ( θ ) > 0 , this implies its distribution is heavily skewed to the right with a power-law tail . Consider the idealized scenario of estimating the mean of w ( θ ) ∼ GPD ( u , σ , k ) . We assume the mean is finite , which is equivalent to assuming k < 1 since b1/kc determines the number of finite moments . Because of the heavy right skew , most of the mass of w ( θ ) is below its mean . Therefore , even after averaging a large number of samples , most empirical estimates ∑S s=1 w ( θs ) will be smaller than the true mean . Figure 3a illustrates this behavior for different values of k : even with 1 million samples , the empirical mean is far below the true mean when k > 0.7 . The highly variable sizes of the confidence intervals based on 10,000 replications further highlight the instability of the estimator . So , even though the empirical mean is an unbiased estimator , in the pre-asymptotic regime ( before the generalized central limit theorem is applicable [ 5 ] ) , in practice the estimates are heavily biased downward with high probability . If w ( θ ) is not a generalized Pareto distribution , we can instead treat k as the tail index k : = inf { ` > 0 : Eθ∼q { w ( θ ) 1/ ` } < ∞ } , which encodes the same tail behavior as GPD ( u , σ , k ) . Crucially , we should expect k to be much larger than 0 when there is a significant mismatch between the target distribution and the variational family . Since selecting a variational family that can match the typical set tends to be more difficult in higher dimensions , we should expect k to be larger for higher-dimensional posteriors . We can generalize our observations about pre-asymptotic estimation bias to the estimators L̂ ( λ ) and Ĝ ( λ ) . For the loss estimator , we replace w ( θs ) with f ( w ( θs ) ) , where f ( w ) is polynomial in w and logw for the class of losses we consider . If the dominant term of f ( w ) is of order wα , the tail behavior will be similar to a generalized Pareto with kα = αk . Thus , L̂ ( λ ) will have larger pre-asymptotic bias as α increases . For example , estimation of the mass-covering inclusive KL ( where α = 1 ) – and , more generally , mass-covering α-divergences with α > 0 – will suffer from a large pre-asymptotic bias . On the other hand , for the mode-seeking exclusive KL , f ( w ) = log ( w ) , so we can expect all moments to be finite and therefore a much smaller pre-asymptotic bias . Similar considerations apply to the gradient estimator , with the details depending on the specific estimator used . However , when using self-normalized weights for α-divergences , we can expect a large pre-asymptotic bias whenever w ( θ ) has such bias since self-normalization involves estimating the mean ofw ( θ ) . This bias will affect the accuracy of the solution found using stochastic optimization . Thus , the quality of the solutions found can only partially be improved by using a smaller step size since smaller step sizes will only reduce the effects of a large estimator variance , but not the effects from a large bias . We provide more details on the behavior of the score function and reparameterized gradients for each of the divergences in ? ? ? ? , following Geffner and Domke [ 11 ] . In summary , our framework makes two key predictions : ( P1 ) Estimates and gradients of mode-seeking divergences ( in particular exclusive KL divergence with log dependence on w ) have lower variance and are less biased than those of mass-covering divergences ( in particular α-divergences with α > 0 , with polynomial dependence on w ) . ( P2 ) The degree of polynomial dependence on w determines how rapidly the bias and variance will increase as approximation accuracy degrades – in particular , in high dimensions . Because the adaptive f -divergence depends directly on the ( ordered ) weights , we expect it to behave similarly to the mass-covering divergences .

Summary: The paper discusses the accuracy of variational inference and the choices that need to be me made with regard to divergence and approximating family, particularly as they relate to the dimensionality of the problem. Based on an analysis of the pre-asymptotic behavior of the density ratios involved, the authors propose a framework based on Pareto smoothed importance sampling (PSIS) to help practitioners analyze the success of their approximate inference procedure. The authors provide experiments demonstrating general findings for variational inference for a range of divergences and approximate posteriors.