Nested Variational Inference

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.


Reviews and author responses:

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.  # Novelty
The method introduced is novel AFAIK, and well situated with regard to prior work.

# Strengths
- The paper is well-written (modulo some minor points below).
- The paper presents a convincing case in Section 3.3 for why NVI's gradient estimates could be expected to be low-variance.
- The experiments are diverse and the results are convincing - particularly for the annealing experiment in 4.1.

# Questions/weaknesses
- It is counter-intuitive to me that "full optimization yields poor results" compared to partial optimization in Section 4.2. Do the authors have any idea why this might be?
- In Table 1, very similar results are reported for several methods (particularly AVO-flow and NVI*-flow). It would be helpful if confidence intervals were included, or if the authors could at least comment on the significance of the differences. Similarly for $\log \hat{Z}$ and ESS in Figure 6.

# Presentation
- In Section 3, I think it should be made clearer than $\gamma_k$ can be learnable. This was a point of confusion for me until reaching the experiments section.
- In Section 3, line 82, it is confusing when you say you can optimize Equation 6 purely with respect to $q_1$ and $\pi_k$, since Equation 6 also depends on $r_k$ and so it is unclear that is done with $r_k$ in this case. After reading Section 4.2, I am assuming that this is only done in problems where $r_k$ is omitted (as in Equation 20), but it would be helpful if this was made clearer.
- I am a bit confused about that is in bold in Table 1 - I am not sure why nothing is in bold for the $K=2$ columns and think the $\log \hat{Z}$ entry for AVO-flow and $K=4$ should be bold.

EDIT: Having read the other reviewer's comments, I agree with them that the paper could benefit from a significantly more focused presentation and more systematic experiments. I have therefore lowered my score to a 6. Yes.

 We understand these comments, yet we respectfully disagree that this paper needs a more focused presentation, or that positioning relative to related work is fundamentally insufficient in a manner that would necessitate another round of review.

This is a paper that builds on work that combines importance sampling and variational inference. There are many examples of such methods, which each have different use cases, and we believe our experiments are representative of these use cases. We therefore disagree that these experiments lack focus. 

We discuss what we consider to be the most directly relevant work in the main text, along with a very extensive related work section in the Appendix. The work on structured variational inference indeed considers structured models, and is somewhat related to NVI in this respect. However, as the reviewer agrees, the technical contributions in this line of work are mostly orthogonal to the contributions to this paper. Moreover, here we consider hierarchical structure that requires "meta" reasoning about variables that are shared across a batch of inputs, which is not really a use case in work on structured variational inference. We are therefore happy to discuss this work, but we do not agree that this discussion requires fundamental reframing of our work. 

Most of the methods in past work on structured variational inference are in fact not directly applicable as baselines for the models in our experiments. This further suggests that connections to this line of work are somewhat indirect. The only method that we believe we can apply is ASVI, which one would naively expect to perform poorly (given the need for score-function estimators). We are more than happy to follow up and perform this comparison, but we emphatically disagree that our experiments lack appropriate baselines.

We do believe we understand the sentiment of this review. The class of objectives considered by NVI is quite general, so we can empathize with the feedback that the paper feels unfocused. Perhaps the reviewer is also concerned that the paper could be interpreted as a bit of a "land grab". This is not our intention. Our goal with this paper is to say "here is a new general way to think about objectives for methods that combine importance sampling and variational inference, and here are some illustrative use cases of such methods". 

The reviewers agree that this work is novel, technically sound, and that there are no fundamental problems with clarity. We believe that the experiments are sufficiently illustrative of use cases and think that further research should determine how useful these methods are in practice, and whether their apparent complexity is worth it in terms of payoffs? Asking for a paper to demonstrate this conclusively seems like setting a very high bar. Very few papers truly meet this bar in practice, especially when looking back a few years onward.


 Thanks for the reply. As things are, I’m inclined to keeping my score. While I think the method has potential, I think the paper would require many changes for me to recommend acceptance. I repeat some of the comments (also in the review) which I consider to be the most relevant:

- 2D example, while relevant due to multi-modality, not enough. I think adding results similar to the ones for experiment 1 in several high dimensional models is necessary.

- I think that an analysis for the biased gradients, even if empirical, would be great. (Taking into account that this bias issues may be significantly more severe in high dimensions.) 

- For plain inference tasks, comparison against AIS with HMC, which is known to work quite well for a wide range of problems, would be good. I agree that running HMC may be expensive, but given the idea behind NVI this seems quite natural to me. (Even if HMC is not fully-tuned; e.g. just the step-size to achieve a reasonable rejection rate.)

- I tend to agree with the comments by reviewer 4j52 about the empirical evaluation and the presentation being too generic. I think the latter point may make it hard to map NVI to each of the settings considered in the experiments.

I’m willing to discuss this points further or provide more details.

 I do agree that these papers are orthogonal to your contribution. My point was that, if you wish to argue that one of main use cases of the approach is structured inference, then you need to discuss the existing literature in the main text and include appropriate baselines in the experiments. I do appreciate the suggestion of including ASVI as baseline and to discuss the relations with the other methods, however I am not sure that such a revision is possible in the context of this submission as it required a decently large revision including new experiments and a more focused presentation that directly addresses the use cases with their related literature. 

 Thank you for taking the time to read our response and for your kind words about the potential of our work.

We agree that these papers are related and should be discussed in our paper; thank you for the suggested references. That being said, we consider this line of work a somewhat orthogonal approach to NVI. While the cited papers are about designing better variational proposals to capture the structure of the generative model, NVI is about breaking down the inference problem itself into multiple ‘local’ inference problems by modifying the variational objective. Therefore, while we do believe that both of these approaches are valuable, it may be hard to compare them directly without confounding the choice of model and objective/inference approach.

In other words, contrary to the cited papers, the choice of variational family/model is not dictated by our approach. In fact, some of the constraints of the cited methods make it difficult/not possible to apply them to the tasks in our experiments:
- SSVI exploits locally conjugate models using conjugate update rules
- Cascading flows (CF) can, to the best of our knowledge, not model discrete variables, which we use in all experiments but the fist
- ASVI can be applied to our time series model and BGMM. However, it requires to compute a score function gradient estimator, that usually has high variance, which can be avoided using NVI

We will add ASVI as a baseline method to our experiments.

 Thank you for the response. 

There is a large literature on structured variational inference (see for example [1-5]) which should be throughly discussed and  included as baseline if the authors think that the structured setting is the main use case. From the current manuscript, it is unclear if the method can improve the state of the art ins structured variational inference.

I do stand by all points I raised in my original review. However, I do think that the paper has merit and that it definitely has the potential to enter in a top conference if the presentation becomes more focused on a well-specified algorithm with clear use cases. 

[1] Rezende, Danilo, and Shakir Mohamed. "Variational inference with normalizing flows." ICML, 2015.
[2] Ranganath, Rajesh, Dustin Tran, and David Blei. "Hierarchical variational models." ICML, 2016.
[3] Hoffman, Matthew, and David Blei. "Stochastic structured variational inference." AISTATS, 2015.
[4] Ambrogioni, Luca, et al. "Automatic structured variational inference." AISTATS, 2021.
[5] Ambrogioni, Luca, Gianluigi Silvestri, and Marcel van Gerven. "Automatic variational inference with cascading flows." ICML, 2021.

 We appreciate the thoughtful and constructive comments.

### Comparison to diffusion models
Thank you for pointing this out! There are indeed many similarities between our current framework and diffusion models. One key difference is that there are no intermediate densities in diffusion models. We believe that diffusion models can benefit from training with NVI and are planning to do some follow-up work on this end.

### Rectification of NVI's complexity compared to other standard VI methods
Our methodology combines methods from importance sampling and variational inference, which allows for more tailored inference approaches in models with structured representation, which are common in probabilistic programming or meta-learning approaches. Here standard amortized VI methods often underperform due to their difficulty of exploiting the underlying structure of the model. The time-series inference task and the BGMM task illustrate such structured models. Generally, in which other applications the benefit outweighs the cost of this technical complexity remains to be shown in future work.

### Overgenerality of the presented method

We present a method for general probabilistic inference. As such, we believe that the generality of our method is something positive and intentionally did not choose to explore only a specific domain or experimental setting, e.g. focusing on exploring different annealing paths, but to showcase the versatility of the proposed methodology.

### Not committing to any specific transformation operators

While our presentation is general (see reasons above), in our experiments we use variational distributions to define the transition kernels. We view the ability to use different transition operators as a strength of our method.
  
### Unfocused non-systematic experiments / baselines
We respectfully disagree with the observation that our experiments are unfocused and non-systematic. On the contrary, we believe that our experiments represent a set of diverse and interesting use cases, which highlight the strength of NVI in structured models.
In the case of the first experiment, while there exist other state-of-the-art methods, we chose  to compare to AVO due to its technical similarity. This allows us to showcase the strengths of our work compared to previous, technically similar contributions.




 We appreciate the thoughtful and constructive comments.

### Usage of MCMC kernels

Running MCMC chains is generally computationally expensive and often comes with certain restrictions, e.g. standard HMC is restricted to continuous domains. Variational kernels are more flexible in this regard and have the advantage of, once sufficiently trained, providing a computationally effective way of generating samples. Nevertheless, we agree that incorporating additional MCMC transition kernels might be useful to further refine samples. As you pointed out, such transition kernels can be easily incorporated when tuning hyperparameters is not desired. The difficulty of tuning hyperparameters of such kernels, due to the tractability issues resulting from the MH correction step, is not unique to our methodology and has long been subject of research in the community (see Appendix A for a brief discussion).
  

### Comparing against SMC / AIS.

We want to point out that we compare our method to variational SMC in the second experiment.

In the first experiment we are specifically interested in learning the parameters of the variational transition kernel and hence mainly compare with AVO and AVO-flow, a custom baseline inspired by AVO. However, we do agree that an additional comparison to an AIS baseline with a hand tuned MH transition operator might be an insightful addition to Table 1 and we will include this baseline in the updated manuscript.

  

### Objective in terms of stochastic lower bounds

The high level intuition behind our work can be described as breaking one large (difficult) inference problem to several smaller (easier) inference problems. In general, NVI objective (Eq 6) does not define a bound on the marginal likelihood given that it is defined with a series of f-divergences. However, if we choose to optimize a KL-divergence, this corresponds to optimizing an ELBO or EUBO (depending on if we are optimizing a reverse or forward KL-divergence) at the last level of nesting. This holds for both, (1) when the final target density $\pi_K(z_K)$ corresponds to some posterior $p(z_K \mid x)$ and $\gamma_K(z_K)=p(z_K, x)$, and (2) when the $z_K$ corresponds to the data $x$ itself (e.g. in the 3rd experiment).

For example, in the case of (1) where $\pi_K(z_K) = p(z_K|x)= p(z_K,x)/p(x)$, we have that

$$\min KL(\pi(z_{K-1})q_K(z_K|z_{K-1})||p(z_K|x)r_{K-1}(z_{K-1}|z_K)) = \min \underbrace {\mathbb E_{\hat \pi_k} \left[ \log \frac{\pi(z_{K-1})q_K(z_K|z_{K-1})}{p(x,z_K)r_{K-1}(z_{K-1}|z_K)} \right]}_{-ELBO}$$

  

### Experiment to quantify the bias in the self-normalized estimators

We think adding an additional experiment that aims to quantify the bias for different numbers of samples is a good suggestion. We will add this to the appendix in the updated version of the manuscript.

### Experiments

 **1) Suggested ablation study**

We tried to do something akin to the requested ablation study in the first experiment.
- Tuning kernels only -> NVI
- Tuning kernels + resampling -> NVIR
- Tuning kernels and intermediate densities -> NVI*
- Tuning kernels and intermediate densities + resampling -> NVIR*
- Effect of IS (‘proper weighting preserving transitions’) -> AVO vs NVI

  

**2) Tasks & Datasets**

While it is true that a 2D example might not be representative of higher dimensional problems, it allows us to run and compare a lot of different models due to the comparatively low computational cost. We consider it a good first experiment to build intuition on how our method behaves in comparison to other technically similar methods. That being said, we will add an additional experiment on a non-synthetic data set, and include further results for the third experiment. Thank you for the reference.


**3) Mapping of BGMM-VAE to NVI**

The intermediate densities are $p(\lambda), \pi_{c}(c), \pi_{z}(z), \hat p(x)$, the forward kernels $q$ (the generative model) are $p(c|\lambda)p(z|c)p(x|c)$, and the reverse kernels $r$ (the inference model) are $q(z|x)q(c|z)q(\lambda|c)$. The reason we switched to this notation was because of the commonality of using $p$ as the generative model and $q$ as the inference model in the VAE literature. We agree, this is confusing and we will use the same notation in Section 3 for the BGMM-VAE experiment in the revised manuscript.

  

**4) Training a 'normal' VAE**

As stated above, In the case of vanilla VAE ($K=2, z_K=x$), optimizing the NVI objective is no different than optimizing the ELBO. The advantages of NVI are more highlighted when the model is hierarchical, which is the reason we decided to employ a BGMM-VAE in our last experiment.

  

### Clarity
Thank you for your suggestions. We will add discussion regarding the transition kernels and more detailed description of the tasks in the updated manuscript.

  
### Additional discussion on limitations
Another limitation which we briefly point out in the self-normalized importance sampling section in the background, arises from the fact that for small sample sizes our estimator might initially be significantly biased. However, as long as the gradient estimator moves the proposal closer to the target the bias and variance of the estimator typically decrease and the corresponding NVI step converges to the correct target. This also improves the gradient estimates for the subsequent NVI steps as they obtain better samples from the previous step. This behaviour can be observed in Figure 1 (on the very left).

We will add additional discussion regarding the limitations in the updated version of the manuscript.


 Thank you for your review. We are delighted by your positive feedback.
  
### Full optimization yields poor results

Computing the full gradient requires computing a score-function gradient estimate, which, despite variance reduction techniques, yields high variance compared to partial optimization. We hypothesize that this is the reason for the poor performance of the joint optimization scheme.

### Significance of results in Table 1

The standard deviations of the numbers reported in Table 1 do not vary significantly for comparable methods, which is why we decided to omit them in Table 1 to fit the table next to the figure. We’ll try to include the standard deviation in the table in the updated manuscript or alternatively comment on the significance of the presented results in the text as suggested.

### Presentation

Thank you for your helpful comments regarding the presentation. We will clarify which parameters are optimized and can be learned in each case in the revised manuscript.

### No bold numbers in table for $K=2$

For $K=2$, AVO, NVI and SVI correspond to the same method (essentially a standard VAE), hence we refrained from highlighting any results here.


The paper proposes Nested Variational Inference (NVI), a method to train (unnormalized) intermediate distributions between some simple initial approximation and an intractable target. The formulation is done in a general way, and compatible with several algorithms, such as annealed importance sampling (AIS) and sequential monte carlo (SMC).  Strengths:
- Up to my knowledge the method proposed in the paper, NVI, is novel.
- In some ways the method is close to AVO (Huang et al 2018), it proposes a sequence of K distributions bridging some simple initial approximation with the target, and proposes to learn forward and reverse kernels between bridging densities. These kernels are conditional densities differentiable in its parameters. NVI has two main advantages over AVO: it can learn parameters of the intermediate densities (e.g. annealing coefficients), and it can be combined with proper weighting preserving operations (such as MCMC kernels). I believe that the latter could possibly be used with AVO too, though the original AVO paper does not address this is any way.
- Compared to AIS (or SMC), an advantage is that NVI allows the learning of intermediate densities' parameters, and the parameters of some additional transition kernels.


Weaknesses:
- As I understand it, NVI cannot tune parameters of MCMC kernels (since these usually require a metropolis-hastings correction). However, the authors show that NVI may be combined with MCMC kernels or re-sampling (without tuning their parameters) in a simple way, since these transitions preserve "proper weighting". However, I'm not completely clear on how much the kernels allowed by NVI (q and r in the paper) would actually help, when compared to the addition of powerful MCMC kernels, such as HMC, used to improve samples while preserving proper weighting. This of course depends on the form chosen for q and r, but could be checked for a few simple choices (e.g. Gaussians as in the first set of experiments).
- Related to the comment above, I think it would be great if the paper included an empirical comparison against powerful methods such as SMC and AIS. Different questions could be addressed: How much tuning the bridging coefficients using NVI helps over AIS (all other things, like MCMC transitions used, kept constant)? How much does having additional learnable kernels (such as Gaussian with learnable mean and covariance) help?
- The method uses biased gradient estimators to tune all of its parameters. I understand that this bias will decrease as the number of intermediate densities is increased, or if powerful "proper weighting" transitions are used. However, I think that a discussion and an analysis to understand the effect of this bias be useful. The analysis could even be a simple empirical evaluation, comparing the method's final performance when using more samples to estimate gradients (or any other way to reduce bias) in a few different settings (preferably high dimensional, self-normalized importance sampling is known to have issues in high dimensions).
- The objective being optimized (eq. 6), while intuitive, is not really justified (can it be mapped to some bound as the typical ELBO?).


Comments on empirical evaluation:
- As I see it, the method does three things jointly: (1) tuning intermediate densities; (2) tuning some transition kernels; (3) allows the use of proper weighting preserving transitions. Other methods achieve only subsets of these (e.g. AVO only (2), AIS and SMC (3), and there are some methods for (1)). It would be nice to have different experiments to verify how much benefits are obtained from each of (1) (2) and (3), and how this changes as more bridging densities are used.
- While the mixture of Gaussians is a nice example and allows for nice visualizations, it is a 2 dimensional toy problem. It would be great to see comparisons in more plain inference tasks using real models in higher dimensions. I really like all the metrics and results reported for the task, but I think that adding other realistic tasks (e.g. could possibly check models in the inference gym [1], hierarchical regression, bayesian linear regression, etc) would improve this part of the empirical evaluation significantly. (Second experiment is also on synthetic data.)
- For the BGMM-VAE, the mapping to the NVI framework was not very clear to me.
- Overall, I think it would be nice to have an empirical evaluation on more tasks used by previous work (and with real models and data), using the typical performance measures. For instance, the method could be used to train a "normal" VAE, that is, using only z as latent variables (with a std Gaussian prior). Does NVI lead to better log-likelihoods than other methods typically used for the task? (e.g. a plain diagonal Gaussian approximate posterior, importance weighting, ...).


Clarity: I think there's room for improvement. Some comments (and personal suggestions) next:
- While the whole formulation is done in a general way, I think that some instantiation of the ideas at some early point in the paper could help. For instance, right before Section 3.1 a paragraph giving a possible choice for the forward and reverse kernels could be given (e.g. Gaussian with mean and variance as learned functions of the previous samples). I think this could help, specially because the terminology forward and reverse kernels is also used in the context of AIS and SMC, where they are taken to be MCMC kernels. As I understand, in this work the kernels are required to be differentiable (MCMC kernels are not), thus the requirements on the kernels for each method is quite different. In addition, MCMC transitions may be combined with NVI, but they would not play the role of the transition kernels used in the paper. Personally, I think adding a few comments on this may help.
- The experimental section could include a more thorough description of the tasks in the main text.
- Could include a more explicit mapping to the NVI framework using the notation from Section 3 for the BGMM-VAE experiment, specifying clearly what are the intermediate distributions (\hat \pi, and \check \pi (inverted hat)), forward/reverse kernels (q and r), etc.

[1] https://github.com/tensorflow/probability/tree/main/spinoffs/inference_gym One limitation the authors mention is that the use of forward KL divergence tends to be unstable in some scenarios. Is there an additional discussion about limitations that I'm missing?

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.   The paper integrates known techniques such as variational inference and nested importance sampling in an interesting way. The sequential nature of the transformation from a simple sampling density to a complex distribution reminds me of diffusion-based generative models and I would be interested in seeing if this kind of important sampling approach can be useful in that context.

However, the proposed approach is rather complicated and the paper fails to show convincing use cases where this complexity is needed. The most obvious use case for the approach is to learn paths for annealed importance sampling. If this were the goal of the authors, the paper should have covered the theory of annealed importance sampling in greater details and compared the new approach with the existing alternatives. The experiment section should have accordingly been focused on comparing the relative performance of different AIS paths in a wide range of relevant inference problems. 

In general, the presentation of the method is too generic. The authors do not commit to any specific transformation operators. The result is a very large family of possible methods, each with very different performances and characteristics. It is really difficult therefore to evaluate the proposed method. 

The experiments are unfocused ad non-systematic. The experiment with multimodal densities is relevant. However, the set of baselines is rather idiosyncratic and it does not represent the current state-of-the-art in density estimation. The experiment with state-space models is interesting but It si in my opinion outside the main scope of the paper. Application of the method to a Bayesian filtering setting introduces several additional complexities that need to be addressed in the body of the paper. I would suggest to cover the filtering application in an independent follow up work.  The authors discuss the limitations of their approach only briefly. 

Meta-review: The paper introduces an approach to sampling from intractable distributions (NVI) by learning sequences of proposal distributions for nested importance sampling, with the proposal at each level of nesting learned by minimizing the (local) reverse KL-divergence.

The reviewers found the method interesting and novel but felt that the presentation was too unfocused, and the evaluation, while broad, insufficiently systematic to be convincing. In particular, they found the annealed sampling study in Section 4.1 informative but thought that it was compromised by the use of a toy dataset. The paper would be strengthened by including a more realistic dataset in the study, a clearer motivation for the choice of the baselines (i.e. making it clear that the baselines correspond to ablations), and providing more detail about how exactly NVI is applied in each setting. In summary, the paper clearly has a lot of potential, but needs some more work to deliver on it.

Decision: accept