Fast Compression of MCMC Output

Abstract

We propose cube thinning, a novel method for compressing the output of an MCMC (Markov chain Monte Carlo) algorithm when control variates are available. It allows resampling of the initial MCMC sample (according to weights derived from control variates), while imposing equality constraints on the averages of these control variates, using the cube method (an approach that originates from survey sampling). The main advantage of cube thinning is that its complexity does not depend on the size of the compressed sample. This compares favourably to previous methods, such as Stein thinning, the complexity of which is quadratic in that quantity.

1. Introduction

MCMC (Markov chain Monte Carlo) remains, to this day, the most popular approach to sampling from a target distribution p, in particular in Bayesian computations [1].

Standard practice is to run a single chain, $X_1,\dots ,X_N$ according to a Markov kernel that leaves invariant p. It is also common to discard part of the simulated chain, either to reduce its memory footprint, or to reduce the CPU cost of later post-processing operations, or more generally for the user’s convenience. Historically, the two common recipes for compressing an MCMC output are:

• burn-in, which allows discarding the b first states;
• thinning, which allows retaining only one out of t (post burn-in) states.

The impact of either recipes on the statistical properties of the subsampled estimates are markedly different. Burn-in reduces the bias introduced by the discrepancy between p and the distribution of the initial state $X_1$ (since $X_b\approx p$ for b large enough). On the other hand, thinning always increases the (asymptotic) variance of MCMC estimators [2].

Practitioners often choose b (the burn-in period) and t (the thinning frequency) separately, in a somewhat ad hoc fashion (i.e., through visual inspection of the initial chain), or using convergence diagnosis such as, e.g., those reviewed in [3].

Two recent papers [4,5] cast a new light on the problem of compressing an MCMC chain by considering, more generally, the problem, for a given M, of selecting the subsample of size M that best represents (according to a certain criterion) the target distribution p. We focus for now on [5], for reasons we explain below.

Stein thinning, the method developed in [5], chooses the subsample $\mathcal{S}$ of size M which minimises the following criterion:

$$D\left(\mathcal{S}\right):=\frac{1}{M^2}\sum _{m,n\in \mathcal{S}}{k}_p(X_m,X_n),\mathcal{S}\subset \{1,\dots ,N\},\left|\mathcal{S}\right|=M$$

(1)

where $k_p$ is a p-dependent kernel function derived from another kernel function k: $\mathcal{X}\times \mathcal{X}\to \mathbb{R}$, as follows:

$$\begin{array}{c}{k}_p(x,y)={\nabla }_x·{\nabla }_{y}k(x,y)+\langle {\nabla }_{x}k(x,y),s_p\left(y\right)\rangle +\langle {\nabla }_{y}k(x,y),s_p\left(x\right)\rangle +k(x,y)\langle s_p\left(x\right),s_p\left(y\right)\rangle \hfill \end{array}$$

with $\langle ·,·\rangle$ being the Euclidean inner product, $s_p\left(x\right):=\nabla logp\left(x\right)$ is the so-called score function (gradient of the log target density), and ∇ is the gradient operator.

The rationale behind criterion (1) is that it may be interpreted as the KSD (kernel Stein discrepancy) between the true distribution p and the empirical distribution of subsample S. We refer to [5] for more details on the theoretical background of the KSD, and its connection to Stein’s method.

Stein thinning is appealing, as it seems to offer a principled, quasi-automatic way to compress an MCMC output. However, closer inspection reveals the following three limitations.

First, it requires computing the gradient of the log-target density, $s_p\left(x\right)=\nabla logp\left(x\right)$. This restricts the method to problems where this gradient exists and is tractable (and, in particular, to $\mathcal{X}={\mathbb{R}}^d$).

Second, its CPU cost is $\mathcal{O}\left(N{M}^2\right)$. This makes it nearly impossible to use Stein thinning for $M\gg 100$. This cost stems from the greedy algorithm proposed in [5], see their Algorithm 1, which adds at iteration t the state $X_i$ which minimises $k_p(X_i,X_i)+{\sum }_{j\in S_{t-1}}{k}_p(X_i,X_j)$, where $S_{t-1}$ is the sample obtained from the $t-1$ previous iterations.

Third, its performance seems to depend in a non-trivial way on the original kernel function k; the authors of [5] propose several strategies for choosing and scaling k, but none of them seem to perform uniformly well in their numerical experiments.

We propose a different approach in this paper, which we call cube thinning, and which addresses these shortcomings to some extent. Assuming the availability of J control variates (that is, of functions $h_j$ with known expectation under p), we cast the problem of MCMC compression as that of resampling the initial chain under constraints based on these control variates. The main advantage of cube thinning is that its complexity is $\mathcal{O}\left(N{J}^3\right)$; in particular, it does not depend on M. That makes it possible to use it for much larger values of M. We shall discuss the choice of J, but, by and large, J should be of the same order as d, the dimension of the sampling space. The name stems from the cube method of [6], which plays a central part in our approach, as we explain in the body of the paper.

The availability of control variates may seem like a strong requirement. However, if we assume we are able to compute $s_p\left(x\right)=\nabla logp\left(x\right)$, then (for a large class of functions $\varphi :{\mathbb{R}}^d\to {\mathbb{R}}^d$, which we define later)

$${\mathbb{E}}_p\left[\varphi \left(x\right)s_p\left(x\right)+{\nabla }_x·\varphi \left(x\right)\right]=0$$

where ${\nabla }_x·\varphi$ denotes the divergence of $\varphi$. In other words, the availability of the score function implies, automatically, the availability of control variates. The converse is not true: there exists control variates, e.g., [7], that are not gradient-based. One of the examples we consider in our numerical examples feature such non gradient-based control variates; as a result, we are able to apply cube thinning, although Stein thinning is not applicable.

The supporting methods of [4] do not require control variates. It is thus more generally applicable than either cube thinning or Stein thinning. On the other hand, when gradients (and thus control variates) are available, the numerical experiments of [5] suggest that Stein thinning outperforms support points. From now on, we focus on situations where control variates are available.

This paper is organised as follows. Section 2 recalls the concept of control variates, and explains how control variates may be used to reweight an MCMC sample. Section 3 describes the cube method of [6]. Section 4 explains how to combine control variates and the cube method to perform cube thinning. Section 5 assesses the statistical performance of cube thinning through two numerical experiments.

We use the following notations throughout: p denotes both the target distribution and its probability density; $p\left(f\right)$ is a short-hand for the expectation of $f\left(X\right)$ under p. The gradient of a function f is denoted by ${\nabla }_{x}f\left(x\right)$, or simply $\nabla f\left(x\right)$ when there is no ambiguity. The $i$-th component of a vector $v\in {\mathbb{R}}^d$ is denoted by $v\left[i\right]$, and it is transposed by $v^t$. The vectors of the canonical basis of ${\mathbb{R}}^d$ are denoted by $e_i$, i.e., $e_i\left[j\right]=1$ if $j=i$, 0 otherwise. Matrices are written in upper-case; the kernel (null space) of matrix A is denoted by $\ker A$. The set of functions $f:\mathsf{\Omega }\to {\mathbb{R}}^d$ that are continuously differentiable is denoted by $C^1(\mathsf{\Omega },{\mathbb{R}}^d)$.

2. Control Variates

2.1. Definition

Control variates are a very well known way to reduce the variance of Monte Carlo estimates—see, e.g., the books of [1,8,9].

Suppose we want to estimate the quantity $p\left(f\right)={\mathbb{E}}_p\left[f\left(X\right)\right]$ for a suitable $f:{\mathbb{R}}^d\to \mathbb{R}$, based on an IID (independent and identically distributed) sample $\{X_1,\cdots ,X_N\}$ from distribution p. The generalisation of control variates to MCMC will be discussed in Section 4.

The usual Monte Carlo estimator of $p\left(f\right)$ is

$$\widehat{p}\left(f\right)={\displaystyle \frac{1}{N}}\sum _{n=1}^{N}f\left(X_n\right).$$

(2)

Assume we know $J\in {\mathbb{N}}^{\star }$ functions $h_j:{\mathbb{R}}^d\to \mathbb{R}$ for $j\in \{1,\cdots ,J\}$ such that $p\left(h_j\right)=0$. Functions with this property are called control variates. We can use this property to build an estimate with a lower variance: let us denote $h\left(X\right)={(h_1\left(X\right),\cdots ,h_J\left(X\right))}^t$ and write our new estimate:

$${\widehat{p}}_{\beta }\left(f\right)={\displaystyle \frac{1}{N}}\sum _{n=1}^{N}f\left(X_n\right)+{\beta }^{t}h\left(X_n\right)$$

(3)

with $\beta \in {\mathbb{R}}^J$. Then it is straightforward to show that $\mathbb{E}\left[{\widehat{p}}_{\beta }\left(f\right)\right]=\mathbb{E}\left[\widehat{p}\left(f\right)\right]=p\left(f\right)$. Depending on the choice of $\beta$, we may have $\mathrm{Var}\left[{\widehat{p}}_{\beta }\left(f\right)\right]\le \mathrm{Var}\left[\widehat{p}\left(f\right)\right]$. The next section discusses how to choose such a $\beta$.

2.2. Control Variates as a Weighting Scheme

The standard approach to choose $\beta$ consists of two steps. First, one shows easily that the value the minimises the variance of estimator (3) is:

$${\beta }^{\star }\left(f\right)=\mathrm{Var}{\left(h\left(X\right)\right)}^{-1}\mathrm{Cov}(h\left(X\right),f\left(X\right))$$

(4)

where $\mathrm{Var}\left(h\right(X\left)\right)$ is the $J\times J$ variance matrix of the vector $h\left(X\right)$ and $\mathrm{Cov}\left(h\right(X),f(X\left)\right)$ is the $J\times 1$ vector such that $\mathrm{Cov}{(h\left(X\right),f\left(X\right))}_{i,1}=\mathrm{Cov}(f\left(X\right),h_i\left(X\right))$.

Second, one realises that this quantity may be estimated from the sample $X_1,\dots ,X_N$ through a simple linear regression model, where the $f\left(X_n\right)$s are the outcome, and the $h_j\left(X_n\right)$s are the predictors:

$$f\left(X_n\right)\approx \mu +{\beta }^{t}h\left(X_n\right)+{ϵ}_n,\mathbb{E}\left[{ϵ}_n\right]=0.$$

(5)

More precisely, let $\gamma \in R^{J+1}$ be the vector such that ${\gamma }^t=(\mu ,{\beta }^t)$, $H=\left(H_{ij}\right)$ the design matrix such that $H_{i1}=1$, $H_{i(j+1)}=h_j\left(X_i\right)$, and $F=(f\left(X_1\right),\dots ,f\left(X_N\right))$. Then, the OLS (ordinary least squares) estimate of $\gamma$ is

$${\widehat{\gamma }}_{\mathrm{OLS}}={\left(H^{t}H\right)}^{-1}{H}^{t}F.$$

(6)

Since $\mathbb{E}\left[f\left(X_n\right)\right]=\mu$ in this artificial regression model, the first component of ${\widehat{\gamma }}_{\mathrm{OLS}}$:

$${\widehat{p}}_{\star }\left(f\right):={\widehat{\gamma }}_{\mathrm{OLS}}\times e_1,$$

(7)

actually corresponds to estimate (3) when $\beta ={\widehat{\beta }}_{\mathrm{OLS}}$.

At first glance, the approach described above seems to require implementing a different linear regression for each function f of interest. Ref. [9] noted, however, that one may re-express (7) as a weighted average:

$${\widehat{p}}_{\star }\left(f\right)=\sum _{n=1}^N{w}_{n}f\left(X_n\right)$$

(8)

where the weights $w_n$ sum to one, and do not depend on f. It is thus possible to compute these weights once from a given sample (given a certain choice of control variates), and then quickly compute ${\widehat{p}}_{\star }\left(f\right)$ for any function f of interest.

The exact expression of the weights are easily deduced from (7) and (6): $w=\left(w_n\right)$ with

$$w=H{\left(H^{t}H\right)}^{-1}{e}_1.$$

2.3. Gradient-Based Control Variates

In this section and the next, we recall generic methods to construct control variates. This section specifically considers control variates that are derived from the score function, $s_p\left(x\right)=\nabla logp\left(x\right)$. (We therefore assume that this quantity is tractable.)

Under the following two conditions:

• The probability density $p\in C^1(\mathsf{\Omega },\mathbb{R})$ where $\mathsf{\Omega }\subseteq {\mathbb{R}}^d$ is an open set;
• Function $\varphi \in C^1(\mathsf{\Omega },{\mathbb{R}}^d)$ is such that ${\oint }_{\partial \mathsf{\Omega }}p\left(x\right)\varphi \left(x\right)·n\left(x\right)S\left(dx\right)=0$ where ${\oint }_{\partial \mathsf{\Omega }}$ denotes the integral over the boundary of $\mathsf{\Omega }$, and $S\left(dx\right)$ is the surface element at $x\in \partial \mathsf{\Omega }$.

The following function:

$$h\left(x\right)={\nabla }_x·\varphi \left(x\right)+\varphi \left(x\right)·s_p\left(x\right)$$

(9)

is a control variate: $p\left(h\right)=0$, see, e.g., [10] or [11] for further details. To gain some insight, note that in dimension 1 and assuming the domain of integration is an interval $]a,b[\subset \mathbb{R}$, this amounts to an integration by parts with the condition that $h\left(b\right)p\left(b\right)-h\left(a\right)p\left(a\right)=0$.

Thus, whenever the score function is available (and the conditions above hold), we are able to construct an infinite number of control variates (one for each function $\varphi$). For simplicity, we shall focus on the following standard classes of such functions. First, for $i=1,\dots ,d$,

$$\begin{array}{cc}\hfill {\varphi }_i:{\mathbb{R}}^d& \to {\mathbb{R}}^d\hfill \\ \hfill x& \mapsto e_i\hfill \end{array}$$

which leads to the following d control variates:

$$h_i\left(x\right)=s_p\left(x\right)\left[i\right].$$

(10)

For a Gaussian target, $N(\mu ,\mathsf{\Sigma })$, the score is $s_p\left(x\right)=-{\mathsf{\Sigma }}^{-1}(x-\mu )$, and the control variates above make it possible to reweight the Monte Carlo sample to make it have the same expectation as the target distribution.

Second, we consider, for $i,j=1,\dots ,d$:

$$\begin{array}{cc}\hfill {\varphi }_{ij}:{\mathbb{R}}^d& \to {\mathbb{R}}^d\hfill \\ \hfill x& \mapsto x\left[i\right]e_j\hfill \end{array}$$

which leads to the following $d^2$ control variates:

$$h_{ij}\left(x\right)=\mathbb{1}\{i=j\}+x\left[i\right]s_p\left(x\right)\left[j\right].$$

(11)

Again, for a Gaussian target $N(\mu ,\mathsf{\Sigma })$, this makes it possible to fix the empirical covariance matrix to true covariance $\mathsf{\Sigma }$.

In our simulations, we consider two sets of control variates: the ‘full’ set, consisting of the d control variates defined by (10), and the $d^2$ control variates defined by (11), and a ‘diagonal’ set of $2d$ control variates, where for (11), we only consider the cases where $i=j$. Of course, the former set should lead to a better performance (lower variance), but since the complexity of our approach will be $\mathcal{O}\left(J^3\right)$, where J is the number of control variates, taking $J=\mathcal{O}\left(d^2\right)$ may be too expensive whenever the dimension d is large. In fact, when d is very large, one might even consider considering only control variates that depend on a few components of x of interest.

2.4. MCMC-Based Control Variates

We mention in passing other ways to construct control variates, in particular in the context of MCMC.

For instance, [7] noted that, for a Markov chain $\left\{X_n\right\}$, the quantity

$$\varphi \left(X_n\right)-\mathbb{E}\left[\varphi \left(X_n\right)|X_{n=1}\right]$$

has zero expectations. In particular, if the MCMC kernel is a Gibbs sampler, it is likely that one is able to compute the conditional expectation of each component, i.e., $\varphi \left(x\right)=x\left[i\right]$ for $i=1,\dots ,d$.

See also [12,13] for other ways to construct control variates when the $X_n$s are simulated from a Metropolis kernel.

3. The Cube Method

We review in this section the cube method of [6]. This method originated from survey sampling and is a way to sample from a finite population under constraints. The first subsection gives some definitions, the second one explains the flight phase of the cube method and the third subsection discusses the landing phase of the method.

3.1. Definitions

Suppose we have a finite population $\{1,\cdots ,N\}$ of N individuals and that to each individual $n=1,\dots ,N$ is associated a variable of interest $y_n$ and J auxiliary variables, $v_n=(v_{n1},\cdots ,v_{nJ})$. Without loss of generality, suppose also that the J vectors $(v_{1j},\cdots ,v_{Nj})$ are linearly independent. We are interested in estimating the quantity $Y={\sum }_{n=1}^N{y}_n$ using a subsample of $\{1,\dots ,N\}$. Furthermore, we know the exact value of each sum $V_j={\sum }_{n=1}^N{v}_{nj}$, and we wish to use this auxiliary information to better estimate Y.

We assign, to each individual n, a sampling probability ${\pi }_n\in [0,1]$. We consider random variables $S_n$ such that, marginally, $\mathbb{P}(S_n=1)={\pi }_n$. We may then define the Horvitz–Thompson estimator of Y as

$$\widehat{Y}=\sum _{n=1}^N\frac{S_n{y}_n}{{\pi }_n}$$

(12)

which is unbiased, and which depends only on selected individuals (i.e., $S_n=1$).

We define similarly the Horvitz–Thompson estimator of $V_j$ as

$${\widehat{V}}_j=\sum _{n=1}^N\frac{S_n{v}_{nj}}{{\pi }_n}.$$

(13)

Our objective is to construct a joint distribution $\xi$ for the inclusion variables $S_n$ such that ${\mathbb{P}}_{\xi }(S_n=1)={\pi }_n$ for all $n=1,\dots ,N$, and

$$\widehat{V}=V\xi -\mathrm{almost} \mathrm{surely}.$$

(14)

where $V=(V_1,\dots ,V_J)$, $\widehat{V}=({\widehat{V}}_1,\dots ,{\widehat{V}}_J)$. Such a probability distribution is called a balanced sampling design.

3.2. Subsamples as Vertices

We can view all the possible samples from $\{1,\dots ,N\}$ as the vertices of the hypercube $\mathcal{C}={[0,1]}^N$ in ${\mathbb{R}}^N$. A sampling design with inclusion probabilities ${\pi }_n={\mathbb{P}}_{\xi }(S_n=1)$ is then a distribution over the set of these vertices such that $\mathbb{E}\left[S\right]=\pi$, where $S={(S_1,\dots ,S_N)}^t$, and $\pi ={({\pi }_1,\dots ,{\pi }_N)}^t$ is the vector of inclusion probabilities. Hence, $\pi$ is expressed as a convex combination of the vertices of the hypercube.

We can think of a sampling algorithm as finding a way to reach any vertex of the cube, starting at $\pi$, while satisfying the balancing Equation (14). However, before we describe such a sampling algorithm, we may wonder if it is possible to find a vertex such that (14) is satisfied.

3.3. Existence of a Solution

The balancing equation, Equation (14), defines a linear system. Indeed, we can re-express (14) as S, as a solution to $As=V$, where $A=\left(A_{jn}\right)$ is of dimension $J\times N$, $A_{jn}=v_{kn}/{\pi }_n$. This system defines a hyperplane Q of dimension $N-J$ in ${\mathbb{R}}^N$.

What we want is to find vertices of the hypercube $\mathcal{C}$ that also belong to the hyperplane Q. Unfortunately, it is not necessarily possible, as it depends on how the hyperplane Q intersects cube $\mathcal{C}$. In addition, there is no way to know beforehand if such a vertex exists. Since $\pi \in Q$, we know that $\mathcal{K}:=\mathcal{C}\cap Q\ne \varnothing$ and is of dimension $N-J$. The only thing we can say is stated Proposition 1 in [6]: if r is a vertex of $\mathcal{K}$, then in general $q=\mathrm{card}\left(\right\{n:0<r\left[n\right]<1\left\}\right)\le J$.

The next section describes the flight phase of the cube algorithm, which generates a vertex in $\mathcal{K}$ when such vertices exist, or which, alternatively, returns a point in $\mathcal{K}$ with most (but not all) components set to zero or one. In the latter case, one needs to implement a landing phase, which is discussed in Section 3.5.

3.4. Flight Phase

The flight phases simulates a process $\pi \left(t\right)$ which takes values in $\mathcal{K}=\mathcal{C}\cap Q$, and starts at $\pi \left(0\right)=\pi$. At every time t, one selects a unit vector $u\left(t\right)$, then one chooses randomly between one of the two points that are in the intersection of the hypercube $\mathcal{C}$ and the line parallel to $u\left(t\right)$ that passes through $\pi (t-1)$. The probability of selecting these two points are set to ensure that $\pi \left(t\right)$ is a martingale; in that way, we have $\mathbb{E}\left[{\pi }_t\right]=\pi$ at every time step. The random direction $u\left(t\right)$ must be generated to fulfil the following two requirements: (a) that the two points are in Q, i.e., $u\left(t\right)\in \ker A$, and (b) whenever $\pi \left(t\right)$ reaches one of the faces of the hypercube, it must stay within that face; thus, $u\left(t\right)\left[k\right]=0$ if $\pi (t-1)\left[k\right]=0$ or 1.

Algorithm 1 describes one step of the flight phase.

Algorithm 1: Flight phase iteration

Input: $\pi (t-1)$  Output: $\pi \left(t\right)$ 1 Sample $u\left(t\right)$ in $kerA$ with $u_k\left(t\right)=0$ if the k-th component of $\pi (t-1)$ is an integer. 2 Compute ${\lambda }_1^{\star }$ and ${\lambda }_2^{\star }$, the largest values of ${\lambda }_1>0$ and ${\lambda }_2>0$ such that: $0\le \pi (t-1)+{\lambda }_{1}u\left(t\right)\le 1$ and $0\le \pi (t-1)-{\lambda }_{2}u\left(t\right)\le 1$. 3 With probability ${\lambda }_2^{\star }/({\lambda }_1^{\star }+{\lambda }_2^{\star })$, set $\pi \left(t\right)\leftarrow \pi (t-1)+{\lambda }_{1}u\left(t\right)$; otherwise, set $\pi \left(t\right)\leftarrow \pi (t-1)-{\lambda }_{2}u\left(t\right)$.

The flight phase stops when Step 1 of Algorithm 1 cannot be performed (i.e., no vector $u\left(t\right)$ fulfils these conditions). Until this happens, each iteration increases by at least one the number of components in $\pi \left(t\right)$ that are either zero or one. Thus, the flight phases completes at most in N steps.

In practice, to generate $u\left(t\right)$, one may proceed as follows: first generate a random vector $v\left(t\right)\in {\mathbb{R}}^N$, then project it in the constraint hyperplane: $u\left(t\right)=I\left(t\right)v\left(t\right)-I\left(t\right)A^t{\left(AI\left(t\right)A^t\right)}^{-}AI\left(t\right)v\left(t\right)$, where $I\left(t\right)$ is a diagonal matrix such that $I_{kk}\left(t\right)$ is 0 if ${\pi }_k\left(t\right)$ is an integer and 1 otherwise, and $M^{-}$ denotes the pseudo-inverse of the matrix M.

The authors of [14] propose a particular method to generate vector $v\left(t\right)$, which ensures that the complexity of a single iteration of the flight phase is $\mathcal{O}\left(J^3\right)$. This leads to an overall complexity of $\mathcal{O}\left(N{J}^3\right)$ for the flight phase, since it terminates in at most N iterations.

3.5. Landing Phase

Denote by ${\pi }^{\star }$ the value of process $\pi \left(t\right)$ when the flight phase terminates. If ${\pi }^{\star }$ is a vertex of $\mathcal{C}$ (i.e., all its components are either zero or one), one may stop and return ${\pi }^{\star }$ as the output of the cube algorithm. If ${\pi }^{\star }$ is not a vertex, this informs us that no vertex belongs to $\mathcal{K}$. One may implement a landing phase, which aims at choosing randomly a vertex which is close to ${\pi }^{\star }$, and such that the variance of the components of $\widehat{V}$ is small.

Appendix A gives more details on the landing phase. Note that its worst-case complexity is $\mathcal{O}\left(2^J\right)$. However, in practice, it is typically either much faster, or not required (i.e., ${\pi }^{\star }$ is already a vertex) as soon as $J\ll N$.

4. Cube Thinning

We now explain how the previous ingredients (control variates, and the cube method) may be combined in order to thin a Markov chain, $X_1,\dots ,X_N$, into a subsample of size M. As before, the invariant distribution of the chain is denoted by p, and we assume we know of J control variates $h_j$, i.e., $p\left(h_j\right)=0$ for $j=1,\dots ,J$.

4.1. First Step: Computing the Weights

The first step of our method is to use the J control variates to compute the N weights $w_n$, as defined at the end of Section 2.2. Recall that these weights sum to one, and that they automatically fulfil the constraints:

$$\sum _{n=1}^N{w}_n{h}_j\left(X_n\right)=0$$

(15)

for $j=1,\dots ,J$, and that we use them to compute

$${\widehat{p}}_{\star }\left(f\right)=\sum _{n=1}^N{w}_{n}f\left(X_n\right)$$

(16)

as a low-variance estimate for $p\left(f\right)$ for any f.

Recall that the control variates procedure we described in Section 2 assume that the input variables, $X_1,\dots ,X_N$, are IID. This is obviously not the case in an MCMC context; however, we follow the common practice [10,11] of applying the procedure to MCMC points as if they were IID points. This implies that the weighted estimate above corresponds to a value of $\beta$ in (3) that does not minimise the (asymptotic) variance of estimator (3). It is actually possible to estimate the value of $\beta$ that minimises the asymptotic variance of an MCMC estimate [7,15]. However, this type of approach is specific to certain MCMC samplers, and, critically for us, it cannot be cast as a weighting scheme. Thus, we stick to this standard approach.

We note in passing that, in our experiments (see Figure 1 and the surrounding discussion), the weights $w_n$ make it easy to visually assess the convergence (and thus the burn-in) of the Markov chain. In fact, since the MCMC points of the burn-in phase are far from the mass of the target distribution, the procedure must assign a small or negative weight to these points in order to respect the constraints based on the control variates. Again, see Section 5.2 for more discussion on this issue. The fact that control variates may be used to assess MCMC convergence has been known for a long time (e.g., [16]), but the visualisation of weights makes this idea more expedient.

4.2. Second Step: Cube Resampling

The second step consists in resampling the weighted sample ${(w_n,X_n)}_{n=1,\dots ,N}$, to obtain a subsample $\mathcal{S}=\{X_n:S_n=1\}$ where $S_n$ are random variables such that (a) $\mathbb{E}\left[S_n\right]=w_n$; (b) ${\sum }_{n=1}^N{S}_n=M$, and (c) for $j=1,\dots ,J$:

$$\sum _{S_n=1}{h}_j\left(X_n\right)=0.$$

Condition (a) ensures that the procedure does not introduce any bias:

$$\mathbb{E}\left[\frac{1}{M}\sum _{S_n=1}f\left(X_n\right)|X_{1:N}\right]=\sum _{n=1}^N{w}_{n}f\left(X_n\right).$$

Condition (b) ensures that the subsample is exactly of size M.

We would like to use the cube method in order to generate the $S_n$’s. Specifically, we would like to assign the inclusion probabilities ${\pi }_n$ to $w_n$, and impose the $(J+1)$ constraints defined above by conditions (b) and (c). There is one caveat, however: the weights $w_n$ do not necessarily lie in $[0,1]$.

4.3. Dealing with Weights Outside of $[0,1]$

We rewrite (16) as:

$${\widehat{p}}_{\star }\left(f\right)={\displaystyle \frac{\mathsf{\Omega }}{M}}\times \sum _{n=1}^N{W}_n\times \mathrm{sgn}\left(w_n\right)f\left(X_n\right)$$

(17)

where $\mathsf{\Omega }={\sum }_{n=1}^N\left|w_n\right|$ and $W_n=M\left|w_n\right|/\mathsf{\Omega }$. We now have $W_n\ge 0$, and ${\sum }_{n=1}^N{W}_n=M$, which is required for condition (b) in the previous section. We might have a few points such that $W_n>1$. In that case, we replace them by $\lfloor W_n\rfloor$ copies, with adjusted weights $W_n/\lfloor W_n\rfloor$.

It then becomes possible to implement the cube method, using as inclusion probabilities the $W_n$s, and as the matrix A that defines the $J+1$ constraints, the matrix $A=\left(A_{jn}\right)$ such that $A_{1n}=1$, $A_{(j+1)n}=\mathrm{sgn}\left(w_n\right)h_j\left(X_n\right)$. The cube method samples variables $S_n$, which may be used to compute the subsampled estimate

$$\widehat{\nu }\left(f\right)=\frac{\mathsf{\Omega }}{M}\sum _{S_n=1}\mathrm{sgn}\left(w_n\right)f\left(X_n\right).$$

(18)

More generally, in our numerical experiments, we shall evaluate to which extent the random signed measure:

$$\widehat{\nu }=\frac{\mathsf{\Omega }}{M}\sum _{S_n=1}\mathrm{sgn}\left(w_n\right){\delta }_{X_n}\left(\mathrm{d}x\right).$$

(19)

is a good approximation of the target distribution p.

5. Experiments

We consider two examples. The first example is taken from [5], and is used to compare cube thinning with KSD thinning. The second example illustrates cube thinning when used in conjunction with control variates that are not gradient-based. We also include standard thinning in our comparisons.

Note that there is little point in comparing these methods in terms of CPU cost, as KSD thinning is considerably slower than cube thinning and standard thinning whenever $M\gg 100$. (In one of our experiments, for $M=1000$, KSD took close to 7 h to run, while cube thinning with all the covariates took about 30 s.) Thus, our comparison will be in terms of statistical error, or, more precisely, in terms of how representative of p is the selected subsample.

In the following (in particular in the plots), “cubeFull” (resp. “cubeDiagonal”) will refer to our approach based on the full (resp. diagonal) set of control variates, as discussed in Section 2.3. “NoBurnin” means that burn-in has been discarded manually (hence, no burn-in in the inputs). Finally, “thinning” denotes the usual thinning approach, “SMPCOV”, “MED” and “SCLMED” are the same names used in [5] for KSD thinning, based on three different kernels.

To implement the cube method, we used R package BalancedSampling.

5.1. Evaluation Criteria

We could compare the three different methods in terms of variance of the estimates of $p\left(f\right)$ for certain functions f. However, it is easy to pick functions f that are strongly correlated with the chosen control variates; this would bias the comparison in favour of our approach. In fact, as soon as the target is Gaussian-like, the control variates we chose in Section 2.3 should be strongly correlated with the expectation of any polynomial function of order two, as we discussed in that section.

Rather, we consider criteria that are indicative of the performance of the methods for a general class of function. Specifically, we consider three such criteria. The first one is the kernel Stein discrepency (KSD) as defined in [5] and recalled in the introduction—see (1). Note that this criterion is particularly favourable for KSD thinning, since this approach specifically minimises this quantity. (We use the particular version based on the median kernel in Riabiz et al. [5].)

The second criterion is the energy distance (ED) between p and the empirical distribution defined by the thinning method, e.g., (19) for cube thinning. Recall that the ED between two distributions F and G is:

$$ED(F,G)=2\mathbb{E}\left|\right|Z-{X\left|\right|}_2-\mathbb{E}\left|\right|Z-Z^{\prime }{\left|\right|}_2-\mathbb{E}\left|\right|X-X^{\prime }{\left|\right|}_2$$

(20)

where $Z^{\prime },Z\stackrel{iid}{\sim }F$ and $X^{\prime },X\stackrel{iid}{\sim }G$, and that this quantity is actually a pseudo-distance: $ED(F,G)\ge 0$, $ED(F,G)=0\Rightarrow F=G$, $ED(F,G)=ED(G,F)$, but ED does not fulfil the triangle inequality [17,18].

One technical difficulty is that (19) is a signed measure, not a probability measure; see Appendix B on how we dealt with this issue.

Our third criterion is inspired by the star discrepancy, a well-known measure of the uniformity of N points $u_n\in {[0,1]}^d$ in the context of quasi-Monte Carlo sampling [9] (Chapter 15). Specifically, we consider the quantity

$$d^{\star }(\widehat{P},\widehat{\nu })=\underset{B\in \mathcal{B}}{sup}\left|{\widehat{P}}_{\psi }\left(B\right)-{\widehat{\nu }}_{\psi }\left(B\right)\right|$$

where $\psi :{\mathbb{R}}^d\to {[0,1]}^d$, ${\widehat{P}}_{\psi }$ and ${\widehat{\nu }}_{\psi }$ are the push-forward measures associated to empirical distributions $\widehat{P}={(N-b)}^{-1}{\sum }_{n=b+1}^N{\delta }_{X_n}\left(dx\right)$, and $\widehat{\nu }$ as defined in (19), and $\mathcal{B}$ is the set of hyper-rectangles $B={\prod }_{i=1}^d[0,b_i]$. In practice, we defined function $\psi$ as follows: we apply the linear transform that makes the considered sample to have zero mean and unit variance, and then we applied the inverse CDF (cumulative distribution function) of a unit Gaussian to each component.

Additionally, since the sup above is not tractable, we replace it by a maximum over a finite number of $b_i$ (simulated uniformly).

5.2. Lotka–Volterra Model

This example is taken from [5]. The Lotka–Volterra model describes the evolution of a prey–predator system in a closed environment. We denote the number of prey by $u_1$ and the number of predators by $u_2$. The growth rate of the prey is controlled by a parameter ${\theta }_1>0$ and its death rate—due to the interactions with the predators—is controlled by a parameter ${\theta }_2>0$. In the same way, the predator population has a death rate of ${\theta }_3>0$ and a growth rate of ${\theta }_4>0$. Given these parameters, the evolution of the system is described by a system of ODEs:

$$\begin{array}{cc}\hfill {\displaystyle \frac{d{u}_1}{dt}}=& {\theta }_1{u}_1-{\theta }_2{u}_1{u}_2\\ \hfill {\displaystyle \frac{d{u}_2}{dt}}=& {\theta }_4{u}_1{u}_2-{\theta }_3{u}_2\end{array}$$

Ref. [5] set $\theta =({\theta }_1,{\theta }_2,{\theta }_3,{\theta }_4)=(0.67,1.33,1,1)$, the initial condition $u_0=(1,1)$, and simulate synthetic data. They assume they observe the populations of prey and predator at times $t_i,i=1,\cdots ,2400$ where the $t_i$ are taken uniformly on $[0,25]$ and that these observations are corrupted with a centered Gaussian noise with a covariance matrix $C=\mathrm{diag}({0.2}^2,{0.2}^2)$. Finally, the model is parametrised in terms of $x=(log{\theta }_1,log{\theta }_2,log{\theta }_3,log{\theta }_4)\in {\mathbb{R}}^4$ and a standard normal distribution as a prior on x is used.

The authors have provided their code as well as the sampled values they obtained by running different MCMC chains for a long time. We use the exact same experimental set-up, and we do not run any MCMC chain on our own, but use the ones they provide instead, specifically the simulated chain, of length $2\times {10}^6$, from preconditionned MALA.

We compress this chain into a subsample of size either $M=100$ or $M=1000$. For each value of M, we run different variations of our cube method 50 times and make a comparison with the usual thinning method and with the KSD thinning method with different kernels, see [5]. In Figure 1, we show the first 5000 weights of the cube method. We can see that after 1000 iterations, the weights seem to stabilise. Based on visual examination of these weights, we chose a conservative burn-in period of 2000 iterations for the variants where burn-in is removed manually.

We plot the results of the experiment in Figure 2, Figure 3 and Figure 4.

First, we see that regarding the kernel Stein discrepancy metric, Figure 2, the KSD method performs better than the standard thinning procedure and the cube method. This is not surprising since, even if this method does not properly minimise the Kernel–Stein Discrepency, this is still its target. We also see that, for $M=1000$, the KSD method performs a bit better than our cube method which in turn performs better than the standard thinning procedure. Note that the relative performance of the KSD method to our cube methods depends on the kernel that is being used and that there is no way to determine which kernel will perform best before running any experiment.

The picture is different for $M=100$: KSD thinning outperforms standard thinning, which in turn outperforms all of our cube thinning variations. Once again, the fact that the KSD method performs better than any other method seems reasonable: since it regards minimizing the Kernel–Stein Discrepancy, the KSD method is “playing at home” on this metric.

If we look at Figure 4, we see that all of our cube methods outperform the KSD method with any kernel. Interestingly, the standard thinning methods has a similar energy distance as the cube methods with “diagonal” control variates. These observations are true for both $M=100$ and $M=1000$. We can also note that the cube method with the full set of control variates tends to perform much better than its “diagonal” counterpart, whatever the value of M.

Finally, looking at Figure 3, it is clear that the KSD method—with any kernel—performs worse than any cube method in terms of star discrepancy.

Overall, the relative performance of the cube methods and KSD methods can change a lot depending on the metric being used and the number of points we keep. In addition, while all the cube methods tend to perform roughly the same, this is not the case of the KSD method, whose performances depend on the kernel we use. Unfortunately, we have no way to determine beforehand which kernel will perform best. This is a problem since the KSD method is computationally expensive for subsamples of cardinality $M\gg 100$.

Thus, by and large, cube thinning seems much more convenient to use (both in terms of CPU time and sensitivity to tuning parameters) while offering, roughly, the same level of statistical performance.

5.3. Truncated Normal

In this example, we use the (random-scan version of) the Gibbs sampler of [1] to sample from 10-dimensional multivariate normal truncated to ${[0,\infty )}^{10}$. We generated the parameters of this truncated normal as follows: the mean was set as the realisation of a 10-dimensional standard normal distribution, while for the covariance matrix $\mathsf{\Sigma }$, we first generated a matrix $M\in {\mathcal{M}}_{10,10}\left(\mathbb{R}\right)$ for which each entry was the realisation of a standard normal distribution. Then, we set $\mathsf{\Sigma }=M^{T}M$.

Since we used a Gibbs sampler, we have access to the Gibbs control variates of [7], based on the expectation of each update (which amounts to simulating from a univariate Gaussian). Thus, we consider 10 control variates.

The Gibbs sampler was run for $N={10}^5$ iterations and no burn-in was performed. We compare the following estimators of the expectation of the target distribution the standard estimator, based on the whole chain ("usualEstim" in the plots), the estimator based on standard thinning ("thinEstim" in the plots), the control variate estimator based on the whole chain, i.e., (7) ("regressionEstim" in the plots), and finally our cube estimator described in Section 4 ("cubeEstim" in the plots). For standard thinning and cube thinning, the thinning sample size was set to $M=100$, which corresponds to a compression factor of ${10}^3$.

The results are shown in Figure 5. First, we can see that the control variates we chose led to a substantial decrease in the variance of the estimates for regressionEstim compared to usualEstim. Second, the cube estimator performed worse than the regression estimator in terms of variance, but this was expected, as explained in Section 4. More interestingly, if we cannot say that the cube estimator performs better than the usual MCMC estimator in general, we can see that for some components it performed as well or even better, even though the cube estimator used only $M=100$ points while the usual estimator used ${10}^5$ points. This is largely due to the good choice of the control variates. Finally, the cube estimator outperformed the regular thinning estimator for every component, sometimes significantly.