Importance Weighted Variational Inference without the Reparameterization Trick

TL;DR

The paper tackles the challenge of performing importance-weighted variational inference when the reparameterization trick is unavailable by providing a rigorous analysis of REINFORCE-based gradients for the VR-IWAE bound. It introduces a generalized VIMCO framework with AM/GM variants and a novel VIMCO-star estimator that achieves favorable variance properties, including an $O(\sqrt{N})$ SNR for α in (0,1) and an $O(1/N^3)$ variance for α = 0. The authors prove asymptotic results and validate them through experiments on Gaussian models, state-space models, and Bayesian phylogenetics, showing superior optimization behavior and posterior accuracy. These results broaden the non-reparameterizable VI toolkit, enabling robust learning in discrete latent settings and likelihood-free contexts.

Abstract

Importance weighted variational inference (VI) approximates densities known up to a normalizing constant by optimizing bounds that tighten with the number of Monte Carlo samples $N$. Standard optimization relies on reparameterized gradient estimators, which are well-studied theoretically yet restrict both the choice of the data-generating process and the variational approximation. While REINFORCE gradient estimators do not suffer from such restrictions, they lack rigorous theoretical justification. In this paper, we provide the first comprehensive analysis of REINFORCE gradient estimators in importance weighted VI, leveraging this theoretical foundation to diagnose and resolve fundamental deficiencies in current state-of-the-art estimators. Specifically, we introduce and examine a generalized family of variational inference for Monte Carlo objectives (VIMCO) gradient estimators. We prove that state-of-the-art VIMCO gradient estimators exhibit a vanishing signal-to-noise ratio (SNR) as $N$ increases, which prevents effective optimization. To overcome this issue, we propose the novel VIMCO-$\star$ gradient estimator and show that it averts the SNR collapse of existing VIMCO gradient estimators by achieving a $\sqrt{N}$ SNR scaling instead. We demonstrate its superior empirical performance compared to current VIMCO implementations in challenging settings where reparameterized gradients are typically unavailable.

Figures (6)

• Figure 1: SNR as a function of $N$ for different values of $\alpha$ when $\theta=0, \;\phi=1$ in the Gaussian example. The solid lines correspond to the SNR estimates for different VIMCO gradient estimators, while the dashed lines correspond to the theoretical predictions.
• Figure 2: SNR as a function of $N$ for different values of $\alpha$ when $\theta=0, \;\phi=0.1$ in the Gaussian example. The solid lines correspond to the SNR estimates for different VIMCO gradient estimators, while the dashed lines correspond to the theoretical predictions.
• Figure 3: Performance of VIMCO gradient estimators when $\alpha=0$ in the Gaussian example. The left and middle plots: the absolute mean and standard deviation as functions of $N$. The right plot: the difference between $\phi$ and $\theta$ in absolute value as a function of the number of iterations. The dashed lines correspond to the theoretical predictions. The shaded area represents the 5% to 95% percentile range across 10 independent runs.
• Figure 4: Final marginal posterior estimates in the Variational Bayesian SSMs example.
• Figure 5: Performance of VIMCO-$\star$, VIMCO-AM and VIMCO-GM as the topological gradient estimator for VBPI on DS1.

Theorems & Definitions (28)

• Theorem 1
• Lemma 1
• Theorem 2
• Theorem 3
• Theorem 4
• Theorem 5
• Theorem 6
• Proposition 1
• proof
• Lemma 2