Control Variate Design for Score-Function Gradient Estimators in Importance-Weighted Variational Inference

The gist The authors provide a comprehensive analysis of REINFORCE gradient estimators in Importance Weighted Variational Inference, introducing a novel VIMCO-⋆ gradient estimator that resolves fundamental signal-to-noise ratio collapse issues observed in state-of-the-art methods.

Importance Weighted VI and Gradient Estimation

Variational inference (VI) recasts sampling as an optimization problem by maximizing a variational bound, such as the Evidence Lower Bound (ELBO) or importance weighted bounds like the VR-IWAE bound. A common method for optimizing these bounds is stochastic gradient ascent (SGA), which relies on unbiased gradient estimators derived from reparameterization trick assumptions, though REINFORCE estimators do not require these differentiability assumptions. The paper examines a generalized family of VIMCO gradient estimators for Monte Carlo objectives (VIMCO) gradients, which are state-of-the-art but lack theoretical justification.

Theoretical Analysis of REINFORCE Gradients

The authors establish rigorous analyses of the REINFORCE gradient of the VR-IWAE bound, showing that state-of-the-art VIMCO gradient estimators exhibit a vanishing signal-to-noise ratio (SNR) as the number of Monte Carlo samples N increases, preventing effective optimization. They prove that without proper variance reduction techniques, REINFORCE gradient estimators of the VR-IWAE bound are unusable in practice as N increases. Theorem 5 introduces the novel VIMCO-⋆ gradient estimator, which averts this SNR collapse by achieving a √N SNR scaling when α = 0.

Novel Estimator Design and Performance

The paper proposes the VIMCO-⋆ gradient estimator, denoted as ψ-g(α,VIMCO⋆) N (θ, ϕ; x), which is defined by a specific control variate function f(α−i). This estimator is shown to enjoy the best asymptotic efficiency by design compared to the VIMCO-AM and VIMCO-GM ones when α ∈ (0, 1). The analysis demonstrates that for all α ∈ (0, 1), SNR scales as √N for the VIMCO-⋆ gradient estimator, while the other estimators scale as O(√N) or slower when α = 0.

Practical Considerations and Annealing

The choice of hyperparameters (N, α, f(α−i)) controls two distinct mechanisms: increasing N reduces the bias relative to the marginal log-likelihood and variance, while α governs a bias-variance trade-off. The VIMCO-⋆ gradient estimator is preferred because it has the best asymptotic variance and its nonasymptotic variance decays quicker than 1/N at optimality. Annealing strategies, such as annealing the likelihood, are presented as a practical solution to mitigate noise in early optimization stages and lower the computational barrier to entry.

Empirical Validation

Numerical experiments on a Gaussian example confirm the theoretical claims, showing that for α > 0, all three gradient estimators exhibit an √N scaling of the SNR, with VIMCO-⋆ consistently attaining the highest SNR. In Bayesian phylogenetic inference settings, VIMCO-⋆ outperforms VIMCO-AM and VIMCO-GM in terms of KL divergence from the ground truth posterior.

Conclusion

The paper concludes that the VIMCO-⋆ gradient estimator is the first REINFORCE gradient estimator within the importance weighted VI framework whose use is theoretically supported across all values of α ∈ [0, 1) in terms of both asymptotic efficiency as N → ∞ and non-asymptotic efficiency at optimality.

How it works

**: The authors introduce a generalized family of VIMCO gradient estimators for Monte Carlo objectives (VIMCO) gradients, parameterized by the number of Monte Carlo samples N, a hyperparameter α ∈ [0, 1), and a flexible control variate function f(α−i). They motivate this family through two theoretical results: Theorem 1 characterizes the asymptotic behavior as N → ∞, showing that unbiased estimators point asymptotically in a theoretically-grounded direction that outperforms standard ELBO-based optimization in terms of bias (the lower α, the better). Theorem 2 establishes that without proper variance-reduction techniques, REINFORCE gradient estimators of the VR-IWAE bound are unusable in practice as N increases. The key insight is that the success of REINFORCE gradient estimators is contingent on suitable variance-reduction mechanisms, warranting a deeper investigation into the choice of the control variate function f(α−i). 5. Importance Weighted VI without Reparameterizing pg3> The authors establish that the REINFORCE gradient of the VR-IWAE bound converges to ∂ψVR(α)(θ, ϕ; x) at a fast 1/N rate as N → ∞. They then derive unbiased estimators like ψ-g(α,NAIVE) N (θ, ϕ; x) and show that the NAIVE gradient estimator's SNR scales as 1/√N when ψ ∈ Σϕ1,..., Σϕb. 3.2 Estimating the Gradient of the VR-IWAE Bound pg9> The authors derive and study unbiased REINFORCE gradient estimators of the VR-IWAE bound, showing that the NAIVE gradient estimator (8) has a variance that scales linearly with N when ψ ∈ Σϕ1,..., Σϕb. This leads to the conclusion that the NAIVE gradient estimator is unusable in practice. 3.3 Practical Considerations pg3> The choice of (N, α, f(α−i)) controls two distinct mechanisms: increasing N reduces the bias of the VR-IWAE bound relative to the marginal log-likelihood and variance of its gradient estimator. 3.3.1 Selecting f(α−i) in the VIMCO gradient estimator pg3> The VIMCO-⋆ gradient estimator (19) is preferred because it has the best asymptotic SNR rate that scales as √N and its nonasymptotic variance decays quicker than 1/N with N at optimality. 3.3.2 Annealing schedule pg3> Annealing strategies, such as annealing the likelihood, are presented as a practical solution to mitigate noise in early optimization stages and lower the computational barrier to entry. 4 Numerical experiments pg4> The Gaussian example confirms that for α > 0, all three gradient estimators exhibit an √N scaling of the SNR, with VIMCO-⋆ consistently attaining the highest SNR. 5. Discussion pg5> The versatility of VIMCO gradient estimators opens several compelling research directions, including developing theoretically grounded annealing schedules analogous to those in Sequential Monte Carlo literature. 6. Technical results pg5> The paper provides an asymptotic characterization of VIMCO gradient estimators that complements established results under the reparameterization trick assumption, highlighting that the VIMCO-⋆ gradient estimator achieves an asymptotic SNR rate of O(√N) for the learning of ϕ, markedly outperforming the O(1/√N) rate of standard reparameterized IWAE gradient estimator. 7. Discussion pg5> Future research is expected to focus on the behavior of leading constants within these asymptotic rates. 8. Technical results pg5> The VIMCO-⋆ gradient estimator enjoys the best asymptotic efficiency by design compared to the VIMCO-AM and VIMCO-GM ones when α ∈ (0, 1). 9. Discussion pg5> The paper's results challenge the prevailing heuristic that reparameterized gradient estimators are inherently superior to REINFORCE-based alternatives.

1. Importance weighted variational inference (VI) systems can utilize REINFORCE gradient estimators, which do not suffer from restrictions on data-generating processes or variational families, to optimize bounds that tighten with more Monte Carlo samples N. This allows for the application of these gradients in challenging settings where reparameterized gradients are unavailable, such as likelihood-free inference and state-space models.

2. The VIMCO star gradient estimator can be employed to achieve a signal-to-noise ratio scaling of 1 over sqrt(N) when alpha is between 0 and 1, which is superior to the O(1/sqrt(N)) rate of standard reparameterized IWAE gradient estimators. This provides more effective optimization directions for variational parameters phi.

3. For importance weighted VI on state-space models, the VIMCO star estimator with alpha equals 0 achieves a variance that converges to zero at a rate of 1 over N cubed, which is significantly faster than the variance scaling of existing estimators when alpha equals 0. This makes it suitable for learning parameters in complex time series models.

4. The VIMCO star gradient estimator can be used in Bayesian phylogenetic inference to achieve the most accurate variational approximation, as evidenced by its superior performance on benchmark data compared to VIMCO-AM and VIMCO-GM when alpha equals 0. This indicates it is the preferred estimator for learning tree topology parameters.

5. A dynamic annealing schedule for alpha can be implemented in VI systems by adaptively reducing alpha based on the effective sample size ESS1/alpha, which helps stabilize weights and increase the SNR of the VIMCO gradient estimators, accelerating training towards a tight variational bound.

Importance-weighted variational inference approximates intractable densities by optimizing multi-sample Monte Carlo objectives whose tightness increases with the number of importance samples N. When reparameterization is unavailable, optimization relies on score-function (REINFORCE) gradient estimators. In this paper, we provide a comprehensive analysis of score-function gradient estimators for importance-weighted objectives by characterizing their asymptotic expectation and variance. We establish that standard variational inference for Monte Carlo objectives (VIMCO) gradient estimators suffer from a vanishing signal-to-noise ratio (SNR) as N increases, which arises from a suboptimal control variate choice and can impair optimization. To address this issue, we derive a general mechanism for constructing VIMCO gradient estimators that avoid SNR collapse. We identify the High-Effective Sample Size optimal variance importance sampling estimator as belonging to this category of estimators and establish its sqrt N SNR scaling. This estimator does not, however, have zero variance at optimality and achieving its sqrt N SNR rate may require prohibitively large N when the importance weights are degenerate. We overcome these limitations by introducing VIMCO gradient estimators that attain zero variance at optimality and support importance weight tempering. Experiments with discrete latent-variable and state-space models validate the theoretical predictions and illustrate the benefits of our estimators.