---
title: REINFORCE Leave-One-Out Gradients
url: https://www.emergentmind.com/topics/reinforce-leave-one-out-gradients
type: topic
---

# REINFORCE Leave-One-Out Gradients

REINFORCE leave-one-out (RLOO) gradients refer to a class of variance-reduced gradient estimators for Monte Carlo integration of expectations with respect to parameterized probability distributions, particularly within variational inference (VI) using the score function (REINFORCE) method. These estimators combine the score-function approach with leave-one-out control variates, yielding unbiased but lower-variance gradient estimates, especially suited for discrete latent variable models and variational autoencoders.

## 1. Score-Function (REINFORCE) Estimators and Their Variants

The standard REINFORCE estimator targets gradients of the evidence lower bound (ELBO)
$$
\mathrm{ELBO}(\theta) = \mathbb{E}_{q_\theta(z)} \left[ \log p(x,z) - \log q_\theta(z) \right].
$$
By the score-function (SF) identity,
$$
\nabla_\theta \, \mathrm{ELBO}(\theta) = \mathbb{E}_{q_\theta(z)}\left[ \left(\log p(x,z) - \log q_\theta(z)\right) \nabla_\theta \log q_\theta(z) \right].
$$
Naive Monte Carlo score-function gradients typically exhibit high variance due to sample correlations, especially for discrete variables. To mitigate this, baseline control variates are often subtracted. The leave-one-out (LOO) REINFORCE variant, also known as RLOO, uses a sample-specific baseline:
$$
\widehat{G}_{\mathrm{RLOO}} = \frac1K \sum_{k=1}^K \left( f(x_k) - \bar{f}_{-k} \right) \nabla_\eta \log q_\eta(x_k),
$$
where $\bar{f}_{-k}$ is the average of $f(x_j)$ over all $j \neq k$ [2111.05300].

## 2. Log-Variance Loss and the VarGrad Estimator

VarGrad, introduced by Richter et al. (2020), generalizes RLOO by reframing ELBO gradient estimation as the gradient of a "log-variance loss." For a reference density $r(z)$,
$$
L_r(\theta) = \frac{1}{2} \operatorname{Var}_r \left[ \log \frac{q_\theta(z)}{p(z|x)} \right ].
$$
When $r = q_\theta$, this loss is a divergence that vanishes if the approximate and true posteriors coincide. The gradient of $L_{q_\theta}$ yields a score-function term plus a mean correction:
$$
\nabla_\theta L_r(\theta) = \mathbb{E}_r[ f(z)\nabla_\theta \log q_\theta(z)] - \mathbb{E}_r[f]\;\mathbb{E}_r[\nabla_\theta \log q_\theta(z)].
$$
At $r=q_\theta$, the mean term cancels, recovering $\nabla_\theta(-\mathrm{ELBO}(\theta))$ [2010.10436].

The empirical approximation,
$$
\widehat{L}(\theta) = \frac{1}{2(S-1)} \sum_{s=1}^S (f(z^{(s)}) - \bar{f})^2, \quad \bar{f} = \frac{1}{S} \sum_{s=1}^S f(z^{(s)}),
$$
yields, after differentiation, the VarGrad estimator:
$$
\widehat{g}_{\mathrm{VarGrad}}(\theta) = \frac{1}{S-1} \left( \sum_{s=1}^S f_s \cdot sgrad_s - \bar{f} \sum_{s=1}^S sgrad_s \right ),
$$
where $sgrad_s = \nabla_\theta \log q_\theta(z^{(s)})$ [2010.10436]. This form precisely recovers a leave-one-out baseline with near-optimal coefficient and rescaling.

## 3. Stein-Based and Double Control Variate Extensions

Further variance reduction is possible by introducing additional control variates. Stein operators, as developed in the "RODEO" framework, yield flexible zero-mean corrections in discrete spaces. For a discrete $q(x)$ and Markov kernel $P$, the Stein operator $A = P - I$ satisfies $\mathbb{E}_q[(A h)(x)] = 0$. The RLOO estimator can thus be augmented with local and global Stein control variates without bias:
$$
\widehat{g}_{\rm RODEO} = \frac{1}{K} \sum_{k=1}^K \left( (f(x^{(k)}) - b_{-k}^{\mathrm{Stein}})\nabla_\eta \log q_\eta(x^{(k)}) + (A \tilde h^*_k)(x^{(k)}) \right),
$$
with 
$$
b_{-k}^{\mathrm{Stein}} = \frac1{K-1} \sum_{j\neq k} (f(x^{(j)}) + (A h_j)(x^{(j)})).
$$
Empirically, RODEO significantly lowers gradient variance in generative modeling, achieving state-of-the-art ELBO and convergence trends with the same computational budget as standard RLOO [2202.09497].

Double control variate (DoubleCV) methods further exploit auxiliary functions—typically constructed from first-order Taylor expansions—to form additional, sample-specific corrections atop the leave-one-out baseline. When optimally combined, these yield strictly lower variance than RLOO, sometimes even surpassing estimators using the unattainable "true mean" baseline [2111.05300].

## 4. Theoretical Properties and Variance Analysis

The variance improvement gained by RLOO and its extensions is rooted in the optimal baseline selection problem. The optimal coefficient for a control variate $a^*_i$ in dimension $i$ is
$$
a^*_i = \frac{\operatorname{Cov}_{q_\theta}(f \, \partial_{\theta_i}\log q,\, \partial_{\theta_i}\log q)}{\operatorname{Var}_{q_\theta}(\partial_{\theta_i} \log q)}.
$$
VarGrad uses the sample mean $\bar{f}$ as its baseline coefficient, which is close to optimal under broad conditions: specifically, when $\operatorname{KL}(q_\theta \| p)$ is either very large (early in training) or very small (late in training), and moments of the score are bounded [2010.10436].

In high-dimensional regimes ($D$ large) and with sufficiently many samples, VarGrad's variance
$$
\operatorname{Var}(\widehat{g}_{\mathrm{VarGrad}})_i \leq \operatorname{Var}(\widehat{g}_{\mathrm{REINFORCE}})_i
$$
provably holds. Stein-augmented estimators (RODEO) and DoubleCV, when optimally tuned, analogously achieve variance lower than RLOO [2202.09497, 2111.05300].

## 5. Practical Implementation and Algorithmic Details

VarGrad and its extensions operate as black-box, parameter-free gradient estimators requiring only standard score-function gradient machinery. The implementation closely follows the following steps, exemplified for VarGrad [2010.10436]:
1. Sample $S$ i.i.d. draws $z^{(s)} \sim q_\theta$ (samples detached for “stop-gradient” semantics).
2. Compute $f_s = \log q_\theta(z^{(s)}) - \log p(x, z^{(s)})$.
3. Assemble sample mean $\bar f$ and gradients $sgrad_s = \nabla_\theta \log q_\theta(z^{(s)})$.
4. Return
   $$
   \widehat{g}_{\mathrm{VarGrad}}(\theta) = \frac{1}{S-1}\left( \sum_{s=1}^S f_s sgrad_s - \bar{f} \sum_{s=1}^S sgrad_s \right).
   $$
Algorithmic enhancements for VAEs and other models leverage autodiff to obtain required derivatives, incur negligible additional computational cost, and introduce no extra passes through the decoder or sampling distribution [2111.05300].

## 6. Empirical Evaluation and Comparison

Empirical results demonstrate that RLOO, VarGrad, RODEO, and DoubleCV provide favorable variance-computation trade-offs in both synthetic and real-world settings:
- In discrete VAEs (Bernoulli latent models, e.g., Omniglot), VarGrad (with 4 samples) achieves learning curves nearly matching those of REBAR, RELAX, and ARM, but at reduced per-step computational cost.
- RODEO achieves up to an order-of-magnitude variance reduction relative to RLOO and DoubleCV, matching or exceeding ELBOs of competing estimators at fixed budget [2202.09497].
- DoubleCV converges in fewer steps and achieves uniformly lower variance and higher ELBOs than RLOO, as illustrated in benchmark studies on MNIST, Fashion-MNIST, and Omniglot [2111.05300].

A summary table of ELBO comparisons for $K=2$ on Bernoulli-likelihood VAEs:

| Estimator    | MNIST            | Fashion-MNIST     | Omniglot         |
|--------------|------------------|-------------------|------------------|
| RLOO         | $-103.11 \pm 0.16$ | $-241.53 \pm 0.24$ | $-116.83 \pm 0.05$ |
| DoubleCV     | $\mathbf{-102.45 \pm 0.13}$ | $\mathbf{-240.96 \pm 0.17}$ | $\mathbf{-116.22 \pm 0.08}$ |
| DisARM       | $-102.56 \pm 0.09$ | $-241.02 \pm 0.20$ | $-116.36 \pm 0.05$ |

## 7. Assumptions, Limitations, and Regimes of Use

These methods assume that the support of $q_\theta$ contains the support of $p$ to ensure well-definedness of the control variate and gradient terms. All estimators require that sampling and baseline construction do not differentiate through the samples (“stop-gradient” on draws). Tight variance guarantees hold under additional tail-regularity and bounded kurtosis conditions on the score, with the most significant improvements observed in large-sample ($S$) and high-dimensional ($D$) settings [2010.10436].

RODEO’s efficiency may be limited in moderate-to-high dimensional discrete spaces when full evaluation of neighbor states for the Stein operators is infeasible, though surrogate strategies mitigate this. DoubleCV, when adapted as recommended, incurs no additional autodiff or sampling cost in standard VAE implementations.

---

The development and refinement of REINFORCE leave-one-out gradient estimators—including VarGrad, Stein-augmented, and double control variate extensions—represent key advances in low-variance, unbiased gradient estimation for variational inference in discrete latent variable frameworks [2010.10436, 2202.09497, 2111.05300].

Source: https://www.emergentmind.com/topics/reinforce-leave-one-out-gradients