Forward $\chi^2$ Divergence Based Variational Importance Sampling

Dear Reviewer T2mb,

Thank you very much for your time and valuable comments on our paper. We highly appreciate your recognition of the strengths of our paper. Hopefully, the following responses could resolve most of your concerns and answer your questions.

Weaknesses

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.

• 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.
• 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)$.

Dear Reviewer QxSX,

Thank you very much for your time and valuable comments on our paper. We highly appreciate your recognition of the strengths of our paper. Hopefully, the following responses could resolve most of your concerns and answer your questions.

Weaknesses

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.

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.

Dear Reviewer 84Pi,

Thank you very much for your time and valuable comments on our paper. We highly appreciate your recognition of the strengths of our paper. Hopefully, the following responses could resolve most of your concerns and answer your questions.

Weaknesses:

1. 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.

2. 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.

3. We apologize for the compressed and dense nature of the method section (Sec 3). Due to the 9-page limit, we acknowledge that crucial materials, such as related works, mathematical derivations (e.g., the gradient estimator), and detailed experiment explanations, may not have been extensively expanded in the main content. Most of the important lengthy derivations are provided in the appendix, allowing readers to verify the correctness of the result equations presented in the main content. Regarding Fig. 2, we appreciate the opportunity to resolve any confusion and are prepared to carefully explain all the details to you now.
   Fig 2: To comprehensively understand why VIS outperforms other methods, it's crucial to visualize three distributions:

   • the true posterior $p(z|x;\theta^{\mathrm{true}})$ (dashed)
   • the learned posterior $p(z|x;\theta)$ (solid)
   • the approximated (variational/proposal) posterior $q(z|x;\phi)$ (dotted).

   This linestyle assignment aligns with the convergence curves shown in Fig. 2B. Given $x\in\\{0,1\\}$, we use two colors: purple for $z|x=0$ and brown for $z|x=1$. We can ignore brown curves first and only focus on purple curves to understand why VIS is better than others.

   • Conclusion 1: the dotted curve of VIS is the widest, providing the widest effective sampling range, which is beneficial for learning $\theta$ through IS.
   • Conclusion 2: the dotted curve (approximated posterior) and the solid curve (learned posterior) of VI exhibit the closest match, resulting in the smallest reverse KL divergence. However, the solid curve (learned posterior) is significantly distant from the dashed curve (true posterior). This suggests that maximizing ELBO in VI for learning $\theta$ could lead to a situation where the learned and approximated posteriors are close in reverse KL divergence (yielding a high ELBO), but both are far from the true posterior (resulting in a low marginal likelihood). This significantly impedes the learning of $\theta$.

4. 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.

Dear Reviewer J8Yz,

Thank you very much for your time and valuable comments on our paper. We would like to make clarifications on your valuable suggestions.

• [1] Jona-Lasinio et al. (1999): We acknowledge the presence of existing works that employ adaptive strategies to minimize the forward $\chi^2$ divergence, aiming to reduce the variance of the IS estimator. However, this paper only provides a general guideline but not a concrete algorithm for achieving the goal of minimizing the forward $\chi^2$ divergence. Additionally, our paper introduces another perspective on minimizing the forward $\chi^2$ divergence by considering the bias of $\ln \hat p(\boldsymbol x;\theta,\phi)$ in log space.
• [2] Akyildiz et al. (2021): This paper primarily discusses minimizing $\chi^2$ divergence for distributions from the exponential distribution family. In contrast, our gradient estimator is not restricted to such a distribution family; it is applicable to all distributions, as demonstrated in the toy model (Sec 4.1) and the POGLM model (Sec 4.3). The true posterior distribution family and the approximated distribution family in these two models are way more complicated than the exponential distribution family, but we can still make VIS work for these arbitrary distribution families.
• [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.

In conclusion, we respectfully hold a different perspective, asserting that our work indeed brings novelty and valuable contributions. We hope that our response has addressed your concerns and provided clarity regarding the significance of our research.

Thank you for the clarification.

1. I agree that it makes sense to minimise the chi-square divergence (or the following quantity that is proportional to it for fixed $\theta$) $E_{z \sim q_\phi(z|x)}\Bigl[\Bigl(\frac{p_\theta(x, z)}{q_\phi(z|x)}\Bigr)^2\Bigr]$ in log-space. And you're right, this makes your $\phi$-gradient slightly different from some existing adaptive-importance sampling settings (the $\theta$-gradient is standard). I will raise my score.
2. Yes, [2] assumes exponential families. I don't think this specific reference needs to necessarily be cited in the paper. Rather, along with [1], I only gave it as an example of the fact that minimising the variance of the importance weights (equivalently, minimising the above-mentioned chi-square divergence) is a well-known technique in the importance-sampling literature.
3. This context needs to be made more clear in the paper, e.g., in the introduction.

There are still some typos in the revised manuscript. I haven't re-read everything carefully but noticed that Eq. 12 should probably be "$\approx$" rather than "$=$"

Dear Reviewer J8Yz

We would like to express our sincere thanks for your acknowledgment of our paper and your new score.

• 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.
• Eq. 12. Thanks, it is $\approx$. We will certainly do further proofreading to make sure there is no typo.
• Regarding [3]. Thanks for this suggestion. We have added some new discussions to the introduction part. Due to the 9-page limit, some detailed texts are located in Appendix 8 (with a hyperlink directing from the introduction). Please check the latest revision.
• Comparing with forward KL. Thanks for this suggestion. We do have the result in our hands. For example, Jerfel et al. (2021) also found the problem of the reverse KL divergence and investigated the forward KL divergence, but their algorithm is only confined to Gaussian distribution. According to your suggestion, we had a new section in Appendix A.7 showing that forward $\chi^2$ divergence is better than forward KL divergence. The reason should be from Sason et al. (2016) and Nishiyama et al. (2020) that: $\operatorname{KL}(p\\|q) \leqslant \chi^2(p\\|q)$. I.e., $\chi^2(p\\|q)$ can bound $\operatorname{KL}(p\\|q)$ but not vice versa. Therefore, minimizing forward KL divergence cannot promise a direct minimization of forward $\chi^2$ divergence.

Please check Fig. 9 in Appendix 7 to see the new results. Thank you very much again for your valuable feedback on our initial response.

New references:
Sason, I., & Verdú, S. (2016). $f$-divergence Inequalities. IEEE Transactions on Information Theory, 62(11), 5973-6006.
Nishiyama, T., & Sason, I. (2020). On relations between the relative entropy and $χ^2$-divergence, generalizations and applications. Entropy, 22(5), 563.