VIMCO-★ Gradient Estimator

• The paper introduces VIMCO-★, which employs an optimally tuned control variate to reverse the SNR collapse seen in classical VIMCO estimators, achieving O(√N) scaling.
• It details an asymptotic analysis that yields an estimator with variance scaling as O(1/N³) for α = 0 and minimizes the leading-order variance constant for α > 0.
• Empirical evaluations in Gaussian, state-space, and variational Bayesian phylogenetics models confirm improved convergence, lower gradient variance, and enhanced performance compared to traditional methods.

The VIMCO-$\star$ gradient estimator is a theoretically grounded extension of the VIMCO family of REINFORCE-based gradient estimators for importance weighted variational inference (IWVI) objectives. Designed to address the signal-to-noise ratio (SNR) collapse observed in classical VIMCO estimators for large numbers of importance samples, VIMCO-$\star$ employs an optimally tuned control variate that yields SNR scaling $\propto \sqrt{N}$ even in the limiting IWAE regime. This property contrasts the degraded SNR in standard score-function and VIMCO estimators, making VIMCO-$\star$ empirically and theoretically preferable for non-reparameterizable models and high-$N$ settings (Daudel et al., 1 Feb 2026, Liévin et al., 2020).

1. Background: Importance-Weighted Bounds and Score-Function Gradients

The IWVI paradigm tightens variational lower bounds on log-marginal likelihood as the number of importance samples $N$ increases, crucially relying on unbiased and low-variance gradient estimators for stochastic optimization. The classical IWAE objective is

$$\mathcal{L}_N(\theta, \phi) = \mathbb{E}_{z_{1:N} \sim q_\phi}\left[ \log \left( \frac{1}{N} \sum_{i=1}^{N} w(z_i) \right) \right], \qquad w(z) = \frac{p_\theta(x, z)}{q_\phi(z|x)},$$

where $q_\phi(z|x)$ is the variational approximation and $p_\theta(x, z)$ is the model joint. Optimization with respect to $\phi$ is challenging in non-reparameterizable models, where the reparameterization trick cannot be applied due to discrete latents or simulator constraints. In these contexts, score-function (REINFORCE) estimators and their control-variate-enhanced variants are necessary (Daudel et al., 1 Feb 2026, Liévin et al., 2020, Mnih et al., 2016).

2. Classical VIMCO Estimators and the SNR Collapse

Mnih & Rezende introduced the original VIMCO estimator as a variance-reducing alternative to the naïve REINFORCE gradient for IWAE objectives (Mnih et al., 2016). VIMCO employs a leave-one-out (LOO) control variate—computed for each sample $z_i$ as a function of the other $N-1$ samples—to construct per-sample learning signals. While unbiased and providing improved credit assignment, the gradient estimator remains susceptible to an SNR that, asymptotically in $N$ for $\alpha = 0$ (IWAE case), decays as $O(1/\sqrt{N})$ (Daudel et al., 1 Feb 2026). This SNR collapse means optimization becomes ineffective for large $N$, despite the bound tightening, because the estimator's noise dominates its signal.

3. Theoretical Development of the VIMCO-$\star$ Estimator

VIMCO-$\star$ is derived via an asymptotic analysis of variance and bias for generalized VIMCO objectives parameterized by $\alpha \in [0, 1)$, which interpolate between IWAE ($\alpha = 0$) and variational Rényi bounds. The core result is a closed-form optimal baseline $f_{-i}^{(\alpha,\star)}$ that, when used in place of VIMCO's ad hoc choices (arithmetic mean or geometric mean), minimizes the leading-order asymptotic variance of the gradient estimator. For $\alpha = 0$, this optimal baseline vanishes, considerably simplifying the estimator. The final VIMCO-$\star$ gradient formula is: $$\hat{g}_\psi^\star = \sum_{j=1}^N \bar{w}_j\, \partial_\psi \log w(z_j) - \frac{1}{1-\alpha} \sum_{i=1}^N \partial_\psi \log q_\phi(z_i)\, \log\left(1-\bar{w}_i + \frac{f_{-i}^{(\alpha,\star)}}{\sum_j w_j^{1-\alpha}}\right)$$ where

$$\bar{w}_j = \frac{w(z_j)^{1-\alpha}}{\sum_\ell w(z_\ell)^{1-\alpha}},$$

and for $\alpha = 0$, $f_{-i}^{(\alpha,\star)} \to 0$, yielding

$$\hat{g}_\psi^\star = \sum_j \bar{w}_j\, \partial_\psi \log w_j - \sum_i \partial_\psi \log q_\phi(z_i)\, \log(1-\bar{w}_i).$$

(Daudel et al., 1 Feb 2026).

4. Signal-to-Noise Ratio Analysis and Optimality

Extensive theoretical analysis establishes that for VIMCO-$\star$:

• For $\alpha=0$ (IWAE), variance of the estimator is $O(1/N^3)$ and SNR scales as $O(\sqrt{N})$.
• For general $\alpha$, VIMCO-$\star$ minimizes the leading-order constant in $\mathrm{Var}=O(1/N)$ among all possible constant control variates. By contrast, standard VIMCO-AM and VIMCO-GM only achieve SNR $\propto 1/\sqrt{N}$ for $\alpha=0$, causing performance to degrade with increasing $N$. These results rigorously justify the superior scaling of VIMCO-$\star$ and explain its empirically observed robustness for large $N$, directly resolving the "stalling" pathology of vanilla VIMCO (Daudel et al., 1 Feb 2026, Liévin et al., 2020).

5. Algorithmic Structure and Implementation

The VIMCO-$\star$ estimator may be implemented in any autodiff-compatible framework with computational cost $O(N)$ per gradient update. For $\alpha=0$, no additional baseline is required beyond VIMCO logic. For $\alpha>0$, baseline estimation can be subsampled or amortized with a running average. The algorithm is as follows:

1. Initialize parameters $(\theta,\phi)$.
2. For each iteration:
   • Sample $z_{1:N} \sim q_\phi(\cdot|x)$.
   • Compute weights $w_i$ and normalized weights $\bar{w}_i$.
   • Compute score-function and baseline-corrected REINFORCE terms.
   • Update parameters with the sum of these two contributions.
3. If $\alpha>0$, update the baseline estimate using current and previous minibatches (Daudel et al., 1 Feb 2026).

6. Empirical Evaluation

Empirical assessments across latent Gaussian models, stochastic volatility (state-space) models, and variational phylogenetics confirm the theoretical predictions:

• In Gaussian models, for $\alpha>0$ all estimators have SNR $\propto \sqrt{N}$, but VIMCO-$\star$ achieves the largest constant, and only VIMCO-$\star$ retains this scaling in $\alpha=0$ regimes.
• In likelihood-free state-space settings, integrating VIMCO-$\star$ into Riemannian-grouped variational inference yields higher effective sample size, lower gradient variance, and improved convergence relative to VIMCO-AM/GM.
• In variational Bayesian phylogenetics (VBPI), VIMCO-$\star$ produces faster convergence and systematically lower KL divergence to MCMC ground truth (Daudel et al., 1 Feb 2026).

A summary that compares key variants:

Estimator	Control Variate	SNR scaling for $\alpha=0$	SNR scaling for $\alpha>0$
VIMCO-AM/GM	Arithmetic/Geometric	$O(1/\sqrt{N})$	$O(\sqrt{N})$
VIMCO-$\star$	Optimal	$O(\sqrt{N})$	$O(\sqrt{N})$ (largest constant)

7. Relation to Prior Work and Practical Considerations

VIMCO-$\star$ generalizes the baseline construction in VIMCO (Mnih et al., 2016) and is closely related to the OVIS estimator developed independently in "Optimal Variance Control of the Score Function Gradient Estimator for Importance Weighted Bounds" (Liévin et al., 2020). Both works demonstrate that, in the large-sample regime, control variate optimized estimators can invert the detrimental SNR scaling of classical REINFORCE approaches, without requiring the reparameterization trick. Limitations include the assumption of finite-variance importance weights and the requirement for adequate minibatch sizes to estimate mean and covariance terms involved in the optimal baseline. Practical guidelines typically set the interpolation parameter (where present) to $\gamma=1$ in low-effective-sample-size regimes and $\gamma=0$ for very high ESS, noting robust empirical performance across broad model classes (Daudel et al., 1 Feb 2026, Liévin et al., 2020).

A plausible implication is that VIMCO-$\star$ enables effective IWVI optimization in previously inaccessible domains (discrete, likelihood-free, phylogenetic), mitigating the tradeoff between estimator variance and bound tightness that has constrained non-reparameterizable VI schemes.

• Importance Weighted Variational Inference without the Reparameterization Trick (Daudel et al., 1 Feb 2026)
• Optimal Variance Control of the Score Function Gradient Estimator for Importance Weighted Bounds (Liévin et al., 2020)
• Variational inference for Monte Carlo objectives (Mnih et al., 2016)