Paper: Forward $\chi^2$ Divergence Based Variational Importance Sampling (ICLR-2024)

Reviewer comment (Evaluation):
There is no discussion on the limitation of the proposed method.

Author response:
**Discussion on limitation of VIS**: The last sentence of the discussion section does point out a limitation of VIS. "Nevertheless, it is worth noting that while this choice of proposal distribution in VIS is statistically optimal for importance sampling, its practical significance in certain real-world applications might require further investigation and validation." Therefore, the proposal distribution learned by our VIS method still functions as an approximated posterior distribution. It is closest to the true posterior in terms of forward $\chi^2$ divergence, which differs from the variational distribution learned by VI (maximizing ELBO) in terms of reverse KL divergence.

Reviewer comment (Evaluation):
Increase in training time based on the additional sampling.

Author response:
* When the problem itself is complicated (e.g., complicated generative process and/or high latent dimensionality such as the POGLM model in Sec 4.3), a large number of Monte Carlo samples is indeed necessary, no matter for VIS or other baseline methods including VI. When working with VAE, people are already accustomed to using one Monte Carlo sample, but more Monte Carlo samples certainly improves the quality of $\widehat{\operatorname{ELBO}}$ since $\operatorname{ELBO}$ is also an integral expression w.r.t. $q(\boldsymbol z|\boldsymbol x;\phi)$. Anyway, however, employing a larger number of Monte Carlo samples always improves the estimation for all sampling-based estimation methods, including VI as depicted in Fig. 1. The main conclusion from Fig. 1 is that increasing the number of Monte Carlo increases the effectiveness of both $\widehat{\operatorname{ELBO}}$ and $\ln \hat p$, but it is only very helpful for tightening the bound of $\ln \hat p$ (with the speed of order $\frac 1 K$) rather than $\widehat{\operatorname{ELBO}}$. As we mentioned below Eq. 8, when the number of Monte Carlo samples $K=1$, $\widehat{\operatorname{ELBO}} = \ln \hat p$. Therefore, we need at least two Monte Carlo samples to distinguish them. **And using larger $K$ could lead to greater improvement to VIS than VI. This doesn't mean VIS requires larger $K$ to work or VIS takes more time (same $K$, then same sampling time for all methods including VIS and VI), so no extra training time.** In practice, for example, if two samples are enough for VI to converge, then two samples are also enough for VIS to converge, and VIS should still outperform others. I.e., the number of Monte Carlo should be suitable to the complexity of the model/problem, rather than which method you choose. In our experiments, we use the same number of Monte Carlo and enough Monte Carlo samples for a fair and insightful comparison of different methods, especially the learning behavior of different methods, for example in Sec 4.1 the toy mixture model (Fig. 2). **New experiment results regarding the number of Monte Carlo are shown in Appendix 6 and Fig. 8 in our latest revision.**

Reviewer comment (Theory):
Instability in training due to the biased approximation (due to Eq. 8 and 11).

Author response:
* **Approximation of Eq. 8 and 11**:
   * **Eq. 8**: Eq. 8 is deriving the bias of the numerical estimator of IS $\ln \hat p(\boldsymbol x;\theta,\phi)$, showing that its distance to the true $\ln \hat p(\boldsymbol x;\theta)$ is asymptotically much closer than $\widehat{\operatorname{ELBO}}(\boldsymbol x;\theta,\phi)$. This tells us one of the reasons why we should prefer to directly maximize $\ln \hat p(\boldsymbol x;\theta,\phi)$ rather than $\widehat{\operatorname{ELBO}}(\boldsymbol x;\theta,\phi)$. The $\approx$ in Eq. 8 is not related to the "biased approximation" since it is not an estimator.
   * **Eq. 11**: It is a derivation of the numerically stable estimator of the forward $\chi^2$ divergence in log space. This quantifies the bias of the importance sampling (IS) estimator in log space and the effectiveness of the IS estimator. Hence, we minimize this to improve IS. It is likely that this approximation is biased, but the only way to approximate Eq. 10 $\ln V(\boldsymbol x|\theta,\phi)$ is to use such an approximation with samples $\boldsymbol z^{(k)}\sim q(\boldsymbol z|\boldsymbol x;\phi)$.

Reviewer comment (Theory):
The ground truth model of the toy experiment (Sec. 4.1) seems unidentifiable. In particular, the distribution of the hidden variable has multiple modes, which seems impossible to be identified only with the binary visible variable.

Author response:
* **Identifiable**: The purpose of the toy example is to provide insights into the distinctions between VIS and other methods. Therefore, we deliberately choose a restricted approximating posterior distribution family $q$ parameterized by $\phi$. This deliberate choice allows us to observe and comprehend the diverse behaviors of different methods, facilitating a comparative analysis to determine their efficacy through visualization.

Reviewer comment (Evaluation):
The observation of unidentifiability explains why, despite the substantial failure of conventional VI to estimate the modes, there is only a tiny difference in log-likelihood compared to other methods (order of 1e-4, see Fig. 2a).

Author response:
* **Log-likelihood**: The magnitude of the log-likelihood is highly contingent on the model itself $p(\boldsymbol x;\theta) = \int p(\boldsymbol x,\boldsymbol z;\theta)\ \mathrm{d} \boldsymbol z$. The error bars presented in Fig. 2a statistically validate that VIS is significantly better than others for the task of learning the model parameter set $\theta$.

Reviewer comment (Theory):
Why do you compare the parameter convergence, HLL, and CLL, which I think are meaningful to compare only if the model is identifiable? A similar argument also applies to the synthetic experiment of spike NN.

Author response:
* **Identifiable of latent does not affect us to compare parameter recovery**: It is not critical to require the latent to be identifiable. Unidentifiable approximating posterior distribution $q$ can still serve as variational/proposal distribution for learning $\theta$. In most cases, we don't know the true posterior distribution family, and we are not trying to find a better approximating posterior distribution family for a particular model. The purpose of VIS is to find the optimal proposal distribution to do IS and to learn $\theta$ under any arbitrarily given approximating posterior distribution family. So, we can still compare the learning results of different methods.
* **Parameter convergence**: Our aim is to utilize parameter convergence as a reliable confirmation that improved test log-likelihood aligns consistently with better parameter estimation.
* **HLL and CLL serve as two supplementary metrics**, aiding our understanding of why the inferred latent is better or worse, regardless of whether the latent is identifiable. Similar to LL, for example, **HLL informs us about the performances of various approximating posterior distributions $q(\boldsymbol z|\boldsymbol x;\phi)$, no matter the approximating posterior distribution family $q(\boldsymbol z|\boldsymbol x;\phi)$ is able to match the true posterior distribution family $p(\boldsymbol z|\boldsymbol x;\theta)$ or not**.
* In a real-world experiment, we will never know what is the distribution family of the generative model and the posterior distribution, but we still use marginal/hidden log-likelihood to compare the performance of different methods.

Reviewer comment (Novelty):
I am not entirely sure how the method proposed in this paper differs from previous work, as the practice of using Monte Carlo samples to sharpen variational bound is a topic that has been explored by previous work. I hope the authors can make some clarifications or have a small section in the manuscript that discusses the difference across different ways to combine IS with VI.

Author response:
In the third paragraph of Section 1, various methods of combining IS with VI are briefly mentioned. Additionally, these methods are summarized as baseline techniques at the beginning of Sec 4. Below, we aim to provide a more detailed explanation. In the context of a latent variable model $p(\boldsymbol x, \boldsymbol z;\theta)$, two main challenges arise: the learning problem, which focuses on learning $\theta$, and the inference problem, which aims to find an approximating posterior distribution $q(\boldsymbol z|\boldsymbol x;\phi)$ that approximates $p(\boldsymbol z|\boldsymbol x;\theta)$ for a specific $\theta$.
* Learning $\theta$: The process of learning $\theta$ involves maximizing the log-likelihood, denoted as $\arg\max_\theta \ln p(\boldsymbol x;\theta)$ in log space. Two distinct optimization approaches exist for this task:
   * (a) Maximizing $\widehat{\operatorname{ELBO}}(\boldsymbol x;\theta,\phi)$ w.r.t. to $\theta$, which is a hard lower-bound of $\ln p(\boldsymbol x;\theta)$, measured by $\operatorname{KL}(q\\|p)$.
   * (b) Maximizing $\ln \hat p(\boldsymbol x;\theta,\phi)$ w.r.t. $\theta$, which is an asymptotically tighter lower-bound of $\ln p(\boldsymbol x;\theta)$, measured by $\chi^2(p\\|q)/(2K)$.  

Although both of them are sampling-based methods, the former one is usually called VI, and the latter one is usually called IS. Both of them requires sampling $K$ samples $\left\\{\boldsymbol z^{(k)}\right\\}_{k=1}^K$ from the current approximating posterior distribution $q(\boldsymbol z|\boldsymbol x;\phi)$.
* Learning $\phi$: The ways to learn $\phi$ under a fixed $\theta$ are various. For example:
   * (1) Minimizing the reverse KL divergence $\operatorname{KL}(q(\boldsymbol z|\boldsymbol x;\phi)\\|p(\boldsymbol z|\boldsymbol x;\phi))$ which happen to also maximizing $\widehat{\operatorname{ELBO}}(\boldsymbol x;\theta,\phi)$ but w.r.t. $\phi$.
   * (2) Minimizing the forward $\chi^2$ divergence $\chi^2(p(\boldsymbol z|\boldsymbol x;\theta)\\|q(\boldsymbol z|\boldsymbol x;\phi))$
   * (3) $\alpha$ divergence-based methods or forward KL divergence

Subsequently, we can comprehend the various ways of combining them:
* VI: (a)*(1). Since they utilize the same target function $\widehat{\operatorname{ELBO}}$.
* CHIVI: (a)*((1) + (2)). Since they assume squeezing the posterior by a lower-bound ELBO and an upper-bound CUBO could improve the posterior approximating, but this may not be the best approximating posterior for learning $\theta$.
* VBIS: (b)*(1). They assume the variational distribution found by VI (1) could serve as a proposal distribution for learning $\theta$ (a).

**For your questions regarding $K$**, please see the following **Questions** part.

Reviewer comment (Theory):
It is neat to see the same chi-squared divergence on both the bias and effectiveness of the IS estimator, but eq. 8 contains an approximation that seems to work mostly for large $K$s, and the effectiveness is evaluated at the estimator $\hat{p}$ and not $\log \hat{p}$, making the connection seem a bit artificial. Is it possible to minimize a divergence in the form of, e.g., the first line in eq. 8?

Author response:
**Regarding your inquiry about Eq. 8**, it's challenging to analytically derive the effectiveness (variance) of the estimator $\ln \hat p(\boldsymbol x;\theta,\phi)$ in log space. However, empirical evidence from Fig. 1 demonstrates that reducing the variance in the original space also leads to a reduction in variance in the log space.

**Minimizing the first line in Eq. 8**: We argue that this is not practical, because the first line of Eq. 8 includes samples. Eq. 8 is deriving an estimator's bias. Bias is one of the metrics to evaluate a statistic, so its expression cannot include concrete samples. Mathematically we have to eliminate concrete samples $\boldsymbol z^{(k)}$ from the expression to obtain a sample-free analytical form of the bias, i.e. the last line of Eq. 8.

Reviewer comment (Experiments):
The experiments presented in the paper show that VIS performs well, but it seems that a large number of Monte Carlo samples are chosen in many of the experiments. This seems that it can present a significant computational cost. Could the authors include some explanation about the effect of the number of Monte Carlo samples on the efficacy of VIS, and clarify if taking many Monte Carlo samples is practical in training latent variable models?

Author response:
When the problem itself is complicated (e.g., complicated generative process and/or high latent dimensionality such as the POGLM model in Sec 4.3), a large number of Monte Carlo samples is indeed necessary, no matter for VIS or other baseline methods including VI. When working with VAE, people are already accustomed to using one Monte Carlo sample, but more Monte Carlo samples certainly improves the quality of $\widehat{\operatorname{ELBO}}$ since $\operatorname{ELBO}$ is also an integral expression w.r.t. $q(\boldsymbol z|\boldsymbol x;\phi)$. Anyway, however, employing a larger number of Monte Carlo samples always improves the estimation for all sampling-based estimation methods, including VI as depicted in Fig. 1. The main conclusion from Fig. 1 is that increasing the number of Monte Carlo increases the effectiveness of both $\widehat{\operatorname{ELBO}}$ and $\ln \hat p$, but it is only very helpful for tightening the bound of $\ln \hat p$ (with speed of order $\frac 1 K$) rather than $\widehat{\operatorname{ELBO}}$. As we mentioned below Eq. 8, when the number of Monte Carlo samples $K=1$, $\widehat{\operatorname{ELBO}} = \ln \hat p$. Therefore, we need at least two Monte Carlo samples to distinguish them. **And using larger $K$ could lead to greater improvement to VIS than VI. This doesn't mean VIS requires larger $K$ to work or VIS takes more time (same $K$, then same sampling time for all methods including VIS and VI), so no extra training time.** In practice, for example, if two samples are enough for VI to converge, then two samples are also enough for VIS to converge, and VIS should still outperform others. I.e., the number of Monte Carlo should be suitable to the complexity of the model/problem, rather than which method you choose. In our experiments, we use the same number of Monte Carlo and enough Monte Carlo samples for a fair and insightful comparison of different methods, especially the learning behavior of different methods, for example in Sec 4.1 the toy mixture model (Fig. 2). **New experiment results regarding the number of Monte Carlo are in Appendix 6 and Fig. 8 in our latest revision.**

Reviewer comment (Novelty):
Using such adaptive-importance-sampling approaches for variational inference is already extensively discussed in [3].

Author response:
* [3] Finke et al. (2019): While this paper addresses the gradient breakdown problem in the IWAE (Burda et al., 2016) paper and by providing a correct gradient estimator for optimizing the encoder or the proposal distribution, it's still important to highlight differences. Firstly, their gradient estimator is in the original space (Dieng et al., 2017), unlike ours, which is in log space. This is nontrivial and an important way to make the gradient estimator numerically stable. **Please check Appendix 5 and Fig 7 in our latest revision for insights into the numerical stability issue.** Besides, their experiments are confined to cases where the generative model has an explicit decomposition of $p(\boldsymbol x, \boldsymbol z;\theta) = p(\boldsymbol x|\boldsymbol z;\theta) p(\boldsymbol z;\theta)$. In contrast, our POGLM model in Sec 4.3 extends the analysis to a complex and general graphical model where such an explicit decomposition does not exist. Even for such a complicated general graphical model, our VIS can still work and outperform other baselines, which proves the practical effectiveness of our method.

Furthermore, to the best of our knowledge, there is a gap in the existing literature concerning extensive experiments that systematically compare the practical learning performances of various methods across a diverse range of latent variable models. We believe it is necessary to visualize the behaviors of different methods (e.g. the toy model in Section 4.1 and the POGLM model in Sec 4.3) for readers to fully understand and convince how VIS works better than others.

Reviewer comment (Novelty):
Algorithm 1 of the present paper is a special case of the generic method described in [3, Section 3] (in particular, see [3, Section 3.4]).

Author response:
* [3] Finke et al. (2019): While this paper addresses the gradient breakdown problem in the IWAE (Burda et al., 2016) paper and by providing a correct gradient estimator for optimizing the encoder or the proposal distribution, it's still important to highlight differences. Firstly, their gradient estimator is in the original space (Dieng et al., 2017), unlike ours, which is in log space. This is nontrivial and an important way to make the gradient estimator numerically stable. **Please check Appendix 5 and Fig 7 in our latest revision for insights into the numerical stability issue.** Besides, their experiments are confined to cases where the generative model has an explicit decomposition of $p(\boldsymbol x, \boldsymbol z;\theta) = p(\boldsymbol x|\boldsymbol z;\theta) p(\boldsymbol z;\theta)$. In contrast, our POGLM model in Sec 4.3 extends the analysis to a complex and general graphical model where such an explicit decomposition does not exist. Even for such a complicated general graphical model, our VIS can still work and outperform other baselines, which proves the practical effectiveness of our method.

Reviewer comment (Novelty):
The $\theta$-gradient from Equation 6 of the present work is already well known (see, e.g. [3]).

Author response:
* [3] Finke et al. (2019): While this paper addresses the gradient breakdown problem in the IWAE (Burda et al., 2016) paper and by providing a correct gradient estimator for optimizing the encoder or the proposal distribution, it's still important to highlight differences. Firstly, their gradient estimator is in the original space (Dieng et al., 2017), unlike ours, which is in log space. This is nontrivial and an important way to make the gradient estimator numerically stable. **Please check Appendix 5 and Fig 7 in our latest revision for insights into the numerical stability issue.** Besides, their experiments are confined to cases where the generative model has an explicit decomposition of $p(\boldsymbol x, \boldsymbol z;\theta) = p(\boldsymbol x|\boldsymbol z;\theta) p(\boldsymbol z;\theta)$. In contrast, our POGLM model in Sec 4.3 extends the analysis to a complex and general graphical model where such an explicit decomposition does not exist. Even for such a complicated general graphical model, our VIS can still work and outperform other baselines, which proves the practical effectiveness of our method.

Reviewer comment (Theory):
From the author's rebuttal, it is now more clear that their $\phi$ gradient is slightly different than in [3] because they derive the $\phi$-gradient in log-space.

Author response:
* **Regarding [2]**.
    * Eq. 10 in our paper shows that $\chi^2(p\\|q) = \frac{1}{p(\boldsymbol x;\theta)^2} V(\boldsymbol x;\theta,\phi) - 1$. So, our target function $V(\boldsymbol x;\theta;\phi) = p(\boldsymbol x;\theta)^2\chi^2(p\\|q) + p(\boldsymbol x;\theta)^2$ in original space is the same as existing works.
    * Theoretically, $p(\boldsymbol x;\theta)^2$ is a constant w.r.t. $\phi$. So, we (and existing works) drop the term of $p(\boldsymbol x;\theta)^2$, to only minimize $V(\boldsymbol x;\theta,\phi) = \int \frac{p(\boldsymbol x,\boldsymbol z;\theta)^2}{q(\boldsymbol z|\boldsymbol x;\phi)}\ \mathrm d \boldsymbol z$. We acknowledge that this step is helpful for numerical stability because the target function $V(\boldsymbol x;\theta,\phi)$ will not be affected by the magnitude of $p(\boldsymbol x;\theta)^2$.
    * We take an additional step to perform the minimization in log space, which further improves the numerical stability significantly.

Reviewer comment (Writing):
The paper is not polished yet; although the major parts are all there, it may require another thorough pass. It has too many mistakes and typos: the notation changes from bold to normal in many places, the title has a typo ('importane'), and there is a mistaken statement 'log function is a convex function'.

Author response:
Thank you for bringing attention to the typo in the title. We apologize for any confusion caused by these errors. Please check the polished latest revision.  
   **Notation**: In the method section, we employ vector forms $\boldsymbol x, \boldsymbol z$ for generality. However, in the experiments, we adapt them to the corresponding typefaces based on the model. For instance, in Sec 4.1 Toy example, both $x$ and $z$ are scalars, so we use the normal typeface. In Sec 4.3 POGLM, every observable $\boldsymbol X \in \mathbb N^{\text{time}\times \text{visible neurons}}$ and every latent $\boldsymbol Z \in \mathbb N^{\text{time}\times \text{hidden neurons}}$ are multivariate time sequences represented in matrix forms, so we use the capital bold typeface.  
   **log**: Thanks! log is a concave function.

Reviewer comment (Novelty):
Some of the references and recent literature are missing, especially those that have looked at the quality of different divergence objectives such as CUBO and ELBO for finding the optimal sampling distribution.

Author response:
**The papers related to CUBO and ELBO had been included in our initial submission. The second baseline method CHIVI (Dieng et al., 2017), explained at the beginning of Sec 4, is the CUBO-ELBO paper.** Specifically, they utilize CUBO as an upper-bound and ELBO as a lower-bound to squeeze the approximated posterior. However, all discussions in the CUBO paper revolve around the inference problem, i.e., approximating the posterior distribution $q(\boldsymbol z|\boldsymbol x;\phi)$, rather than determining the optimal proposal distribution for learning the parameter $\theta$ in $p(\boldsymbol x;\theta)$ as we were doing.

Reviewer comment (Theory):
It is the bane of chi-square divergence methods that they do not scale well with dimensions, as covered in the papers https://arxiv.org/pdf/2010.09541.pdf and https://arxiv.org/abs/1802.02538, and it seems that this method may not scale well either since it uses chi-squared divergence minimization.

Author response:
**Scale to work for high dimensional latent space: This stands as one of our significant contributions.** The two papers you mentioned, along with Pradier et al. (2019) cited in our paper, acknowledge the instability associated with minimizing $\chi^2$ divergence, particularly in high-dimensional scenarios. **The lack of success in minimizing the $\chi^2$ divergence could be attributed to the numerically unstable form of their gradient estimator. In contrast, our gradient estimator (Equations 12 and 24) is derived in log space and presented in a simple and numerically stable form. Please check Appendix 5 and Fig 7 in our latest revision.** Our succinct gradient estimator makes VIS always work, even for the most general case that $p(\boldsymbol x,\boldsymbol z;\theta)$ cannot be explicitly decomposed into $p(\boldsymbol x|\boldsymbol z;\theta) p(\boldsymbol z;\theta)$. The POGLM experiments indeed confirm that VIS is robust across different models, numerically stable, and scalable to high latent dimensionality (addressing Question 1). Having successfully overcome the challenges in previous research related to $\chi^2$ divergence, we consider this a strength and a noteworthy contribution to our work, rather than a weakness.

Reviewer comment (Experiments):
What is the dimensionality of the POGLM model? Do the authors intend to use this method as a tool for low-dimensional complex posteriors, given that both IS and CUBO do not scale well with dimensions, and even a large sample size as done in this paper will not help?

Author response:
* **The dimensionality of the latent space in POGLM** is the number of hidden neurons $\times$ the number of time bins, so it is a **high dimensional** problem.
   * **Scaling issue: Anyway, employing a larger number of Monte Carlo samples always improves the estimation for all sampling-based estimation methods, including VI as depicted in Fig. 1.** But when the problem itself is complicated, a large number of Monte Carlo samples is indeed necessary, no matter for VIS or other baseline methods including VI (ELBO). **Please check Appendix 6 and Fig 8 in our latest revision.**
   * Our main goal is learning $\theta$ to achieve high marginal log-likelihood. We are not studying dimensionality reduction. And we are not trying to find the best variational/proposal distribution family for a particular model. The POGLM experiment is to show that for such a complicated generative model, VIS obtains the most optimal approximated posterior, which serves for recovering the model parameter $\theta$ the best.

Reviewer comment (Experiments):
Did the authors use any other optimizers other than ADAM? Did it have any effect? How did you choose the optimization algorithm hyperparameters like learning rate, etc.?

Author response:
* **Optimizer**: We opt for ADAM as our choice of optimizer primarily because it has become the default preference across the board in current practices. While the selection of the optimizer falls outside the specific focus of this paper, we did experiment with Adagrad, and the comparative results did not exhibit significant changes. Notably, people commonly employ ADAM, Adagrad, or other optimizers to train VAE using VI, and the same applies to training VAE using our VIS. In the interest of a fair and effective comparison, we adhere to the widely used ADAM optimizer.
   * **Learning rate**: To determine the learning rate, we employ cross-validation, aligning with standard practices in many machine learning tasks. It's crucial to emphasize that the learning rate is a hyperparameter associated with the optimizer rather than being intrinsic to the method or the model. For the sake of a fair comparison, we set the learning rate to be consistent and, more importantly, suitable for the specific model and problem under consideration.