Interpolating Between Sampling and Variational Inference with Infinite Stochastic Mixtures
Abstract
Sampling and Variational Inference (VI) are two large families of methods for approximate inference that have complementary strengths.
Sampling methods excel at approximating arbitrary probability distributions, but can be inefficient.
VI methods are efficient, but may misrepresent the true distribution.
Here, we develop a general framework where approximations are stochastic mixtures of simple component distributions.
Both sampling and VI can be seen as special cases: in sampling, each mixture component is a delta-function and is chosen stochastically, while in standard VI a single component is chosen to minimize divergence.
We derive a practical method that interpolates between sampling and VI by analytically solving an optimization problem over a mixing distribution.
Intermediate inference methods then arise by varying a single parameter.
Our method provably improves on sampling (reducing variance) and on VI (reducing bias+variance despite increasing variance).
We demonstrate our method’s bias/variance trade-off in practice on reference problems, and we compare outcomes to commonly used sampling and VI methods.
This work takes a step towards a highly flexible yet simple family of inference methods that combines the complementary strengths of sampling and VI.
INTRODUCTION
We are concerned with the familiar and general case of approximating a probability distribution, such as occurs in
Bayesian inference when both the prior over latent variables and the likelihood function connecting them to data are known, but computing the posterior exactly is intractable.
There are two largely separate families of techniques for approximating such intractable inference problems: Markov
Chain Monte Carlo (MCMC) sampling, and Variational
Inference (VI)
[Bishop, 2006, Murphy, 2012]
Sampling-based methods, including MCMC, approximate a distribution with a finite set of representative points.
MCMC methods are stochastic and sequential, generating a sequence of sample points that, given enough time, become representative of the underlying distribution increasingly well.
MCMC sampling is (typically) asymptotically unbiased, at the expense of high variance, leading to long run times in practice.
Similar to the approach we take here, sampling methods are studied at different scales: both in terms of their asymptotic limit (i.e. their bias at infinitely many samples) and their practical behavior for finite samples or other resource limits [Korattikara et al., 2014, Angelino et al., 2016].
Variational Inference (VI) refers to methods that produce an approximate distribution by minimizing some quantification of divergence between the approximation and the desired posterior distribution [Blei et al., 2017, Zhang et al., 2019].
For the purposes of this paper, we will use VI to refer to the most common flavor of variational methods, namely minimizing the Kullback-Leibler ( KL ) divergence between an approximate distribution from a fixed family and the desired distribution [Bishop, 2006, Wainwright and Jordan, 2008,
Murphy, 2012, Blei et al., 2017].
The best-fitting approximate distribution is often used directly as a proxy for the true posterior in subsequent calculations, which can greatly simplify those downstream calculations if the approximate distribution is itself easy to integrate.
In contrast to MCMC,
VI is often used in cases where speed is more important than asymptotic bias [Angelino et al., 2016, Blei et al., 2017,
Zhang et al., 2019].
Our goal is to develop an intermediate family of methods that “interpolate” between MCMC and VI, inspired by a simple and intuitive picture (Figure 1): we propose applying sampling methods in the space of variational parameters such that the resulting approximation is a stochastic mixture of variational “component” distributions [Yin and
This extends sampling by replacing the sampled points with extended components, and it extends VI by replacing the single best-fitting variational distribution with a stochastic mixture of more localized components.
This is qualitatively distinct from previous variational methods that use stochastic optimization : rather than stochastically optimizing a single variational approximation [Hoffman et al., 2013, Salimans et al., 2015], we use stochasticity to construct a random mixture of variational components that achieves lower asymptotic bias than any one component could.
As we will show below, this framework generalizes both sampling and VI, where sampling and VI emerge as special cases of a single optimization problem.
This paper is organized as follows.
In section 2, we set up the problem and our notation, and describe how both classic sampling and classic VI can be understood as special cases of stochastic mixtures.
In section 3, we introduce an intuitive framework for reasoning about infinite stochastic mixtures and define an optimization problem that captures the trade-off between sampling and VI.
Section 4 introduces an approximate objective and closed-form solution and describes a simple practical algorithm.
Section 5 gives empirical and theoretical results that show how our method interpolates the bias and variance of sampling and VI.
Finally, section 6 concludes with a summary, related work, limitations, and future directions.
SETUP AND NOTATION
Let p ∗ ( x ) =
Zp( x ) denote the unnormalized probability distribution of interest, with unknown normalizing constant
Z .
For instance, in the common case of a probabilistic model with latent variables x , observed data D , and joint distribution p( x , D ) , we are interested in approximations to the posterior distribution p( x |D ) .
This is intractable in general, but we assume that we have access to the un-normalized posterior p ∗ ( x |D ) =
1Z p
D| x )p( x ) .
1 Let q( x ; θ ) be any “simple” distribution that may be used used in a classic VI context (such as mean-field or Gaussian), and let m T ( x ) be a mixture containing T of these simple distributions as components, defined by a set of T parameters { θ (1) , . . .
, θ ( T
[Equation]
[Equation] For example, if q is a multivariate normal with mean µ and covariance Σ , then θ ( t ) = { µ ( t ) , Σ ( t ) } and m T ( x ) would be a mixture of T component normal distributions [Gershman et al., 2012, Zobay, 2014].
We will study properties of mixing distributions, or distributions over component parameters, which we denote ψ ( θ )
[Ranganath et al., 2016].
If the set of θ ( t ) is drawn randomly from ψ ( θ ) , then as T → ∞ , m T ( x ) approaches the idealized infinite mixture,
[Equation]
Sampling and VI as special cases of the mixing distribution.
Let θ ∗ = arg min θ KL(q( x ; θ ) || p( x )) be the parameters corresponding to the classic single-component variational solution.
VI corresponds to the special case where the mixing distribution ψ ( θ ) is a Dirac delta around θ ∗ , or ψ ( θ ) =
δ ( θ − θ ∗ ) , in which case the mixture m T ( x ) is equivalent to q( x ; θ ∗ ) regardless of the number of components T .
Sampling can also be seen as a special case of ψ ( θ ) in which each component narrows to a Dirac delta ( ψ ( θ ) places negligible mass on regions of θ -space where components have appreciable width), and the means of the components are distributed according to p( x ) .
This requires that the component family q( x ; θ ) is capable of expressing a Dirac-delta at any point x , such as a location-scale family.
Thus, both sampling and VI can be seen as limiting cases of stochastic mixture distributions, m T ( x ) , defined by a distribution over component parameters, ψ ( θ ) .
In what follows, we will show how designing the mixing distribution ψ ( θ ) allows us to create mixtures that trade-off the complementary strengths of sampling and VI.
CONCEPTUAL FRAMEWORK 3.1 DECOMPOSING KL(m || p) INTO MUTUAL INFORMATION AND EXPECTED KL
The idealized infinite mixture m( x ) is fully defined by the chosen component family q( x ; θ ) and the mixing distribution ψ ( θ ) .
Consider the variational objective with respect to the entire mixture, KL(m || p) :
[Equation]
where Z is the normalizing constant of p ∗ ( x ) and is irrelevant for constructing m( x ) .
Instead of (3), one can use the equivalent objective of maximizing the Evidence Lower
BOund or ELBO [Bishop, 2006, Murphy, 2012, Blei et al., 2017].
Regardless, minimizing (3) or maximizing the ELBO for mixtures is intractable in general.
However, as first shown by Jaakkola and Jordan [1998] for finite mixtures, it admits the following useful decomposition:
[Equation]
(dropping log Z ).
The first term, (i), is the Expected KL
Divergence for each component when the parameters are drawn from ψ ( θ ) .
This term quantifies, on average, how well the mixture components match the target distribution.
In isolation, Expected KL is minimized when all components individually minimize KL(q || p) , i.e. when ψ ( θ ) → δ ( θ − θ ∗ ) .
This tendency to concentrate ψ ( θ ) to the single best variational solution is balanced by the second term, (ii), which is the Mutual Information between x and θ , which we will write I [ x ; θ ] , under the joint distribution q( x ; θ ) ψ ( θ ) .
This term should be maximized , and, importantly, it does not depend on p ∗ ( x ) .
Mutual Information is maximized when the components are as diverse as possible, which encourages the components to become narrow and to spread out over diverse regions of x regardless of how well they agree with p( x ) .
This decomposition of KL(m || p) into Mutual
Information (between x and θ ) and Expected KL (between q and p ) is convenient because approximations to Mutual
Information are well-studied, and minimizing Expected KL can leverage standard tools from VI.
3.2 TRADING OFF BETWEEN MUTUTAL INFORMATION AND EXPECTED KL
We will refer back to this decomposition of the KL(m || p) objective into Expected KL and Mutual Information throughout.
Figure 2 depicts a two-dimensional space with Expected KL on the x-axis and Mutual Information on the y-axis.
Any given mixing distribution ψ ( θ ) can be placed as a point in this space, but in general many ψ ( θ ) ’s may map to the same point.
Sampling and VI live at extreme points in this space.
Classic
VI, where ψ ( θ ) =
δ ( θ − θ ∗ ) , corresponds to the blue point (c), because by definition θ ∗ achieves the minimum possible
KL , and I [ x ; θ ] is zero.
Classic sampling corresponds to the green point (d), with ψ ( θ ) placing mass only on Diracdelta-like components, and selecting each component with probability p( µ ) , where µ is the mean of q determined by θ .
Towards the goal of constructing mixtures that trade-off properties of sampling and VI, we propose to view the two terms in (4) as separate objectives that may be differently weighted, and maximizing the objective
[Equation]
for a given hyperparameter λ with respect to the mixing distribution ψ ( θ ) .
This objective may alternatively be viewed as the Lagrangian of a constrained optimization problem over the mixing density ψ ( θ ) , where Mutual Information is maximized subject to a constraint on Expected KL.
This is a concave maximization problem with linear constraints, defining a Pareto front of solutions that each achieve a different balance between Expected KL and Mutual Information.
Maximizing Mutual Information necessitates approximations [Poole et al., 2019], so there may be good mixture approximations that are not found in practice, such as the yellow point (e) in Figure 2.
In section 4 below, we use an approximation to Mutual Information that has the property, illustrated by the orange curve (f) in Figure 2, of connecting VI (c) to sampling (d), controlled by varying λ .
As shown on the right of Figure 2, our method produces mixtures that behave like classic samples when λ = 1 , that behave like classic VI when λ → ∞ , and that exhibit intermediate behavior at intermediate values of λ .
We emphasize that this frame is quite general: any stochastic mixture can be reasoned about in terms of its Expected
KL and Mutual Information, and this is a natural space in which to think about interpolating sampling and VI.
A similar decomposition of KL(m || p) (or the ELBO) has been used by previous methods that optimize mixtures [Zobay, 2014, Jaakkola and Jordan, 1998, Gershman et al., 2012,
Yin and Zhou, 2018].
The primary difference between these previous methods is how they approximate (or lower-bound)
Mutual Information.
In the next section, we introduce a new approximation that is particularly efficient, and is the first to our knowledge that replicates sampling-like behavior with finitely many components.
APPROXIMATE OBJECTIVE
Maximizing Mutual Information, as is required by (5), is a notoriously difficult problem that arises in many domains, and there is a large collection of approximations and bounds in the literature [Jaakkola and Jordan, 1998, Brunel and
Nadal, 1998, Gershman et al., 2012, Wei and Stocker, 2016,
Kolchinsky and Tracey, 2017, Poole et al., 2019].
Previous work has optimized finite mixtures by considering how each of T components interacts with the other T − 1 components, resulting in quadratic scaling with T [Gershman et al., 2012,
Zobay, 2014, Guo et al., 2016, Miller et al., 2017, Kolchinsky and Tracey, 2017, Nalisnick and Smyth, 2017, Yin and Zhou, 2018, Poole et al., 2019].
Beginning instead with infinite mixtures, we find that the local geometry of θ -space is sufficient to provide an approximation to Mutual Information that can be evaluated independently for each value of θ , leading to linear scaling with T .
4.1 STAM’S INEQUALITY
Mutual Information between x and θ can be written as
[Equation]
where H [ θ ] is the entropy of ψ ( θ ) and H [ˆ θ | x ] is the entropy of q(ˆ θ | x ) =
q( x ;ˆ θ ) ψ (ˆ θ ) m( x ) , i.e. the distribution of inferred θ values for a given x .
The second line follows simply from expanding the definition of m( x ) in the outer expectation.
The term H [ˆ θ | θ ] can be thought of in terms of a statistical estimation problem: ˆ θ is the “recovered” value of θ after passing through the “channel” x .
Bounding the error of such estimators is a well-studied problem in statistics.
From (6), a lower-bound on Mutual Information can be derived from an upper bound on H [ˆ θ | θ ] for each θ .
For this, we draw inspiration from Stam’s inequality [Stam, 1959,
Dembo et al., 1991, Wei and Stocker, 2016], which states
[Equation]
where | · | is a determinant, and F ( θ ) is the Fisher Information Matrix, defined as
[Equation]
The Fisher Information Matrix is also the local metric on the statistical manifold with coordinates θ [Amari, 2016]; it is used to quantify how “distinguishable” θ is from θ + dθ .
Note that (7) can be viewed as the entropy of a Gaussian approximation to q(ˆ θ | x ) with precision matrix F ( θ ) ; this approximation is most accurate when q( x ; θ ) itself is narrow and approximately Gaussian [Wei and Stocker, 2016].
Combining (6) and (7), we propose to use
[Equation]
as a proxy for the intractable I [ x ; θ ] in (5), having used log |F − 1 | = − log |F| .
Note that I F [ x ; θ ] has not been proven to be a strict bound on I [ x ; θ ] , but may be seen as an approximation to it
[Wei and Stocker, 2016].
Briefly, this is because the original Stam’s inequality, as stated in (7), assumes θ is a scalar location parameter, and assumes the high-precision limit where q( ˆ θ | x ) is well-approximated by a Gaussian.
Despite this, I F [ x ; θ ] is well-suited for our purposes, since (i) it leads to a remarkably simple and easy to implement expression for ψ ( θ ) below; (ii) we can prove that it leads to sampling when λ = 1 and VI when λ → ∞ ; and (iii) we suspect that the inequality in (7) is nonetheless strict, since we neglect the prior information contained in ψ ( θ ) and therefore overestimate the conditional entropy H [ˆ θ | θ ] .
By analogy to the
Bayesian Cramér-Rao bound [Gill and Levit, 1995, Fauß et al., 2021], a tighter variant of (7) could be derived that takes into account the prior, though possibly at the expense of added complexity; we leave this to future work.
4.2 CLOSED-FORM MIXING DISTRIBUTION
Substituting I F [ x ; θ ] for I [ x ; θ ] in (5) gives the following approximate objective,
[Equation]
having dropped additive constants.
This now resembles a maximum-entropy problem with an expected-value constraint, which has the following simple closed-form solution:
[Equation]
again dropping additive constants.
Equation (10) is strikingly simple, and amenable to many existing MCMC sampling methods for drawing samples of θ from ψ ( θ ) .
Despite being derived from an approximation to our original objective, (10) nonetheless contains both sampling and VI as special cases.
As λ → ∞ , the KL term dominates and ψ ( θ ) concentrates to δ ( θ − θ ∗ ) , reproducing VI.
When λ = 1 , the resulting mixture recovers the behavior of “sampling” in the following sense:
Definition 1 (Sampling)
A stochastic mixture, defined by the component family q( x ; θ ) and mixing distribution ψ ( θ ) , is considered to be “sampling” if it is unbiased and it consists of non-overlapping components . An unbiased mixture is one where m( x ) =
p( x ) . A mixture consists of T nonoverlapping components if (cid:80) Tt =
1 q( x ; θ t ) ≈ max t q( x ; θ t ) with high probability.
Lemma 4 in Appendix A.2 establishes that ψ ( θ ) with λ =leads to sampling as defined here, assuming mixture components q are Gaussian.
However, we conjecture that sampling arises from a broader class of q components as well, though computing and differentiating through F ( θ ) for nonGaussian component families poses additional challenges.
4.3 IMPLEMENTATION
Equation (10) provides a closed-form unnormalized log probability density, which is straightforward to sample from using any of a large number of existing sampling methods. For example, discrete Langevin dynamics are [Equation]
{+) (cid:124) (cid:123)(cid:122) (cid:125) ( i ) + γ 2 ∇ θ log |F| (cid:124) (cid:123)(cid:122) (cid:125) ( ii ) + (cid:112) 
+}
γη t (cid:124) (cid:123)(cid:122) (cid:125) ( iii ) where γ is the step size and η t is unit isotropic Gaussian noise. This update rule is remarkably simple: ( i ) is equivalent to gradient descent of KL(q || p) , as done in ADVI [Kucukelbir et al., 2017], ( ii ) biases the updates towards regions where |F ( θ ) | is large (i.e. narrower components), and ( iii ) adds noise. For our experiments below, we implemented sampling from (10) in Stan [Carpenter et al., 2017], an open-source framework for probabilistic models and approximate inference algorithms. We set q to be a multivariate Gaussian with diagonal (axis-aligned) covariance, and sampled θ from ψ ( θ ) using Stan’s default implementation of the No U-Turn Sampler (NUTS)
[Hoffman and Gelman, 2014], but we emphasize that samples can be drawn from (10) using a variety of off-the-shelf sampling methods.
We computed the KL(q || p) term using 200 random samples from q per evaluation, using the reparameterization trick to compute the gradient ∇ θ KL(q || p) and resampling the reparameterization noise only once per NUTS trajectory.
This incurred a high cost in terms of number of function evaluations per sample of θ , but this cost can in principle be significantly reduced by using a sampler that accepts stochastic function evaluations [Korattikara et al., 2014, Ma et al., 2015].
All comparisons to existing methods were with Stan’s built-in NUTS sampler (over x ) and its built-in Automatic-Differentiation VI (ADVI)
[Kucukelbir et al., 2017].
NAVIGATING BIAS/VARIANCE TRADE-OFFS FOR FINITE T 5.1 REDUCING MEAN SQUARED ERROR (MSE)
In this section, we expound the sense in which our method “interpolates” sampling and VI in terms of bias and variance.
In our experiments, we quantify bias and variance in terms of the Mean Squared Error (MSE) of the expectation of an arbitrary f ( x ) using a random mixture of T components, m T ( x ) .
For ADVI, we measured variance across runs with different random initializations.
In Figure 3, we show empirically that by increasing λ one can interpolate between the low bias but high variance solution, equivalent to sampling, and the low variance but high bias solution, equivalent to VI.
Between these extremes, our method smoothly interpolates both bias and variance.
Further implementation details can be found in Appendix B.
Computing bias and variance requires choosing a class of functions f ( x ) .
We construct random smooth functions by discrete Fourier synthesis.
Specifically, we select a series of sinusoid plane waves in the space of x with increasing frequency ω , random directions t and phase ϕ , such that f ( x ) =
(cid:80)
Nω =1 a ω sin( ω t ⊤ x + ϕ ω ) .
The amplitudes a ω are set according to a power law: a ω = ω − α .
Note that by choosing a random direction t for each frequency, it is easy to apply this definition of a random function to arbitrarily high-dimensional inference problems.
An example of a 2dimensional f ( x ) is shown in Figure 3b with α = − 1 . 5 .
We vary the smoothness of the integrated function in Figure 4 by varying α [Stein and Shakarchi, 2011].
We also tested our algorithm on three reference problems from posteriordb [Magnusson et al., 2021], as well as on a 32-dimensional regression problem with synthetic data and known ground-truth parameters, shown in Supplemental Figure B.1.
The conclusion is similar: across many random f s, our algorithm performs on average as well as or better than both sampling (by reducing variance) and VI (by reducing bias).
5.2 CONSIDERATIONS FOR SELECTING λ
A first practical consideration for the choice of λ is the particular function f ( x ) to be integrated.
Since MSE can be decomposed into the sum of squared bias and variance, the value of λ that minimizes MSE occurs when ∂ bias
[Equation]
Any factor that increases the variance but not the bias for a fixed number of components T will push the optimal λ towards higher values (closer to VI).
One such factor is the smoothness of f ( x ) .
Classic sampling can have problematically high variance when f ( x ) is very jagged, as single points are not very representative of the surrounding function.
Intuitively, then, higher λ (more VI-like mixtures) is preferred when f ( x ) is more “wiggly.”
To show this, we generated a random function with varying smoothness and computed their expectations them over random mixtures. The resulting MSE, bias, and variance are shown in Figure 4.
We adjusted smoothness by varying the power law decay, α , for a fixed set of phases and wave directions.
At any value of λ , variance can be seen to increase as f is made more wiggly ( α → − 1 ).
With all else held equal, it is better to trade some variance for bias when the integrand changes quickly with x .
In Figures 3 and 4, we evaluated bias and variance on the “banana” distribution.
Similar results on higher-dimensional problems can be found in Supplemental Figure B.1. For illustration purposes, we chose T so that the variance of
NUTS was the same order of magnitude as the bias of ADVI, and estimated bias and variance by randomly subsampling sets of size T from much longer chains.
This approach provides theoretical insights on bias/variance trade-offs for T independent samples, but in practice bias will be higher due to burn-in time, variance will be higher due to sampler autocorrelations, and each of these may depend nontrivially on λ . Because MCMC samplers are most effective when they have been tuned to the problem at hand, a challenge for future work will be to adapt the sampler parameters on the fly as λ changes. Another factor that affects the optimal λ is the computational budget.
In our experiments we set a fixed T to demonstrate our algorithm’s properties.
However, if the time budget is not known in advance, a practitioner may wish to decrease λ adaptively over time.
Since variance is O ( T − 1 ) , for sufficiently large T error is dominated by bias, and so the optimal λ will decay towards 1 .
This would result in VI-like behavior for small T and sampling-like behavior for large T .
How quickly λ should decay will depend on the particular problem, specifically, on ∂ Bias 2 ∂λ , and on how efficiently one can produce T independent samples of θ .
5.3 ANALYTICAL RESULTS
While the MSE of the expected value of some f ( x ) is a useful way to compare approximate inference methods, it depends on the somewhat arbitrary choice of f , and in practice, the f ’s of interest are often not known at the time of inference.
This motivates using the following alternative definition of error that is independent of f and closely related to the variational objective of minimizing KL divergence:
[Equation]
That is, we define KL bias as the KL divergence from the infinite mixture m( x ) to the true distribution, and KL variance as the average KL , over realizations of T independent mixture components, from m T ( x ) to the infinite mixture m( x ) .
Note that KL bias is identical to the infinite-mixture objective we started with in (4).
The following theorem establishes that, for all finite T , we can reduce the KL error relative to sampling using some λ > 1 .
Theorem 1 (Improve on sampling)
If a mixture is sampling as in Definition 1, then dd λ KL bias = 0 and dd λ KL variance < 0 .
Thus, dd λ KL error < 0 .
Proof: see Appendix A.2. This theorem establishes the intuitive result that the variance of sampling can be reduced, minimally impacting its bias, by replacing samples with narrow mixture components.
Importantly, Theorem 1 is based on how ψ ( θ ) changes with λ when using the closed-form expression for ψ ( θ ) we derived based on the approximate L F objective.
For this theorem to apply, we must further show that both conditions of “sampling” (Definition 1) are met by ψ ( θ ) when λ = 1 .
This is proved in Lemma 4 in Appendix A.2 for Gaussian components, though we suspect it holds for other component families as well.
We can also improve on VI using our method.
However, this result is slightly more subtle, as there are three cases where one should expect VI to be (locally) optimal.
First, if q is in the same family as p , then q( x ; θ ∗ ) = p( x ) , then is no benefit to increasing T , and reducing λ only adds variance.
Similarly, if q is not in the same family but the VI solution is sufficiently close to p , then a mixture of nearby q s will add variance to m [Lindsay, 1983], potentially making the match to p worse. Third, if T is small – in the most extreme case, if T = 1 – then reducing λ will again only add variance without reducing bias.
With these three cases in mind, the following theorem establishes conditions where we expect to reduce KL error relative to VI by using a large but finite λ < ∞ .
Theorem 2 (Improve on VI) Assume that q( x ; θ ∗ ) is poorly matched to p( x ) , in the sense that
Tr (cid:0) ( ∇ 2 θ KL(q || p)) − 1 F (cid:1) > | θ | , and that λ is sufficiently large to use a Laplace approximation to ψ ( θ ) around θ ∗ .
Then, there exists some finite T 0 > 1 such that for all T ≥ T 0 , dd λ KL error > 0 .
Proof: see Appendix A.3.
Here, | θ | is the dimensionality of θ . The requirement that “ q( x ; θ ∗ ) is poorly matched to p( x ) ” is expressed in terms of the curvature of KL(q || p) around θ ∗ ; if this curvature is small, then many “nearby” q s will also fit p well, and a mixture of them can improve on VI despite adding variance.
On the other hand, the case where this curvature is not small corresponds to the earlier intuition that VI cannot be improved upon if the VI solution is already close to p .
For further details, see the full proof in Appendix A.3.
DISCUSSION
Summary: Our work provides a new perspective on the relationship between the two dominant frameworks for approximate inference – sampling and VI – by viewing both as special cases of inference using a broader class of stochastic mixtures.
Our main theoretical contribution is the framework shown in Figure 2, where mixtures that “interpolate” sampling and VI are analyzed in terms of how they trade off
Mutual Information and Expected KL.
We then derived an easy-to-use method based on an approximation to Mutual
Information that uses the local geometry of the space of variational parameters.
To demonstrate the ease and effectiveness of our method, we implemented it in the popular
Stan language and demonstrated using a small set of reference problems how we “interpolate” sampling and VI by varying a single parameter, λ .
Finally, we showed why such an intermediate inference scheme is useful in practice in terms of trading off bias and variance.
On one hand, we proved that it is always possible to improve on classic sampling ( λ = 1 ) by increasing λ : our method provably reduces the variance of sampling while minimally impacting its bias.
Our method also provably reduces the bias of VI under certain intuitive conditions.
Time and space complexity:
By approximating Mutual
Information using only local geometric information in (8), in our method each component can be selected independently of the others.
This means we can select and evaluate T components in O ( T ) time and
either O ( T ) space
(if all are stored) or O (1) space (if components are evaluated online) – identical to traditional MCMC sampling algorithms.
Further, we can run independent chains sampling θ ∼ ψ ( θ ) for a constant factor speedup.
This improves on past work using mixture approximations, which incurred O ( T 2 ) time and
O ( T ) space complexity, since the optimization problem for the T th component depends on the location of the other
T − 1 components, all of which must be in memory at once [Jaakkola and Jordan, 1998, Gershman et al., 2012, Zobay, 2014, Guo et al., 2016, Nalisnick and Smyth, 2017, Miller et al., 2017, Acerbi, 2018, Yin and Zhou, 2018]
(but the O ( T 2 ) complexity may be hardware-accelerated).
Related Work: The trade-offs between sampling and VI are well-studied, and many methods have been proposed to “close the gap” between them (see [Angelino et al., 2016,
Zhang et al., 2019] for general reviews).
Like these other methods, we aim to provide good approximations with high computational efficiency and low variance.
There are many methods that use mixture models to reduce the bias of variational inference.
Theorem 2 shows that our method only “beats” classic VI
when T > T 0 for some finite but potentially large T 0 .
This is the price we pay for drawing mixture components stochastically [Salimans et al., 2015, Yin and Zhou, 2018].
When a mixture of T components is optimized rather than sampled , bias is reduced and variance remains near zero, as in previous work [Jaakkola and Jordan, 1998, Gershman et al., 2012, Zobay, 2014, Guo et al., 2016, Miller et al., 2017], but in previous work this optimization has incurred a O ( T 2 ) cost while our method is O ( T ) and can be further parallelized.
Further, with some notable exceptions [Anaya-Izquierdo and Marriott, 2007, Salimans et al., 2015], most mixture VI methods make strong assumptions about the family of components [Jaakkola and Jordan, 1998, Gershman et al., 2012, Acerbi, 2018, Miller et al., 2017].
Our framework and method is somewhat agnostic to the family of q , though we have only rigorously proved that is asymptotically unbiased when using Gaussian components.
Many methods use sampling in the service of variational inference, or vice versa, but do not provide a unifying approach to both.
These typically use the samples to compute expectations used to update a variational approximation [Acerbi, 2018, Miller et al., 2017, Kucukelbir et al., 2017], rather than to generate the mixture components themselves.
There is also a large number of sampling approaches that aim to improve the efficiency of sampling by reducing its variance at the cost of some bias.
Some of these use variational approaches as proposal distributions, but ultimately the posterior is approximated by a set of (possibly weighted) samples of the latent variables [de Freitas et al., 2001, Korattikara et al., 2014, Ma et al., 2015, Zhang et al., 2021].
By expanding each sample from a point to a distribution, our approach allows each sample to cover more space with less variance and greater efficiency.
Despite some high-level similarities to other approaches, our framework is unusual in approximating the posterior by a sampled mixture of variational approximations.
The
Mixture Kalman filter [Chen and Liu, 2000] is a special case of this, which uses a sampled mixture of Gaussians, each constructed as a Kalman filter.
A related approach is to optimize a parameterized function that generates mixture components [Salimans et al., 2015, Wolf et al., 2016, Yin and Zhou, 2018], and generative diffusion models can also be seen as a case of such mixtures [Sohl-Dickstein et al., 2015, Ho et al., 2020].
Our work differs in that we derived a closed-form mixing distribution that requires no additional learning or optimization and that is readily implemented in existing inference software (Stan, [Carpenter et al., 2017]). The trade-offs described in this section are summarized in Table C.1 in the Appendix.
Limitations and future work: Using I F [ x ; θ ] to approximate I [ x ; θ ] reduces the generality of our method, since the former is most appropriate for narrow and Gaussianlike components
[Wei and Stocker, 2016].
Incorporating prior information from ψ ( θ ) into this bound, generalizing to other kinds of components, or even starting with alternative bounds on I [ x ; θ ] are all interesting avenues for future work.
Our proof of Theorem 2 requires an assumption about the curvature of KL(q || p) near the VI solution; identifying cases where this assumption holds may also be an interesting future direction.
We currently only study mixtures with T independent mixture components without taking into account the cost of producing independent samples of θ .
In reality, this cost depends on the quality of the sampler, warm-up and burn-in time, and a potentially large number of calls to log p( x ) [Zhang et al., 2021].
Further, λ dramatically changes the shape of log ψ ( θ ) , which may affect the efficiency of the sampler – we mitigated this slightly by scaling the mass parameter of NUTS with λ .
Other sampling algorithms besides
NUTS may also be beneficial, such as ULA, which is known to have favorable scaling properties in higher dimensions [Durmus and Moulines, 2017].
We have so far considered λ to be constant for a run of our algorithm, and this can lead to asymptotic bias even when
T is large.
A simple adjustment to make our method effective at both small and large T would be to decay λ as T grows, but note that this may require adapting the sampler parameters on the fly.
Our method also requires evaluating
KL(q || p) many times per sample of θ .
This could be made more efficient by adapting the number of Monte Carlo evaluations (fewer samples from q are sufficient when λ is low and components are narrow), by accounting for stochastic likelihood evaluations [Ma et al., 2015], or by extending our method to mean-field message-passing [Jaakkola and Jordan, 1998], where ∇ θ KL(q || p) can be computed in closed form [Hoffman et al., 2013].
Acknowledgements
This material is based upon work supported by the Air Force
Office of Scientific Research under award number FA955021-1-0422, NSF CAREER IOS-1552868, and support from the McNair Foundation to XP, as well as support from the
[-Lemma 
-]
[-Lemma 
-]
[-Wherever this holds for some  , it will additionally hold for all narrower components at the same µ . 
-]
[-Despite its impropriety, we are free to draw samples of ✓ from this improper when  = 
-]
NIH R01 EY028811 to RMH.
[-TT  
-]
Thanks also to Daniel Lee, who acted as our guide to the Stan codebase.