Rao-Blackwellized Estimation from MCMC Output

Replacing Latent Draws by Conditional Expectations

MCMC output often contains a latent variable whose contribution to the quantity of interest is noisy even after the other variables are known. If its conditional expectation can be calculated, replacing the sampled contribution by that expectation removes some simulation noise. This procedure is called Rao–Blackwellization. The variance reduction is exact for conditional expectations, but in MCMC one must also check what has been conditioned on: dependence between iterations means that a per-iteration replacement is not automatically the conditional expectation of the entire sample average.

Let $(\Omega,\mathcal{F},\mathbb{P})$ be a probability space, and let $(X_n)_{n\ge0}$ be a discrete-time Markov chain with finite state space $S$ and transition matrix $P=(p_{ij})$, with rows summing to $1$. Suppose its stationary distribution is $\pi_X$. For the variance comparison below, assume $X_0\sim\pi_X$, so the chain is stationary. Given the states $X_1,\ldots,X_N$, suppose $Z_1,\ldots,Z_N$ are conditionally independent draws from a specified conditional distribution of a latent variable $Z$ given $X_n$. This construction gives a legitimate MCMC output with an auxiliary variable freshly sampled at each iteration.

Let $f(X_n,Z_n)$ be a real-valued quantity with finite second moment, and define

H(x)=\mathbb{E}[f(x,Z)\mid X=x],\qquad \theta=\mathbb{E}_{\pi_X}[H(X)].

The ordinary estimator averages the sampled contributions, $\widehat\theta_N=N^{-1}\sum_{n=1}^N f(X_n,Z_n)$. Its Rao–Blackwellized version is $\widetilde\theta_N=N^{-1}\sum_{n=1}^N H(X_n)$. Both target $\theta$ when the chain is stationary. The conditional independence assumption is stated because it permits a direct finite-sample variance comparison for these two averages.

Theorem (Rao–Blackwell variance comparison for conditionally resampled output): Under the assumptions above, $$ \mathbb{E}[\widehat\theta_N]=\mathbb{E}[\widetilde\theta_N]=\theta,\qquad \operatorname{Var}(\widehat\theta_N) =\operatorname{Var}(\widetilde\theta_N) +\frac{1}{N^2}\sum_{n=1}^N\mathbb{E}\!\left[\operatorname{Var}(f(X_n,Z_n)\mid X_1,\ldots,X_N)\right]. $$ In particular, $\operatorname{Var}(\widetilde\theta_N)\le\operatorname{Var}(\widehat\theta_N)$. Equality holds exactly when every conditional variance in the sum is zero almost surely.

Proof. Write $\mathcal{G}=\sigma(X_1,\ldots,X_N)$, the information in the sampled state path. Conditional independence and the specified conditional distributions give $\mathbb{E}[\widehat\theta_N\mid\mathcal{G}]=\widetilde\theta_N$. The conditional variance of the sum is the sum of the conditional variances, since the $Z_n$ are conditionally independent given $\mathcal{G}$. The law of total variance, $\operatorname{Var}(Y)=\operatorname{Var}(\mathbb{E}[Y\mid\mathcal{G}])+\mathbb{E}[\operatorname{Var}(Y\mid\mathcal{G})]$, therefore gives the displayed identity. Stationarity gives $\mathbb{E}[H(X_n)]=\theta$, and conditional expectation gives $\mathbb{E}[f(X_n,Z_n)]=\mathbb{E}[H(X_n)]=\theta$. Each conditional variance is nonnegative, proving the inequality. The sum is zero exactly when each of its nonnegative terms is zero almost surely. $\square$

The identity separates two sources of variation. The state path contributes $\operatorname{Var}(\widetilde\theta_N)$, which remains after Rao–Blackwellization and includes MCMC autocorrelation. The conditional latent draws contribute the additional nonnegative term. Thus Rao–Blackwellization removes conditional simulation noise; it does not remove dependence in the state chain.

A Discrete Latent-Variable Example

Worked Example: A Two-State MCMC Target

Let $X_n\in S=\{0,1\}$ have transition matrix $$ P=\begin{bmatrix}\tfrac34&\tfrac14\\[2pt]\tfrac14&\tfrac34\end{bmatrix}. $$ Each row sums to $1$. Its stationary distribution is $\pi_X=(\tfrac12,\tfrac12)$: multiplying gives $(\tfrac12,\tfrac12)P=(\tfrac12,\tfrac12)$, and its entries sum to $1$.

Given $X_n=0$, let $Z_n$ be Bernoulli with success probability $\tfrac14$; given $X_n=1$, let $Z_n$ be Bernoulli with success probability $\tfrac34$. Given the state path, generate the $Z_n$ independently. The target is $\theta=\mathbb{E}[Z_n]$. The conditional expectation is $H(0)=\tfrac14$ and $H(1)=\tfrac34$, so $$ \theta=\tfrac12\cdot\tfrac14+\tfrac12\cdot\tfrac34=\tfrac12. $$ This also checks the target using the marginal distribution of $Z_n$.

Worked Example: Both Estimates on One Realized Path

Consider the four-state path $X_1,X_2,X_3,X_4=(0,0,1,0)$, with latent draws $Z_1,Z_2,Z_3,Z_4=(0,1,1,0)$. The Rao–Blackwellized contributions are $H(X_n)=(\tfrac14,\tfrac14,\tfrac34,\tfrac14)$, so $$ \widetilde\theta_4=\frac{\tfrac14+\tfrac14+\tfrac34+\tfrac14}{4} =\frac{\tfrac32}{4}=\tfrac38. $$ The ordinary estimate is $$ \widehat\theta_4=\frac{0+1+1+0}{4}=\tfrac12. $$ The two realized estimates need not be ordered: here the ordinary estimate is larger. The theorem compares their variances over repeated output, not their values on every particular path.

For a conditional check, the four Bernoulli probabilities are $(\tfrac14,\tfrac14,\tfrac34,\tfrac14)$. Each conditional variance is $p(1-p)=\tfrac3{16}$; hence the conditional variance of $\widehat\theta_4$ is $$ \operatorname{Var}(\widehat\theta_4\mid X_1,\ldots,X_4) =\frac{4(\tfrac3{16})}{4^2}=\tfrac3{64}. $$ The conditional mean is $\tfrac38$, so the conditional second moment of the ordinary estimate is $(\tfrac38)^2+\tfrac3{64}=\tfrac9{64}+\tfrac3{64}=\tfrac3{16}$. By contrast, the average of the four individual conditional second moments is $(\tfrac14+\tfrac14+\tfrac34+\tfrac14)/4=\tfrac38$, because $Z_n^2=Z_n$. These are different quantities: the average of individual second moments is not the second moment of the sample average.

Checking the Variance Reduction with MCMC Dependence

Worked Example: Exact Variances for Four Iterations

For the chain and conditional draws above, compute both unconditional variances for $N=4$. The identity in the theorem gives the variance difference directly. Since each state has conditional Bernoulli variance $\tfrac3{16}$, the expected conditional variance of the ordinary average is $4(\tfrac3{16})/16=\tfrac3{64}$.

To find the variance of $\widetilde\theta_4$, write $H(X_n)=\tfrac14+\tfrac12\mathbf{1}_{\{X_n=1\}}$. Under the stationary chain, $\operatorname{Var}(H(X_n))=\tfrac1{16}$. The transition matrix shows that the centered indicator has conditional expectation one half its current value after one step; iterating gives $\operatorname{Cov}(H(X_n),H(X_{n+k}))=\tfrac1{16}(\tfrac12)^k$. Consequently, $$ \begin{aligned} \operatorname{Var}(\widetilde\theta_4) &=\frac{1}{16}\left[4+2\left(3\cdot\tfrac12+2\cdot\tfrac14+1\cdot\tfrac18\right)\right]\frac1{16}\\ &=\frac{1}{256}\left[4+2\left(\tfrac32+\tfrac12+\tfrac18\right)\right] =\frac{33}{1024}. \end{aligned} $$ The bracket counts the four variance terms and both covariance terms for each positive lag. Adding the conditional-draw contribution gives $$ \operatorname{Var}(\widehat\theta_4)=\frac{33}{1024}+\frac3{64} =\frac{33}{1024}+\frac{48}{1024}=\frac{81}{1024}. $$ Thus the Rao–Blackwellized average has smaller variance, but its variance is not zero: the MCMC state path is still random and dependent.

Check: Both estimates have mean $\tfrac12$. The difference of their computed variances is $\tfrac{81}{1024}-\tfrac{33}{1024}=\tfrac{48}{1024}=\tfrac3{64}$, exactly the expected conditional variance calculated above.

What Can Be Conditioned On?

For any square-integrable estimator $Y$ and any information $\mathcal{G}$, the conditional expectation $\mathbb{E}[Y\mid\mathcal{G}]$ has no greater variance than $Y$. This follows from the same total-variance identity and does not require independent MCMC iterations. However, for a dependent latent-variable chain, conditioning the whole sample average on the state path can use more than the one-time conditional expectations $H(X_n)$. In general, $\mathbb{E}[f(X_n,Z_n)\mid X_1,\ldots,X_N]$ need not equal $H(X_n)$, because other states in the path may reveal information about $Z_n$.

Therefore, the per-iteration estimator $N^{-1}\sum_n H(X_n)$ has the stated finite-sample variance guarantee when the latent draws are conditionally independent given the state path, as assumed here. In a joint MCMC run where latent states themselves persist and are correlated, one must either establish the relevant conditional-expectation relation for that sampler or use the conditional expectation of the entire estimator given the chosen information. Merely replacing each term by a one-time conditional mean is not, by itself, a proof of a lower variance for the full MCMC average.

Common pitfall: A lower conditional variance does not mean every realized Rao–Blackwellized estimate is closer to the target. In the path example, the raw estimate is $\tfrac12$ and the conditional estimate is $\tfrac38$, while the target is $\tfrac12$. The guarantee concerns expected squared deviation, equivalently variance when the estimators are unbiased. Also, do not identify the average of individual second moments with the second moment of an average; the latter includes cross-products and depends on the average’s variance.

Using the Method in Practice

Before using a Rao–Blackwellized MCMC estimate, identify the quantity being averaged, the conditioning information, and the exact conditional distribution needed to calculate its expectation. Then verify that the conditional expectation has the same target mean. If the output uses fresh conditional latent draws, the variance decomposition above quantifies exactly what is removed. If the latent variables are persistent parts of the chain, account for the dependence rather than assuming that a per-iteration replacement automatically yields the same theorem.

Check Your Understanding

Use the conditional-variance identity and the two-state example to answer these questions.

For $X_n=1$, what are $\mathbb{E}[Z_n\mid X_n]$ and $\operatorname{Var}(Z_n\mid X_n)$?
For the path $(0,0,1,0)$, compute the Rao–Blackwellized estimate and the raw estimate if the latent draws are $(1,0,1,0)$.
Why does the finite-sample variance comparison not imply that the Rao–Blackwellized estimate is smaller on every realized path?
For the four-iteration path in the worked example, distinguish the average of the individual conditional second moments from the conditional second moment of the raw sample average.
What additional conditional-independence property made the per-iteration conditional expectations equal the conditional expectation of the full sample average?