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 authors propose a conceptual framework and design guidelines for variational inference based on the pre-asymptotic tail behaviour of the unnormalized importance weight. This framework is used to study the behaviour of commonly variational families and f-divergences in mid-to-high dimensions. The authors observe that mass-covering divergences, while often superior theoretically and for low dimensional problems, are hard to optimize in mid-to-high dimensions compared to the mode-seeking exclusive KL-divergence. These observations are consistent with the prediction made using their framework based on estimating the tail-index of the best-fit generalized pareto distribution.