###### tags: `vrmc` `monte carlo` `expository`

# Variance Reduction in Monte Carlo Methods: Rao-Blackwellisation

**Overview**: In this note, I will describe the technique of 'Rao-Blackwellisation' (also *derandomisation*) for variance reduction in the context of Monte Carlo methods. I will also describe an interpretation of this technique as a way of constructing control variates, which enables its use in Monte Carlo algorithms which generate dependent samples.

## Sampling under Conditional Independence Assumptions

When employing Monte Carlo methods, it is often the case that the distribution from which we are sampling admits some kind of simplifying structure. One particular structure which arises in many scenarios is that of *locality*, or the presence of *conditional independence assumptions*.

Suppose that our method involves sampling draws $x$ from a probability measure $p$. In interesting settings, it is often the case that $x$ admits some decomposition into coordinates, i.e. one can write out

\begin{align}
x = (x_1, x_2, \ldots, x_N)
\end{align}

where each of the $x_i$ hold some meaning. For example,

* In materials science, $x$ may represent the state of a molecule, with each of the $x_i$ describing the position of the $i^\text{th}$ constituent atom.
* In statistical physics, $x$ may represent the state of some spatially-varying field, and with each of the $x_i$ describing the intensity of that field at the location $i$.
* In statistics, $x$ may represent all of the random variables in describing a statistical model, and each $x_i$ denotes a separate component of that model.

Strictly speaking, the mere existence of such a decomposition does not say much about the random variable in question. Such a decomposition becomes more meaningful when it is compatible with the probability measure in some way.

One way for this compatibility to manifest is in the form of *conditional independence assumptions*. In short, these impose certain locality conditions on the measure, in the sense that for each $x_i$, knowing the behaviour of only a few components of $\{ x_j \}_{j \neq i}$ is sufficient to understand the conditional behaviour of $x_i$.

Stated more precisely, we would say that the measure $p$ exhibits a form of locality if, for each $i \in [N]$, one can write

\begin{align}
p (dx_i | \{ x_j \}_{j \neq i}) = p_i (dx_i | x_{\partial i})
\end{align}

where $\partial i \subseteq [N]$ has cardinality much less than $N$.

Locality of this form abounds in applications. It is particularly prevalent in physical models, as well as in statistical modelling, where locality often arises from imposing simplifying assumptions on the data-generating process.

It should also be emphasised that in these cases, the conditional distribution $p_i (dx_i | x_{\partial i})$ often has a very simple form, i.e. parametric. When this is the case, it follows that certain expectations under the law of $p_i (dx_i | x_{\partial i})$ can be computed easily in closed form. We will use this to develop our technique.

## Variance Reduction by Derandomisation

Consider now the task of estimating the integral

\begin{align}
I = \int p(dx) f(x)
\end{align}

where $x$ can be decomposed as $x = (x_1, x_2)$, such that the conditionals $\{ p(dx_1 | x_2), p(dx_2 | x_1) \}$ are easy to work with. In particular, suppose that we can easily compute the functions

\begin{align}
f_1 (x_1) = \int p(dx_2 | x_1) f(x_1, x_2) \\
f_2 (x_2) = \int p(dx_1 | x_2) f(x_1, x_2).
\end{align}

A standard example of this would be when one is sampling from a measure for which the conditional distributions take a parametric form (e.g. Gaussian, tilted Rademacher, Gamma, ...), and $f(x)$ is a simple polynomial in $x$, in which case the corresponding $f_1$ is usually available in closed form.

The standard Monte Carlo method would estimate $I$ by:

1. Draw samples $x^1, \cdots, x^M \sim p$.
2. Form the estimator

\begin{align}
\hat{I}_M = \frac{1}{M} \sum_{i = 1}^M f(x^i_1, x^i_2).
\end{align}

However, due to the availability of $f_1, f_2$, we have some other tricks up our sleeves. While we already know, essentially by definition, that $I = \mathbf{E}_p [f]$, one can check that $I$ can also be expressed in both of the following ways:

\begin{align}
I = \mathbf{E}_p [f_1] \\
I = \mathbf{E}_p [f_2].
\end{align}

As such, we can consider forming estimators like

\begin{align}
\hat{I}_M^{(1)} = \frac{1}{M} \sum_{i = 1}^N f_1(x_1^i) \\
\hat{I}_M^{(2)} = \frac{1}{M} \sum_{i = 1}^N f_2(x_2^i)
\end{align}

and these will also constitute reasonable estimators of $I$. These are known as *Rao-Blackwellised* estimators, due to a link to a classical result in mathematical statistics.

## Variance Reduction (?)

Having derived these new estimators, we now ask what we should do with them. A priori, they are just new estimators, and they may not provide any improvements over the original $\hat{I}_M$. Fortunately, this is not the case: both of the new estimators have variance no worse than what we started with.

To prove this (for $f_1$, the same analysis works for $f_2$), assume without a loss of generality that $I = 0$. Then,

\begin{align}
M \cdot \text{Var}_p \left( \hat{I}_M \right) &= \mathbf{E}_p \left[ f(x)^2 \right] \\
&= \mathbf{E}_p \left[ \left(f_1(x_1) + \left\{f(x) - f_1 (x_1) \right\} \right)^2 \right] \\
&= \mathbf{E}_p \left[ f_1(x)^2 \right] + 2 \cdot \mathbf{E}_p \left[ f_1(x_1) \cdot \left\{f(x) - f_1 (x_1) \right\} \right] + \mathbf{E}_p \left[ \left(f(x) - f_1 (x_1)\right)^2 \right].
\end{align}

Now, consider the term

\begin{align}
\mathbf{E}_p \left[ f(x) f_1(x_1) \right] &= \int p(dx_1, dx_2) f(x_1, x_2) f_1 (x_1) \\
&= \int p(dx_1) \int p(dx_2 | x_1)f(x_1, x_2) f_1 (x_1) \\
&= \int p(dx_1) f_1 (x_1) f_1 (x_1) \\
&= \mathbf{E}_p \left[ f_1(x_1)^2 \right].
\end{align}

This tells us that $f - f_1$ and $f_1$ are uncorrelated, and hence the middle term vanishes. Thus,

\begin{align}
M \cdot \text{Var}_p \left( \hat{I}_M \right) &= \mathbf{E}_p \left[ f_1(x)^2 \right] + \mathbf{E}_p \left[ \left(f(x) - f_1 (x_1)\right)^2 \right] \\
&\geqslant \mathbf{E}_p \left[ f_1(x)^2 \right] \\
&= M \cdot \text{Var}_p \left( \hat{I}_M^{(1)} \right)
\end{align}

as desired.

## Control Variate Interpretation

While the above construction validates that the Rao-Blackwellised estimators are never worse in the i.i.d. setting, there is an unfortunate caveat that the same does not hold if the $x^i$ are generated in a correlated fashion, e.g. through the simulation of a Markov chain with equilibrium measure $p$. [Geyer](https://www.jstor.org/stable/1390763) describes some scenarios in which one can explicitly demonstrate that the resulting estimator has **worse** variance. As such, it is desirable to remedy this situation.

As the title of this section suggests, the trick is to interpret Rao-Blackwellisation as a control variate method in which the regression step is carried out in secret. To be more concrete, recall that

\begin{align}
I = \mathbf{E}_p [f] = \mathbf{E}_p [f_1] \\\implies \mathbf{E}_p [f_1 - f] = 0 .
\end{align}

We thus deduce that $f_1 - f$ is a control variate for $p$! As such, for any $\beta$, it holds that

\begin{align}
I = \mathbf{E}_p [f + \beta(f_1 - f)].
\end{align}

Following our previous post about control variates, we should choose $\beta$ to minimise the variance of the resulting estimator. In fact, it turns out that for control variates derived via Rao-Blackwellisation, the optimal value of $\beta$ is always exactly $1$. The eager reader is encouraged to carry out this computation for themselves.

The utility of this observation is that it can be neatly adapted to the setting in which our $x^i$ are correlated, e.g. drawn from a Markov chain. In this setting, the asymptotic variance of the estimator

\begin{align}
\hat{I}_T = \frac{1}{T} \sum_{t = 1}^T f(x_t)
\end{align}

can be written as

\begin{align}
\sum_{k \in \mathbf{Z}} \delta_k
\end{align}

where $\delta_k$ is the lag-$k$ autocovariance of the sequence $\{f_t\}_{t \geqslant 1}$, i.e.

\begin{align}
\delta_k = \mathbf{E} [ f_t f_{t + k}].
\end{align}

One can then consider setting the value of $\beta$ by minimising the asymptotic variance implied by this expression, which will in general be different from $1$, even as $T$ grows to infinity. As one might expect, for a properly-set $\beta$, the resulting Rao-Blackwellised Control Variate estimator is never worse than the original estimator in terms of asymptotic variance.

## Conclusion

Rao-Blackwellisation is a technique for Monte Carlo variance reduction which typically applies when the measure from which we are sampling admits certain tractable conditional distributions. In the i.i.d. setting, it provides a guaranteed variance reduction, and can be quite simple to implement. When the samples are correlated, a naive implementation of the technique can actually make matters worse, but by viewing the technique as an instance of control variates, it is straightforward to remedy this.