---
title: Rao-Blackwellised Gradient Estimator
url: https://www.emergentmind.com/topics/rao-blackwellised-gradient-estimator
type: topic
---

# Rao-Blackwellised Gradient Estimator

A Rao-Blackwellised gradient estimator is a stochastic gradient estimator that leverages the Rao-Blackwell theorem to reduce variance by analytically marginalizing over one or more latent random variables, rather than resorting purely to Monte Carlo sampling. This approach has found broad application across variational inference, discrete and continuous stochastic optimization, probabilistic deep learning, and reinforcement learning, yielding unbiased (or lower-bias) and lower-variance estimators that often require only moderate additional computation. It subsumes and generalizes several classical techniques and is particularly effective in high-dimensional or structured latent-variable models where naive score-function estimators or standard reparameterization gradients suffer from excessive variance.

## 1. Definition and Fundamental Principle

The core principle behind Rao-Blackwellised gradient estimators is the reduction of the variance of an unbiased (or otherwise useful) Monte Carlo gradient estimator by conditioning on a convenient sufficient statistic, performing exact integration over a tractable subset of latent variables (or labels). In a stochastic gradient setting, suppose the gradient estimator $g(z)$ is unbiased for the gradient $\nabla_\theta L(\theta)$, with $z$ a random variable sampled from distribution $q_\theta$. The Rao-Blackwell theorem guarantees that the conditional expectation $\tilde{g}(y) = E[g(z) \mid y]$ satisfies $E[\tilde{g}(y)] = \nabla_\theta L(\theta)$ but with lower variance: $\operatorname{Var}[\tilde{g}(y)] \leq \operatorname{Var}[g(z)]$, with strict inequality unless $g(z)$ is almost surely a function of $y$ only.

In practice, this is realized by drawing a full sample from $q_\theta$, then, for the subset of variables most closely associated with a given parameter block or computational bottleneck, replacing Monte Carlo expectations with analytically tractable marginalizations (such as exact summation over all values of a discrete random variable or one-dimensional numerical quadrature in continuous cases) [1503.01494].

## 2. Applications Across Latent Variable Models

### 2.1. Variational Inference—Local Expectation Gradients

In variational inference, particularly for latent-variable models $p(y, x)$ and structured variational families $q_\nu(x)$, stochastic optimization of the evidence lower bound (ELBO) typically employs either the log-derivative (score-function) gradient estimator or the reparameterization trick. The local expectation gradient estimator—explicitly a Rao-Blackwellised score-function gradient—achieves lower variance by integrating over the latent variable $x_i$ most directly associated with parameter $v_i$:

\[
\nabla_{v_i}\mathcal{F}(\nu) = \mathbb{E}_{q(x_{\setminus i})}\left[\mathbb{E}_{q(x_i|\mathrm{mb}_i)}[f(x)\nabla_{v_i} \log q_{v_i}(x_i|\mathrm{pa}_i)]\right]
\]

This estimator is applicable for both discrete and continuous $x_i$ (using quadrature or summation), is unbiased, and empirically lowers variance by about an order of magnitude versus standard reparameterization at fixed computational budget. Local expectation can be computed for all parameters in parallel, at per-iteration complexity $O(nK)$, yielding fast, stable convergence in high-dimensional and structured latent variable models [1503.01494].

### 2.2. Discrete Distributions and "Top-K" Summing Out

For gradients involving expectations over large discrete spaces (e.g., categorical variables with $K \gg 1$), the Rao-Blackwellized estimator can be implemented by summing out the $k$ categories with the largest mass in $q_\eta$, and sampling over the complement:

\[
\hat{G}_{RB} = \sum_{z \in C_k} q_\eta(z) g(z) + q_\eta(\bar{C}_k) g(v)
\]
where $v \sim q_\eta(\cdot|\bar{C}_k)$.

This technique preserves unbiasedness and leverages extreme concentration of probability mass (low-entropy settings) to drive variance down rapidly as $k$ increases, requiring only minimal extra cost over vanilla REINFORCE. Practically, even $k=1$ yields a $>90\%$ reduction in variance in highly concentrated cases [1810.04777].

### 2.3. Rao-Blackwellisation with Sampling Without Replacement

Variants based on sampling unordered sets without replacement (e.g., using Plackett-Luce samplers) further lower variance in expectations over discrete random variables. By conditioning on the unordered set and analytically computing leave-one-out ratios for each element, unbiased estimators with substantially reduced variance—especially in low-entropy or peaky distributions—are obtained, often outperforming independent-sampling baselines [2002.06043].

### 2.4. Continuous Latent Variables—Reparameterisation Gradients

For continuous latent Gaussian variables with linear or neural parameterizations, Rao-Blackwellising the reparameterization gradient (e.g., as in the R2-G2 estimator) analytically marginalizes over sampled pre-activation noise, yielding gradients of the form

\[
\widehat{\nabla_{\phi}f}^{\,R2\textrm{-}G2} := \mathbb{E}_{\varepsilon | w}\left[ \widehat{\nabla_{\phi}f} \right ]
\]
where $w$ is a linear function of the Gaussian variable.

This construction strictly reduces variance compared to standard reparameterization, preserves unbiasedness, and generalizes the local reparameterization trick for BNNs [2506.07687].

### 2.5. Augment-REINFORCE-Swap-Merge (ARSM) Estimator

The ARSM estimator applies Rao-Blackwellisation to categorical expectations, re-expressing the gradient as an expectation over a Dirichlet distribution, and leveraging permutation symmetry (swap/merge) to construct variance-reduced estimators. The resulting estimator matches the true gradient up to Monte Carlo error, with variance $10^2$--$10^4$ times lower than REINFORCE or RELAX [1905.01413].

## 3. Algorithmic Implementations

Across settings, the general recipe for a Rao-Blackwellised gradient estimator is:

- Draw a full or partial pivot sample from the variational or target distribution.
- For each sub-block/parameter/latent variable of interest, analytically compute the expectation of the gradient estimator over the variable(s) most directly coupled to the parameter, conditioned on the rest of the sample.
- Aggregate those conditional expectations—typically one per parameter block—in a parallelizable or batched fashion.
- Update parameters using the resulting low-variance gradient blocks.

If the conditional expectation is analytic (finite sum for discrete, tractable integral for continuous, or closed-form Dirichlet/conditional Gumbel, etc.), then no additional stochasticity is introduced, and the computational cost is only marginally higher than the naive estimator. Example implementations include:

- Local expectation gradients for variational inference [1503.01494]
- Top-K category summing for large discrete support [1810.04777]
- Sampling-without-replacement/leave-one-out methods [2002.06043]
- R2-G2 for latent Gaussians [2506.07687]
- Gumbel-Rao or Rao-Blackwellised ST-Gumbel-Softmax [2010.04838]
- ReinMax-Rao for categorical/straight-through gradients [2603.08257]

## 4. Variance Reduction, Bias, and Theoretical Properties

The variance reduction property of Rao-Blackwellisation is guaranteed by the law of total variance: conditioning on an analytically tractable sufficient statistic always yields an estimator with variance no larger than the original, with equality only if the original estimator is already deterministic given the statistic. Multiple studies quantify specific reductions; for instance, summing out one category in a highly concentrated discrete distribution can cut REINFORCE variance by over 90% [1810.04777]; local expectation gradients are empirically about an order of magnitude lower in variance than vanilla reparameterization in both synthetic and real-world examples [1503.01494].

Bias properties are preserved in almost all Rao-Blackwellisation cases, since the estimator is constructed as a conditional expectation of an unbiased estimator. Some advanced constructions (e.g., Gumbel-Softmax or ReinMax-Rao) operate on already biased surrogate estimators, but variance is still strictly reduced at fixed bias [2010.04838][2603.08257].

## 5. Empirical Performance and Application Case Studies

Extensive empirical studies consistently demonstrate the effectiveness of Rao-Blackwellised estimators:

- High-dimensional Gaussian variational inference: order-of-magnitude variance improvements, faster stable ELBO convergence, ability to use smaller quadrature/single sample compared to thousands of MC samples for naive estimators [1503.01494].
- Discrete-latent VAEs: top-K RB estimators achieve rapid ELBO convergence and high classification accuracy with minimal computational overhead [1810.04777].
- Categorical VAEs, RL policy gradients, and structured prediction: ARSM and unordered-set RB estimators match or surpass all prior baselines, show 10–1,000× gradient variance reduction, and avoid additional control-variate or baseline tuning [1905.01413][2002.06043].
- Risk-sensitive control: Rao-Blackwellised score climbing with particle filters yields unbiased, low-variance policy gradients within SMC frameworks for non-Gaussian, nonlinear stochastic dynamical systems [2312.14000].

## 6. Comparisons with Alternative Gradient Estimators

The Rao-Blackwellised approach generalizes and often strictly dominates classical methods:

| Estimator        | Applicability               | Variance      | Unbiasedness        | Structure / Notes          |
|------------------|----------------------------|---------------|---------------------|----------------------------|
| Score-function   | Arbitrary (discrete, cont.) | High          | Unbiased            | Simple, needs control-variate |
| Reparameterized  | Cont., differentiable $f$   | Moderate      | Unbiased            | Low-variance if appl.       |
| Rao-Blackwellised| Discrete/cont./graph-factored| Lowest        | Unbiased            | Tractable marginalization   |

In many settings, Rao-Blackwellised estimators inherit the generality of score-function methods but match or outperform reparameterized estimators when the latter are available. Rao-Blackwellisation can also be applied to various straight-through, surrogate, or policy-gradient estimators as a meta-optimization step [1503.01494][2010.04838][2603.08257][1905.01413].

## 7. Limitations and Practical Considerations

While the computational cost of marginalizing over a single variable per parameter block is typically moderate (e.g., $O(K)$ for K quadrature points per iteration), applications to high-arity or high-dimensional discrete variables can be challenging unless strong structural factorization or probability mass concentration is present [1810.04777][2002.06043]. In Bayesian neural networks and deep VAEs, efficient matrix solves or blockwise analytic marginalization suffice for practical deployment [2506.07687]. For particle-based sequential models, per-iteration cost scales with the number of particles but the variance reduction justifies the additional computation [2312.14000].

## References

- "Local Expectation Gradients for Doubly Stochastic Variational Inference" [1503.01494]
- "Rao-Blackwellized Stochastic Gradients for Discrete Distributions" [1810.04777]
- "Estimating Gradients for Discrete Random Variables by Sampling without Replacement" [2002.06043]
- "Rao-Blackwellised Reparameterisation Gradients" [2506.07687]
- "Rao-Blackwellizing the Straight-Through Gumbel-Softmax Gradient Estimator" [2010.04838]
- "Beyond ReinMax: Low-Variance Gradient Estimators for Discrete Latent Variables" [2603.08257]
- "ARSM: Augment-REINFORCE-Swap-Merge Estimator for Gradient Backpropagation Through Categorical Variables" [1905.01413]
- "Risk-Sensitive Stochastic Optimal Control as Rao-Blackwellized Markovian Score Climbing" [2312.14000]

Source: https://www.emergentmind.com/topics/rao-blackwellised-gradient-estimator