Nested Variational Inference
Year: 2021

Abstract:
We develop nested variational inference (NVI), a family of methods that learn proposals for nested importance samplers by minimizing an forward or reverse KL divergence at each level of nesting. NVI is applicable to many commonly-used importance sampling strategies and provides a mechanism for learning intermediate densities, which can serve as heuristics to guide the sampler. Our experiments apply NVI to (a) sample from a multimodal distribution using a learned annealing path (b) learn heuristics that approximate the likelihood of future observations in a hidden Markov model and (c) to perform amortized inference in hierarchical deep generative models. We observe that optimizing nested objectives leads to improved sample quality in terms of log average weight and effective sample size.

heading: 1 Introduction
text: Deep generative models provide a mechanism for incorporating priors into methods for unsupervised representation learning. This is particularly useful in settings where the prior defines an inductive bias that reflects the structure of the underlying domain. Training models with structured priors, however, poses some challenges. A standard strategy for training deep generative models is to maximize a reparameterized evidence lower bound with respect to both the generative and an inference model [Kingma, Welling, 2013; Rezende et al., 2014]. This approach works well in variational autoencoders with isotropic Gaussian priors, but often fails for models with more structured priors or likelihoods.
In recent years, a range of strategies for improving upon standard reparameterized variational inference have been put forward. These include wake-sleep style variational methods that minimize the forward KL-divergence [Bornschein, Bengio, 2015; Le et al., 2019] as well as sampling schemes that incorporate annealing [Huang et al., 2018], Sequential Monte Carlo [Le et al., 2018; Naesseth et al., 2017; Maddison et al., 2017], Gibbs sampling [Wu et al., 2019; Wang et al., 2018], and MCMC updates [Salimans et al., 2015; Hoffman, 2017; Li et al., 2017]. While these methods offer flexible inference, typically resulting in better approximations to the posterior compared to traditional variational inference methods, they are either model-specific, requiring specialized sampling schemes and gradient estimators, or can not be easily composed with other techniques.
In this paper, we propose nested variational inference, a framework for combining nested importance sampling and variational inference. Nested importance sampling formalizes the construction of proposals by way of calls to other importance samplers [Naesseth et al., 2015; Naesseth et al., 2019], Many existing importance samplers are instances of nested samplers, including methods based on annealed importance sampling [Neal, 2001] and sequential Monte Carlo [Del Moral et al., 2006].
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
NVI learns proposals by optimizing a divergence at each level of nesting. Additionally, we combine nested variational objectives with importance resampling to further improve sampling quality, without the need to undergo extra steps to maintain differentiability due to the local nature of the objective. Resampling also allows to compute gradient estimates based on incremental weights, which depend only on variables that are sampled locally, rather than on all variables in the model. Doing so yields lower variance weights, whilst maintaining a high sample diversity relative to existing methods.

heading: 2 Background

heading: 2.1 Stochastic Variational Inference.
text: Stochastic variational methods approximate a target density π(z; θ) = γ(z; θ)/Z with parameters θ using a variational density q(z;φ) with parameters φ. Two common variational objectives are the forward and reverse Kullback-Leibler (KL) divergence, which are both instances of f -divergences
Df ( π ∥∥ q) = E
q
[ f ( π(z; θ)
q(z;φ)
)] , (1)
with f(w) = w logw and f(w) = − logw respectively. We are typically interested in the setting where π(z; θ) is the posterior pθ(z |x) of a model with latent variables z and observations x. In this case, γ(z; θ) = pθ(x, z) is the joint density of the model, and Z = pθ(x) is the marginal likelihood.
Reverse KL-divergence. When optimizing KL (q ‖π), known as the reverse or exclusive KL divergence, it is common practice to maximize a lower boundL = Eq[log(γ/q)] = logZ−KL (q ‖π) with respect to θ and φ. The gradient of L can be approximated using reparameterized samples [Kingma, Welling, 2013; Rezende et al., 2014], likelihood-ratio estimators [Wingate, Weber, 2013; Ranganath et al., 2014], or a combination of the two [Schulman et al., 2015; Ritchie et al., 2016a].
Forward KL-divergence. In the case of the forward divergence KL (π ‖ q), also known as the inclusive KL divergence, stochastic variational methods typically optimize separate objectives for the inference and generative model. A common strategy is to train the generative model by either maximizing the ELBO or likelihood on the data. Optimizing the inference model requires samples from π, which itself requires approximate inference. A common strategy, which was popularized in the context of reweighted wake-sleep (RWS) methods [Bornschein, Bengio, 2015; Le et al., 2019], is to use q as a proposal in an importance sampler.

heading: 2.2 Importance Sampling
text: Self-Normalized Importance Samplers. An expectation Eπ[g(z)] with respect to π can be rewritten with respect to a proposal density q by introducing an unnormalized importance weight w,
E π
[ g(z) ] = 1
Z E q
[ w g(z) ] , w = γ(z; θ)
q(z;φ) . (2)
Self-normalized estimators use weighted samples (w(s), z(s))Ss=1 to both approximate the expectation with respect to q, and to compute an estimate Ẑ of the normalizing constant,
E π [g(z)] ' ĝ = 1 Ẑ 1 S S∑ s=1 w(s) g ( z(s) ) , Ẑ = 1 S S∑ s=1 w(s), z(s) ∼ q(·;φ). (3)
The resulting estimator is consistent, i.e. ĝ a.s.−→ Eπ[g(z)] as the number of samples S increases , but it is not unbiased, since it follows from Jensen’s inequality that 1/Z = 1/Eq[Ẑ] ≤ Eq[1/Ẑ] . If the discrepancy between q and π is large, the importance weights will have large variance, which also has an impact on the bias. When q = π, the importance weight w = Z has zero variance, and the Jensen’s inequality is tight. In the context of stochastic variational inference, this means that gradient estimates might initially be strongly biased, since there will typically be a large discrepancy between q and π. However, the variance will typically decrease as the quality of the approximation improves.
Nested importance sampling and proper weighting Nested Importance sampling formalizes the construction of proposals by way of calls to other importance samplers. Extending the example above, we now assume q is some unnormalized density itself that we cannot directly generates samples from. We can employ another importance sampler with proposal η to simulate samples from q,
z(l) ∼ η(·, ψ), w(l) = q(z (l);φ)
η(z(l);ψ) , Ẑq =
1
L L∑ l=1 w(l)
The resulting weighted samples (w(l), z(l))Ll=1 can subsequently be used compute the normalizing constant Zq of q and to generate samples that are approximately distributed w.r.t. q via resampling. This allows us to use q as a proposal in a subsequent importance sampling step, whose weighted samples can again be used to compute consistent estimates w.r.t. the target density of interest π. This can be generalized, by introducing the notion of proper weighting. Definition 2.1 (Proper weighting). Let π be a probability density. For some constant c > 0, a random pair (w, z) ∼ Π is properly weighted (p.w.) for an unnormalized probability density γ ≡ Zπ if w ≥ 0 and for all measurable functions g it holds that
E w,z∼Π [w g(z)] = c
∫ dz γ(z) g(z) = cZ E
z∼π [g(z)] .
As long as we can generate properly weighted samples for the target density π from some sampler Π, we can compute consistent estimates
1 S ∑S s=1 w
(s) g(z(s)) 1 S ∑S s=1 w (s) = cZ Ez∼π [g(z)] cZ = E z∼π [g(z)] . (4)

heading: 3 Nested Variational Inference
text: A widely used strategy in importance sampling is to decompose a difficult sampling problem into a series of easier problems. A common approach is to define a sequence of unnormalized densities {γk}Kk=1 that interpolate between an initial density π1 = γ1/Z1, for which sampling is easy, and the final target density πK = γK/ZK . At each step, samples from the preceding density serve to construct proposals for the next density, which is typically combined with importance resampling or application of a Markov chain Monte Carlo (MCMC) operator to improve the average sample quality.
Nested variational objectives. NVI defines objectives for optimizing importance samplers that target a sequence of densities. At every step, it minimizes the discrepancy between a forward density π̂k = γ̂k/Zk−1, which acts as the proposal, and a reverse density π̌k = γ̌k/Zk, which defines an intermediate target. We define the forward density by combining the preceding target γk−1 with a forward kernel qk, and the reverse density by combining the next target γk with a reverse kernel rk−1,
γ̌k(zk, zk−1) = γk(zk) rk−1(zk−1|zk), γ̂k(zk, zk−1) = qk(zk|zk−1) γk−1(zk−1). (5)
Our goal is to learn pairs of densities π̌k and π̂k that are as similar as possible. To do so, we minimize a variational objective D that comprises an f -divergence for each step in the sequence, along with a divergence between the first intermediate target π1 and an initial proposal q1. Concretely, the NVI objective can be formulated as,
DNVI = Df ( π1 ∣∣∣∣ q1)+ K∑
k=2
Df ( π̌k ∣∣∣∣ π̂k). (6)
Since each intermediate density πk occurs in both π̌k and in π̂k+1, this defines a trade-off between maximizing the similarity to πk+1 and the similarity to πk−1. To optimize this objective, we need to be able to simulate samples from the intermediate densities.
Sampling from intermediate densities. Given a pair (wk−1, zk−1) that is properly weighted for the previous target density γk−1, we can use a sequential Monte Carlo sampling [Del Moral et al., 2006] construction to define a pair (wk, zk) that is properly weighted for γk,
zk ∼ qk(· | zk−1), wk = vk wk−1, vk = γ̌k(zk, zk−1)
γ̂k(zk, zk−1) . (7)
We refer to the ratio vk as the incremental weight. In this construction, (wk, zk−1:k) is properly weighted for γ̌k which implies that (wk, zk) is also properly weighted for γk, since∫
dzk−1 γ̌k(zk, zk−1) = ∫ dzk−1 γk(zk)rk−1(zk−1 | zk) = γk(zk). (8)
Sequential importance sampling can be combined with other operations that preserve proper weighting, including rejection sampling, application of an MCMC transition operator, and importance resampling. This defines a class of samplers that admits many popular methods as special cases, including sequential Monte Carlo (SMC) [Doucet et al., 2001], annealed importance sampling (AIS) [Neal, 2001], and SMC samplers [Del Moral et al., 2006]. In this work we only consider variational transition kernels and importance resampling. These samplers vary in the sequences of densities they define. For example, AIS and SMC samplers both define intermediate density over a fixed common domain. Here, a common strategy is to define an annealing path γk(zk) = γ1(zk)1−βkγK(zk)βk for 0 = β1 < β2 < . . . < βK = 1. In contrast, when using SMC for state-space models we can define the intermediate target πk to be the the filtering distribution on the first k states. In this setting, the dimensionality of the support increases at each step and we can omit the reverse kernel rk−1. We are exploring both of these cases in Section 4.

heading: 3.1 Computing Gradient Estimates
text: The NVI objective can be optimized with respect to three sets of densities. We will use θk, φ̂k, and φ̌k to denote the parameters of the densities πk, qk, and rk respectively. For notational convenience, we use ρ̌k = {θk, φ̌k−1} to refer to the parameters of the reverse density π̌k, and ρ̂k = {φ̂k, θk−1} to refer to the parameters of the forward density π̂k,
Gradients of the Forward KL divergence. When we employ the forward KL as the objective, the derivative with respect to ρ̂k can be expressed as (see Appendix E.3),
− ∂ ∂ρ̂k
KL ( π̌k ∥∥ π̂k) = Ě
πk
[ ∂
∂ρ̂k log γ̂k
( zk, zk−1; ρ̂k )] − E πk−1 [ ∂ ∂ρ̂k log γk−1 ( zk−1; θk−1 )] . (9)
This case is the nested analogue of RWS-style variational inference. We can move the derivative into the expectation, since π̌k does not depend on ρ̂k. We then decompose log π̂k = log γ̂k − logZk−1 and use the identity from Equation 21 to express the gradient logZk−1 as an expectation with respect to πk−1. The resulting expectations can be approximated using self-normalized estimators based on the outgoing weights wk and incoming weights wk−1 respectively.
The gradient of the forward KL with respect to ρ̌k is more difficult to approximate, since the expectation is computed with respect to π̌k, which depends on the parameters ρ̌k. The gradient of this expectation has the form (see Appendix E.3)
− ∂ ∂ρ̌k
KL ( π̌k ∥∥ π̂k) = − Ě
πk
[ log vk ∂
∂ρ̌k log π̌k
( zk, zk−1; ρ̌k )] (10)
= − Ě πk
[ log vk ∂
∂ρ̌k log γ̌k
( zk, zk−1; ρ̌k )] + Ě πk [ log vk ] E πk [ ∂ ∂ρ̌k log γk ( zk; θk )] .
In principle, we can approximate this gradient using self-normalized estimators based on the outgoing weight wk. We experimented extensively with this estimator, but unfortunately we found it to be unstable, particularly for the gradient with respect to the parameters of the reverse kernel rk−1. For this reason, our experiments employ the reverse KL when learning reverse kernels.
Our hypothesis is that the instability in this estimator arises because the gradient decreases the probability of high-weight samples and increases the probability of low-weight samples, rather than the other way around. This could lead to problems during early stages of training, when the estimator will underrepresent low-weight samples, for which the probability should increase.
Gradients of the Reverse KL divergence. When computing the gradient of the reverse KL with respect to ρ̂k, we obtain the nested analogue of methods that maximize a lower bound. Here we can either use reparameterized samples [Kingma, Welling, 2013; Rezende et al., 2014] or likelihood-ratio estimators [Wingate, Weber, 2013; Ranganath et al., 2014]. We will follow Ritchie et al. (2016b) and define a unified estimator in which proposals are generated using a construction
wk = vkwk−1, zk = g(z̃k, φ̂k), z̃k ∼ q̃k(z̃k | zk−1, φ̂k), wk−1, zk−1 ∼ Πk−1. (11) This construction recovers reparameterized samplers in the special case when q̃k does not depend on parameters, and recovers non-reparameterized samplers when zk = z̃k. This means it is applicable to models with continuous variables, discrete variables, or a combination of the two. The gradient of the reverse KL for proposals that are constructed in this manner becomes (see Appendix E.2)
− ∂ ∂ρ̂k
KL ( π̂k ∥∥ π̌k) = Ê
πk
[ ∂
∂zk log γ̌k
( zk, zk−1; ρ̂k )∂zk ∂ρ̂k ] + Ê πk [ log vk ∂ ∂ρ̂k log γ̂k ( zk, zk−1; ρ̂k )] (12)
− Ê πk
[ log vk ] E
πk−1
[ ∂
∂ρ̂k log γk−1
( zk−1; θk−1 )] .
In this gradient, the first term represents the pathwise derivative with respect to reparameterized samples. The second term defines a likelihood-ratio estimator in terms of the unnormalized density γ̂k, and the third term computes the contribution of the gradient of the log normalizer logZk−1.
Computing the gradient of the reverse KL with respect to ρ̌k is once again straightforward, since we are computing an expectation with respect to π̂k, which does not depend on ρ̌k. This means we can move the derivative into the expectation, which yields a gradient analogous to that in Equation 9,
− ∂ ∂ρ̌k
KL ( π̂k ∥∥ π̌k) = Ê
πk
[ ∂
∂ρ̌k log γ̌k
( zk, zk−1; ρ̌k )] − E πk [ ∂ ∂ρ̌k log γk ( zk; θk )] . (13)
Variance Reduction. To reduce the variance of the gradient estimates we use the expected logincremental weight as a baseline for score function terms and employ the sticking-the-landing trick [Roeder et al., 2017] when reparameterizing the forward kernel as described in Appendix E.

heading: 3.2 Relationship to Importance-Weighted and Self-Normalized Estimators
text: There exists a large body of work on methods that combine variational inference with MCMC and importance sampling. We refer to Appendix A for a comprehensive discussion of related and indirectly related approaches. To position NVI in the context of the most directly related work, we here focus on commonly used importance-weighted and self-normalized estimators.
NVI differs from existing methods in that it defines an objective for pairs of variables (zk, zk−1) at each level of nesting, rather than a single objective for the entire sequence of variables (z1, . . . , zK). A standard approach for combining importance sampling and variational inference is to define an “importance-weighted” lower bound L̂K = log ẐK [Burda et al., 2016]. By Jensen’s inequality, E[L̂K ] ≤ logE[ẐK ] = logZK , which implies that we can define a stochastic lower bound using any properly-weighted sampler for γK , including samplers based on SMC [Le et al., 2018; Naesseth et al., 2018; Maddison et al., 2017]. For purposes of learning the target density γK , this approach is equivalent to computing an RWS-style estimator of the gradient,
∂
∂θK L̂K =
1
ẐK
1
S S∑ s=1 wsK ∂ ∂θK log γK(z s K ; θK). (14)
However, these two approaches are not equivalent for purposes of learning the proposals. We can maximize a stochastic lower bound to learn qk, but this requires doubly-reparameterized estimators
[Tucker et al., 2018] in order to avoid problems with the signal-to-noise ratio in this estimator, which can paradoxically deteriorate with the number of samples [Rainforth et al., 2018]. The estimators in NVI do not suffer from this problem, since we do not compute the logarithm of an average weight.
NVI is also not equivalent to learning proposals with RWS-style estimators. If we use sequential importance sampling (SIS) to generate samples, a self-normalized gradient for the parameters of qk that is analogous to the one in Equation 24 has the form
E πK ,rK−1,...,r1
[ ∂
∂φ̂k log qk(zk | zk−1 ; φ̂k) ] ' 1 ẐK 1 S S∑ s=1 wsK ∂ ∂φ̂k log qk(z s k | zsk−1 ; φ̂k). (15)
Note that this expression depends on the final weight wK . By contrast, a NVI objective based on the forward KL yields a self-normalized estimator that is defined in terms of the intermediate weight wk
E πk,rk−1
[ ∂
∂φ̂k log qk(zk | zk−1 ; φ̂k) ] ' 1 Ẑk 1 S S∑ s=1 wsk ∂ ∂φ̂k log qk(z s k | zsk−1 ; φ̂k). (16)
If instead of SIS we employ sequential importance resampling (i.e. SMC), then the incoming weight wk−1 is identical for all samples. This means that we can express this estimator in terms of the incremental weight vk rather than the intermediate weight wk
E πk,rk−1
[ ∂
∂φ̂k log qk(zk | zk−1 ; φ̂k)
] ' S∑ s=1 vsk∑S s′=1 v s′ k ∂ ∂φ̂k log qk(z s k | zsk−1 ; φ̂k). (17)
We see that NVI allows us to compute gradient estimates that are localized to a specific level of the sampler. In practice, this can lead to lower-variance gradient estimates.
Having localized gradient computations also offers potential memory advantages. Existing methods typically perform reverse-mode automatic differentiation on an objective that is computed from the final weights (e.g. the stochastic lower bound). This means that memory requirements scale as O(SK) since the system needs to keep the entire computation graph in memory. In NVI, gradient estimates at level k do not require differentiation of the incoming weights wk−1, This means that it is possible to perform automatic differentiation on a locally-defined objective before proceeding to the next level of nesting, which means that memory requirements would scale as O(S). It should therefore in principle be possible to employ a large number of levels of nesting K in NVI, although we do not evaluate the stability of NVI at large K in our experiments.

heading: 4 Experiments
text: We evaluate NVI on three tasks, (1) learning to sample form an unnormalized target density where intermediate densities are generated using annealing, (2) learning heuristic factors to approximate the marginal likelihood of future observations in state-space models, and finally (3) inferring distributions over classes from small numbers of examples in deep generative Bayesian mixtures.

heading: 4.1 Sampling from Multimodal Densities via Annealing
text: A common strategy when sampling from densities with multiple isolated modes is to anneal from an initial density γ1, which is typically a unimodal distribution that we can sample from easily, to the
target density γK , which is multimodal [Neal, 2001]. Recent work on annealed variational objectives (AVOs) learns forward kernels qk and reverse kernels rk−1 for an annealing sequence by optimizing a variational lower bound at each level of nesting [Huang et al., 2018].
max qk,rk LAVOk , LAVOk = E q1,...,qk
[ log vk ] , γk(z) = π1(z) 1−βkγK(z) βk , k = 1, . . . ,K. (18)
NVI allows us to improve upon AVO in two ways. First, we can perform importance resampling at every step to optimize an SMC sampler rather than an annealed importance sampler. Second, we can learn the annealing schedule (β1, . . . , βK) as part of the intermediate densities γk such that intermediate densities are scheduled more equidistantly in terms of KL-divergence.
We illustrate the effect of these two modifications in Figure 3, in which we compare AVO to NVI with resampling and a learned path, which we refer to as NVIR∗. Both methods minimize the reverse KL at each step 1. The learned annealing path in NVIR∗ results in a smoother interpolation between the initial and final density. We also see that AVO does not assign equal mass to all 8 modes, whereas NVIR∗ yields a more even distribution.
In Figure 4a we compute the reverse KL between targets at each step in the sequence. For the standard linear annealing path, the KL decreases with k, suggesting that later intermediate densities are increasingly redundant. By constrast, in NVIR∗ we see that the KL is approximately constant across the path, which is what we would expect when minimizing D with respect to βk. This is also the case in an ablation without resampling, which we refer to as NVI∗.
1As noted in Section 3.1, we found optimization of the forward KL to be unstable when learning rk.
Figure 5: Ground truth densities (GT) and samples from final target density for NVIR∗ with 2 intermediate densities (K=4).
Figure 4b shows a rolling average of the ESS and its variance during training. We compare NVI-based methods to SVI and a variational SMC sampler [Le et al., 2018; Maddison et al., 2017; Naesseth et al., 2017]. NVIR* has consistently higher ESS and significantly lower variance compared to baselines. These plots also provides insight into the role of resampling in training dynamics. In NVI∗, we observe a cascading convergence, which is absent in NVIR∗. We hypothesize that resampling reduces the reliance on high-quality proposals from step k − 1 when estimating gradients at step k. Annealed NVI has similar use cases as normalizing flows [Rezende, Mohamed, 2015]. Inspired by concurrent work of Arbel et al. (2021), which explores a similar combination of SMC samplers and normalizing flows, we compare flow-based versions of NVI to planar normalizing flows, which maximize a standard lower bound (SVI-flow). We find that a normalizing flows can be effectively trained with NVI, in that samplers produce better estimates of the normalizing constant and higher ESS compared to SVI-flow. We also find that flow based models are able to produce high-quality samples with fewer intermediate densities (Figure 8). Moreover, we see that combining a flow-based proposal with learned βk values (NVI∗-flow) results in a more accurate approximation of the target than in an ablation with a linear interpolation path (AVO-flow)2.
In Table 1 we report sample quality in terms of the stochastic lower bound L̂K = log ẐK and the effective sample size ESS = ( ∑ s w s K) 2/ ∑ s(w s K)
2. The first metric can be interpreted as a measure of the average sample quality, whereas the second metric quantifies sample diversity. We compare NVI with and without resampling (NVIR∗ and its NVI∗) to ablations with a linear annealing path (NVIR and NVI), AVO, and a standard SVI baseline in which there are no intermediate densities. We additionally compare against AVO-flow and NVI∗-flow, which employ flows. We observe that NVIR∗ and NVI∗-flow outperform ablations and baselines in terms of log Ẑ, and are competitive in terms of ESS. We show qualitative results for two additional target densities in Figure 5. For more details we refer to Appendix G.1.

heading: 4.2 Learning Heuristic Factors for State-space Models
text: Sequential Monte Carlo methods are commonly used in state-space models to generate samples by proposing one variable at a time. To do so, they define a sequence of densities πk = γk/Zk on the first k time points in a model, which are also known as filtering distributions,
γk(z1:k, η) = p(x1:k, z1:k, η) = p(η) p(x1, z1 | η) ∏k
l=2 p(xl, zl | zl−1, η), Zk = p(x1:k). (19)
Here z1:k and x1:k are sequences of hidden states and observations, and η is a set of global variables of the model. The densities γk differ from those in the annealing task in Section 4.1 in that the dimensionality of the support increases at each step. We define a forward density γ̂k that combines the preceding target γk−1 with a proposal qk for the time point zk, and define a reverse density γ̌k = γk that is equal to the next intermediate density (which means that we omit rk−1),
γ̌k(z1:k, η) = γk(z1:k, η), γ̂k(z1:k, η) = qk(zk |xk, z1:k−1, η) γk−1(z1:k−1, η). (20)
A limitation of this construction is that the filtering distribution πk−1 is not always a good proposal, since it does not incorporate knowledge of future observations. Ideally, we would like to define
2AVO-flow is itself a novel combination of AVO and flows, albeit an incremental one.
intermediate densities γk(z1:k, η) = p(x1:K , z1:k, η) that correspond to the smoothing distribution, but this requires computation of the marginal likelihood of future observations p(xk+1:K | zk, η), which is intractable. This is particularly problematic when sampling η as part of the SMC construction. The first density π1(z1, η) = p(z1, η |x1) will be similar to the prior, which will result in poor sample efficiency, since the smoothing distribution p(z1, η |x1:K) will typically be much more concentrated. To overcome this problem, we will use NVI to learn heuristic factors ψθ that approximate the marginal likelihood of future observations. We define a sequence of densities (γ0, . . . , γK),
γ0(η) = p(η) ψθ(x1:K |η), γk(z1:k, η) = p(x1:k, z1:k, η) ψθ(xk+1:K | η), k = 1, 2, ...,K. Our goal is to learn parameters θ of the heuristic factor such that that intermediate densities approximate the smoothing distribution. This approach is similar to recently proposed work on twisted variational SMC [Lawson et al., 2018], which maximized a stochastic lower bound.
To evaluate the this approach, we will learn heuristic factors for a hidden Markov model (HMM). While HMMs are a well-understood model class, they are a good test case, in that they give rise to significant sample degeneracy in SMC and allow to compute tractable heuristics. We optimize an NVI objective based on the forward KL with respect to the heuristic factor ψθ, an initial proposal q0(η |x1:K ;φ) and a forward kernel qk(zk |xk, zk−1, η;φk). Figure 6 shows qualitative and quantitative results. We compare NVIR∗ with a neural GMM-heuristic to a baseline without a heuristic and a baseline that uses a hand-coded GMM-heuristic. In addition, we train models which employ the convex update operator proposed in the Automatic Structured Variational Inference (ASVI) [Ambrogioni et al., 2021]. Training ASVI by optimizing the ELBO failed due to the high variance of the score function estimator, which is needed for the discrete latent variables. However, we were able to train ASVI-based models with NVIR∗. We refer to this novel combination as NASVI. Because NVI let us treat each step as a separate optimization problem, we also compare partial optimization with respect to γ̂k only to full optimization with respect to both γ̂k and γ̌k. While full optimization yields poor results, partial optimization learns a neural heuristic whose performance is similar to the GMM heuristic, which is a strong baseline in this context. For more details see Appendix G.2.

heading: 4.3 Meta Learning with Deep Generative models
text: In this experiment, we evaluate NVI in the context of deep generative models with hierarchicallystructured priors. Concretely, we consider the task of inferring class weights from a mini-batch of images in a fully unsupervised manner. For this purpose, we employ a variational autoencoder with a prior in the form of a Bayesian Gaussian mixture model. Unlike the previous experiments, here we both train a generative and an inference model. The forward kernels in this setting act as the generative model and the reverse kernels are the inference model. Moreover, in this model sampling simplifies because we design the intermediate densities to be tractable, hence no importance sampling construction is needed. Regardless, we still can construct an NVI objective to train hierarchical deep generative models more efficiently and with a more accurate posterior inference.
Model Description. We define a hierarchical deep generative model for batches of N images of the form (see Appendix G.3 for a graphical model and architecture description)
λ ∼ Dir(· ; α) cn ∼ Cat(· |λ) zn ∼ N (· |µcn , 1/τcn) xn ∼ p(· | zn; θ) for n = 1 . . . N. Here λ, cn, zn, and xn refer to the cluster probabilities, cluster assignments, latent codes, and observed images respectively. In the likelihood p(x|z; θx), we use a convolutional network to parameterize a continuous Bernoulli distribution [Loaiza-Ganem, Cunningham, 2019]. We define
proposals q(z|x;φz), q(c|z;φc), and q(λ|c;φλ), which are also parameterized by neural networks. We refer to this model as Bayesian Gaussian mixture model VAE (BGMM-VAE).
Objective. To construct an NVI objective, we define intermediate densities for c and z. Unlike in previous experiments, we employ tractable densities in the form of a categorical πc(c; θc) for cluster assignments and a 8-layer planar flow πz(z; θz) for the latent codes. Subsequently, we can define the forward and reverse densities at each step as:
π̂2(λ, c) = π(λ)p(c|λ), π̂3(c, z) = π(c)p(z|c), π̂4(z, x) = π(z)p(x|z), π̌2(λ, c) = π(c)q(λ|c), π̌3(c, z) = π(z)q(c|z), π̌4(z, x) = π̃(x)q(z|x),
where π̃(x) is an empirical distribution over mini-batches of training data. We minimize the forward KL for the first two steps and the reverse KL at the final step.
Since the intermediate densities are tractable in this model, no nested importance sampling is required to optimize this nested objective; we can compute gradient estimates based on a (single) sample from p(λ)p(c|λ) in the first term, πc(c)q(λ|c) in the second, and p̂(x)q(z|x) in the final term. To learn the parameters {µ, τ, θx} of the generative model, we maximize a single-sample approximation of a lower bound L = Eq [ log ( p(x, z, c, λ) / q(z, c, λ|x) )] .
Results. We evaluate NVI for the BGMM-VAE using the following procedure. We generate minibatches with a sampled λ∗ (for which we make use of class labels that are not provided to the model). We then compute the expectation of λ under q(λ, c, z|x) by sampling from the inference model, and compare this value against λ∗. Additionally, we generate a new mini-batch given the inferred λ by running the generative model forward. We compare NVI against RWS, where we use 10 samples to estimate the gradients at each iteration. The results are shown in Figure 7. The cluster indices are rearranged based on the association of clusters to true classes. We observe that for RWS, even though the sample quality is reasonable, the posterior approximation is quite poor. When we train with NVI however, both the inferred λ and the generated samples match the test mini-batch reasonably well.

heading: 5 Conclusion
text: We developed NVI, a framework that combines nested importance sampling and variational inference by optimizing a variational objective at every level of nesting. This formulation allows us to learn proposals and intermediate densities for a general class of samplers, which admit most commonly used importance sampling strategies as special cases. Our experiments demonstrate that samplers trained with NVI are able to outperform baselines when sampling from multimodal densities, Bayesian state-space models, and hierarchical deep generative models. Moreover, our experiments show that learning intermediate distributions results in better samplers.
NVI is particularly useful in the context of deep probabilistic programming systems. Because NVI can be applied to learn proposals for a wide variety of importance samplers, it can be combined with methods for inference programming that allow users to tailor sampling strategies to a particular probabilistic program. Concretely, NVI can be applied to nested importance samplers that are defined using a grammar of composable inference combinators [Stites et al., 2021], functions that implement primitive operations which preserve proper weighting, such as using samples from one program as a proposal to another program, sequential composition of programs, and importance resampling.

heading: Acknowledgements
text: This work was supported by the Intel Corporation, the 3M Corporation, NSF awards 1835309 and 2047253, startup funds from Northeastern University, the Air Force Research Laboratory (AFRL), and DARPA.

Review summaries:
The paper describes a framework for learning and inference in probabilistic models with structured latent spaces. It operates in the tradition of IWAE and VSMC-like methods that use importance sampling to tighten a variational bound (or equivalently, uses a variational bound as an objective for learning importance sampler proposals). Unlike methods that optimize a single global bound, NVI is defined in terms of local pairwise divergences between a sequence of target distributions, which acts to decouple the stochastic optimization, reducing gradient variance and potentially offering a strategy for parallel inference.

The author proposes a nested variational inference (NVI) framework that combines nested importance sampling and variational inference. This framework allows the training of the important sampling proposal by minimizing KL divergences at each level of nesting. This also provides us a way to learn intermediate densities, which serves as a heuristics for guiding the importance sampler. Empirically, the author evaluates NVI by (1) sampling from a multimodal distribution with a learned annealing path; (2) learn heuristics for future observations in the state-space model; (3) perform amortized inference in the hierarchical deep generative model.

This paper proposes an inference procedure that combines ideas from variational and nested inference, resulting in a sequence of forward and reverse proposal distributions that are optimized with local objectives that encourage the forward and backward edge marginals to be similar. This allows the objective and therefore gradients to decouple. The addition of resampling, afforded by the connection to nested inference, additionally improves sample diversity.

The authors present a method to jointly learn a series of proposals and intermediate densities for e.g. SMC, or annealed importance sampling. They use a novel objective which sums divergences between distributions over each adjacent pair of variables in the chain, and argue that this leads to lower-variance gradient estimates than prior work which uses a single objective for the entire sequence of variables. They show experimentally the benefits of this method in three different settings.

The authors introduce a new inference and approximate weighted sampling framework that combines f-divergence variational inference and nested importance sampling. An original variation sampling density is transformed through a series of weight preserving operators that map the proposal to an approximation of an unnormalized target distribution. On the other hand, reverse transformations are use to transform the target back to the sampling distribution. The transformations are trained by minimizing an f-divergence at each nesting level.