Technical Summary: Learning with Importance Weighted Variational Inference

Problem Statement

Variational Inference (VI) seeks to approximate an intractable target posterior density $p_\theta(z|x)$ using a tractable family of distributions $q_\phi(z|x)$. While the standard Evidence Lower BOund (ELBO) is the traditional objective, recent research has focused on importance-weighted bounds, such as the Importance-weighted Auto-Encoder (IWAE), Variational Rényi (VR), and the unifying VR-IWAE bound. These bounds depend on the number of Monte Carlo samples $N$ and a hyperparameter $\alpha \in [0, 1)$.

Despite the popularity of these bounds, the theoretical understanding of how the joint choice of the bound and the gradient estimator impacts the optimization trajectory remains limited. Specifically, there is a lack of comprehensive analysis regarding:

1. The asymptotic behavior of Reparameterized (REP) and Doubly-Reparameterized (DREP) gradient estimators as $N \to \infty$.
2. The bias-variance tradeoff introduced by the hyperparameter $\alpha$ in these estimators.
3. The performance of these estimators in "challenging regimes" where the variational approximation deteriorates (i.e., the Kullback-Leibler divergence between the variational and posterior densities is large).

Current literature offers fragmented insights, often restricted to specific cases (e.g., $\alpha=0$ for IWAE) or relying on informal arguments. Empirical studies suggest DREP outperforms REP and that tuning $\alpha$ improves performance, but a unified theoretical justification is missing.

Methodology

The authors develop a unified theoretical framework to analyze the Signal-to-Noise Ratio (SNR) of gradient estimators for the VR-IWAE bound. The methodology relies on two primary asymptotic analyses:

1. Asymptotics as $N \to \infty$ (Fixed Mismatch)

The paper first establishes a general theoretical result for the asymptotic expectation and variance of ratios of sample means.

• Theorem 1: Provides the limiting behavior of $E[X_N / (W_N)^\mu]$ and its variance as $N \to \infty$, under mild moment conditions. This generalizes existing results for self-normalized importance sampling.
• Application to VI: By mapping the REP and DREP gradient estimators to this ratio-of-means framework (where $X$ and $W$ are functions of the importance weights and their derivatives), the authors derive the asymptotic SNR for both estimators.
• Key Metric: The SNR is defined as $|E[\text{estimator}]| / \sqrt{\text{Var}[\text{estimator}]}$. A high SNR indicates a reliable gradient direction.

2. Asymptotics in Deteriorating Regimes ($N, \text{KL} \to \infty$)

The paper analyzes a regime where the variational approximation is poor, characterized by the KL divergence $KL(\theta, \phi; x) \to \infty$, while $N \to \infty$ such that $\log N = o(KL)$.

• Assumption (B1): The log-importance weights are modeled as a Gaussian random field. This allows the authors to link the KL divergence directly to the variance of the log-weights.
• Analysis: The authors investigate whether increasing $N$ beyond 1 provides any benefit when the variational density is far from the posterior. They derive conditions under which the gradient estimators revert to the behavior of the $N=1$ (ELBO) case.

Key Contributions

1. Unified Comparison of REP and DREP Estimators

The paper provides the first unified asymptotic characterization of REP and DREP gradient estimators for the VR-IWAE bound across the entire range $\alpha \in [0, 1)$.

• Bias-Variance Tradeoff: The analysis reveals that the hyperparameter $\alpha$ governs a tradeoff. In the REP estimator, decreasing $\alpha$ (moving toward IWAE) improves the bias (closer to the marginal log-likelihood) but increases the variance.
• Theoretical Justification for DREP: The authors formally prove that the DREP estimator possesses a "zero-variance" property at optimality (when $q_\phi$ matches $p_\theta$). Specifically, the asymptotic variance term for DREP vanishes when the variational approximation is perfect, whereas the REP variance remains non-zero due to the score function term. This provides the first rigorous theoretical justification for the empirically observed superiority of DREP.

2. Analysis of Challenging Regimes

The paper characterizes the behavior of importance-weighted VI when the variational approximation deteriorates.

• Weight Collapse: The authors show that if the KL divergence is large, the sample size $N$ required to achieve the benefits of importance weighting grows exponentially ($N \sim e^{KL}$).
• Reversion to ELBO: If $N$ does not grow exponentially with the KL divergence, the gradient estimators effectively revert to the $N=1$ ELBO baseline. Crucially, the paper demonstrates that even in this regime, the gradient estimators point in a valid direction that minimizes the KL divergence, thereby improving the variational approximation and eventually allowing the algorithm to enter the "well-behaved" regime where $N > 1$ becomes effective.

3. General Monte Carlo Tools

Beyond VI, the paper establishes general theoretical tools for analyzing sample mean ratios (Theorem 1). These tools operate under strictly weaker assumptions than previous works (e.g., Rainforth et al., 2018; Daudel et al., 2023) and accommodate general powers $\mu > 0$, extending the applicability of importance-weighting theory.

Results

Asymptotic Behavior ($N \to \infty$)

• REP Estimator: The SNR scales as $\sqrt{N}$ for $\theta$ parameters. For $\phi$ parameters, the SNR scales as $\sqrt{N}$ only if $\alpha \in (0, 1)$. If $\alpha = 0$ (IWAE), the SNR for $\phi$ decays as $1/\sqrt{N}$ due to high variance.
• DREP Estimator: The SNR scales as $\sqrt{N}$ for both $\theta$ and $\phi$ parameters across $\alpha \in [0, 1)$.
• Variance Reduction: The asymptotic variance of the DREP estimator is strictly lower than that of the REP estimator. In Example 2 (Linear Gaussian), the DREP variance is reduced by a factor of $\alpha^2/16$ compared to REP for $\alpha \in (0, 1)$.

Deteriorating Regime Results

• Threshold for $N$: The paper establishes that $N$ must satisfy $\log N \ll KL(\theta, \phi; x)$ to avoid weight collapse.
• Optimization Trajectory: In the high-mismatch regime, increasing $N$ beyond 1 does not significantly alter the optimization direction compared to $N=1$. The gradient estimator behaves as if it were optimizing the ELBO, which is beneficial because it drives the reduction of the KL divergence.
• Two-Phase Training: The results suggest a two-phase training dynamic:
  1. Initial Phase (High Mismatch): $N=1$ is optimal to maximize the number of gradient steps and reduce KL divergence rapidly.
  2. Final Phase (Low Mismatch): Once the variational approximation is sufficiently close to the posterior, increasing $N$ reduces variance and improves convergence.

Significance and Claims

The paper claims to advance the understanding of importance-weighted VI by:

1. Formalizing the DREP Advantage: Providing the first rigorous proof that DREP estimators outperform REP estimators in importance-weighted VI, specifically by eliminating the variance contribution of the score function at optimality.
2. Clarifying the Role of $\alpha$: Demonstrating that $\alpha$ controls a bias-variance tradeoff, explaining why tuning $\alpha$ empirically improves performance.
3. Bridging Training Phases: Unifying the understanding of VI behavior from poor initialization (high KL) to convergence (low KL). The paper argues that the "weight collapse" phenomenon in high-mismatch regimes is not a failure but a mechanism that naturally reverts the algorithm to the robust ELBO baseline, ensuring progress until the variational family becomes expressive enough to support larger $N$.
4. Methodological Contribution: Offering a general framework for analyzing sample mean ratios that extends beyond VI to other Monte Carlo applications, such as the fine-tuning of Large Language Models.

The authors validate these theoretical findings through numerical experiments on Gaussian models, Linear Gaussian models, and a Variational Auto-Encoder (VAE) on MNIST, showing that the observed SNR trends align with the derived asymptotic predictions.