Challenges for BBVI with Normalizing Flows
Abstract
Current black-box variational inference (BBVI) methods require the user to make numerous design choices—such as the selection of variational objective and approximating family—yet there is little principled guidance on how to do so.
We develop a conceptual framework and set of experimental tools to understand the effects of these choices, which we leverage to propose best practices for maximizing posterior approximation accuracy.
Our approach is based on studying the pre-asymptotic tail behavior of the density ratios between the joint distribution and the variational approximation, then exploiting insights and tools from the importance sampling literature.
We focus on normalizing ﬂow models, though we are not limited to them.
1. Introduction
A great deal of progress has been made in black-box var iational 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 high-dimensional posteriors (Wang et al., 2018; Geffner &
Domke, 2020; Agrawal et al., 2020).
While the main focus in the machine learning literature has 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.
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
Copyright 2021 by author(s).
E xc l u s i v e K L
I n c l u s i v e K L 0 5 10 15 0 5 10 15 0 5 10 15 0 5 100.
Distance from the mode
Exclusive KL
Inclusive KL
Target 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 (MacKay, 2003; Vehtari et al., 2019).
Thus, as illustrated in Fig.
1, the beneﬁts of heavy-tailed approximate families and divergences favoring mass-covering diminish as dimensionality of the target distribution increases.
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 (Vehtari et al., 2019) as a simple and practical approach for estimating both the required minimal sample size and obtaining empirical and conceptual insights into the pre-asymptotic convergence rates of estimators of common divergences and their gradients.
Although normalizing ﬂow models offer high capacity to cover posterior mass in high dimensions, they can be notoriously hard to optimize and our proposed framework is general and can be helpful to make design choices.
2. Preliminaries and Background
Let p ( θ, Y ) be a joint distribution of a probabilistic model, where θ ∈ R D 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 ) :
= (cid:82) p ( θ, Y ) dθ ) is often 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 distributions Q .
The best approximation is determined by minimizing a divergence D ( p (cid:107) q ) , which measures the discrepancy between p and q :
[Equation]
where λ ∈ R K 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 .
Let w ( θ ) :
=
p ( θ, Y ) /q
( θ ) denote the density ratio between the joint and approximate distributions.
For a function φ : R D
→ R , the self-normalized importance sampling estimator for the posterior expectation E θ ∼ p
[ φ ( θ )] is given by ˆ I (
[Equation]
[Equation]
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.
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 (Vehtari et al., 2019).
Variational families.
Let q λ ( θ ) be an approximating family parameterised by a K -dimensional vector λ ∈ R K for
D -dimensional inputs θ ∈ R D .
We focus on the most popular mean-ﬁeld and normalizing ﬂow families.
Normalising ﬂows provide more ﬂexible families that can capture correlation and non-linear dependencies.
f -divergences.
The most commonly used divergences are f -divergences.
For a convex function f satisfying f (1) = 0 , [Figure] the f -divergence is given by
[Equation]
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.
Loss estimation and stochastic optimization.
In all cases, minimizing the f -divergence is equivalent to minimizing the loss function L f ( p (cid:107) q ) :
=
E θ ∼ q [ f ( w ( θ ))]
(although, see Wan et al.
(2020) for a different approach).
Let L ( λ ) :=
L f ( p (cid:107) q λ ) denote the loss as a function of the variational parameters λ .
The loss and its gradient G ( λ ) :
= ∇ λ L ( λ ) can both be approximated using the Monte Carlo estimates (cid:98) L ( λ )
=
1 S (cid:80) Ss =1 f ( w ( θ s )) and (cid:98) G ( λ ) = 1 S (cid:80) Ss =1 g ( θ s ) , (2)
where θ 1 , . .
. , θ S are independent draws from q λ and g : R K
→ R K is a gradient function that depends on f and w .
The gradient estimates can be used in a stochastic gradient optimization scheme such that
[Equation]
where η t is the step size.
In practice, more stable adaptive stochastic gradient optimisation methods such as RMSProp or Adam (Hinton & Tieleman, 2012) are often used.
3. Assessing the Reliability of BBVI 3.1. Conceptual framework
The most common variational divergences and their Monte
Carlo gradient estimators can be expressed in terms of the density ratio w ( θ ) .
Reliable BBVI ultimately depends on the behavior of w ( θ ) since (1) accurate optimization requires low-variance and (nearly) unbiased gradient estimates (cid:98) G ( λ ) , and (2) determining convergence and validating the quality of variational approximations can require [Figure] ̂
[Figure]
accurate estimates (cid:98) L ( λ ) of variational divergences (Kucukelbir et al., 2015; Huggins et al., 2020).
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 BBVI operates in the preasymptotic regime.
Therefore, the reliability of BBVI 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, ﬁrst 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.
Li & Turner, 2016; Wang et al., 2018; Bamler et al., 2017).
Pickands (1975) proved, under commonly satisﬁed conditions, that for u tending to inﬁnity, the distribution of w ( θ ) | w ( θ ) >
u is wellapproximated 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 ﬁnite, which is equivalent to assuming k < 1 since (cid:98) 1 /k (cid:99) determines the number of ﬁnite 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 (cid:80) Ss =1 w ( θ s ) will be smaller than the true mean.
Figure 5 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 .
The highly variable sizes of the conﬁdence 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 (Chen &
Shao, 2004)), 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 { (cid:96) > 0 : E θ ∼ q { w ( θ ) 1 /(cid:96) } < ∞} , 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 signiﬁcant mismatch between the target distribution and the variational family.
Since selecting a variational family that can match the typical set tends to be more difﬁcult 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 (cid:98) L ( λ ) and (cid:98) G ( λ ) .
For the loss estimator, we replace w ( θ s ) with f ( w ( θ s )) , where f ( w ) is polynomial in w and log w 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, (cid:98) L ( λ ) will have larger pre-asymptotic bias as α increases.
For example, estimation of the masscovering 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 ﬁnite and therefore a much smaller pre-asymptotic bias.
Similar considerations apply to the gradient estimator.
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 of w ( θ ) .
This bias will affect the accuracy of the solution found.
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.
3.2. Experimental framework
In the light of potentially large non-asymptotic bias arising from the heavy right tail of w ( θ ) , it is important to verify the pre-asymptotic behavior of the Monte Carlo estimators used in variational inference.
We follow the approach developed by Vehtari et al.
for importance sampling and compute an empirical estimate ˆ k of the tail index k by ﬁtting a generalized Pareto distribution to the observed tail draws.
Vehtari et al.
(2019) show that the minimal sample size to have a small error with high probability scales as
S =
O (exp { k/ (1 − k ) 2 } ) .
Two key propositions here are: (P1) Mode-seeking divergences are more stable and reliable than mass-covering ones.
Fig.
8 in Appendix shows that gradient bias and variance increases with dimension for inclusive KL and χ 2 but not exclusive KL.
(P2) Degree of polynomial dependence on w determines sensitivity to approximation–target mismatch.
Fig.also shows that gradient estimates (divergence estimates in Appendix A.1) of divergences that involve higher polynomials of w (e.g. χ 2 ) become more and more unstable as dimensions increase.
[Figure]
c
[Figure]
r
[Figure] [Figure]
4. Experiments
In this section, we describe a series of experiments to show how our framework works.
For all posteriors, we ﬁt mean-ﬁeld Gaussian and Student- t families, a planar ﬂow (Rezende & Mohamed, 2015) with 6 layers and a nonvolume preserving (NVP) ﬂow (Dinh et al., 2017) withstacked neural networks with 2 hidden layers of 10 neurons each for both the translation and scaling operations.
We use Stan (Stan Development Team, 2020) for model construction.
Mode-seeking divergences are easier to optimize.
In particular, the ˆ k values when using normalizing ﬂows, which are more challenging to optimize, is low for D < 20 when using exclusive KL, but inﬁnite when using either the inclusive KL or f -divergence.
We can see that exclusive KL provides also more accurate and reliable posterior approximations than the inclusive KL and adaptive f -divergence, particularly for the normalizing ﬂows.
This is despite the fact that we used 20 times as many MC samples to estimate the gradients for the inclusive KL and the f -divergence compared to the exclusive KL.
To further illustrate the relative difﬁculty of optimizing the inclusive KL divergence,
[Figure]
Fig.
4 shows the largest dimension for which the stochastic optimization converged as a function of the step-size.
For most step-sizes, the combination of normalizing ﬂows and the inclusive KL divergence only converged for D = 2 , whereas convergence is possible in higher dimensions for simpler variational families.
4.1. Realistic models and datasets
We compare variational approximations for models and datasets from posteriordb ( github.com/stan-dev/ posteriordb ) in terms of accuracy of the estimated moments and predictive likelihood.
We use T max = 15 ,and robust optimisation algorithm given in Dhaka et al.
(2020) for stochastic optimisation, for cases where the divergence objectives are analytical, we used BFGS optimisation.
Exclusive KL remains the most reliable for realistic posteriors.
The results are summarized in Figs.
5 and 6: exclusive KL is superior for higher-dimensional posteriors (e.g.,
D > 10 ) or when combined with normalizing ﬂows, while inclusive KL is better for the large values for ˆ k indicate that ﬁtting approximations based on normalizing ﬂows remains a challenge in high dimensions.
The performance for the adaptive f -divergence is comparable to the inclusive KL divergence.
Normalizing ﬂows can be effective but are challenging to optimize.
Fig.
6 also shows that normalizing ﬂows can be quite effective when used with exclusive
KL to ensure stable optimization However, as can be seen in Fig.
5 and Fig.
3, when using out-of-the-box optimization with no problem-speciﬁc tuning (as we have done for a fair comparison), the normalizing ﬂows approximations can have pathological features – even in low dimensions.
A. Appendix A.1. Divergence estimates
We add an additional visualization that demonstrates the unstable behaviour of the f -divergence estimates as dimensions increase, particularly for divergences that involve higher polynomials of w .
B. PosteriorDB datasets
In Table 1 we show the dimensionality of the datasets we use for our real experiments.
C. Additional results for the pre-asymptotic reliability case study
In Fig.
9 and Fig.
10 we show additional results for the preasymptotic reliability case study for different objectives and mean ﬁeld Gaussian approximation.
The results from opti-
[Table]
mising χ 2 , 1 / 2 -divergence and tail adaptive f -divergence follow similar trends as those resulting from optimising exclusive and inclusive KL.
Approximations obtained by optimising χ 2 and 1 / 2 -divergence are more unstable and end up diverging in similar ways as inclusive KL even for moderately low dimensional problems.
We use a warm start procedure for χ 2 , 1 / 2 -divergence and inclusive KL, starting at the solution of exclusive KL for a given problem.
On the other hand, optimising tail adaptive f -divergence seems to be more robust and behave similarly to exclusive KL even in higher dimensions.
[Figure]
Exclusive
[Figure]
Divergence
Inclusive KL
[Figure]
−1.
Dimension
P a r e t o k ^ log(w) w^(1/2) w wlog(w) w^
[Figure]
Inclusive KL
[Figure]
Dimension l og 10 ( D f )
[Figure]