Nested adaptation of MCMC algorithms

Markov chain Monte Carlo (MCMC) has become a widely used approach for simulation from an arbitrary distribution of interest, typically a Bayesian posterior distribution, known as the target distribution. MCMC really represents a family of sampling methods. Generally speaking, any new sampler that can be shown to preserve the ergodicity of the Markov chain such that it converges to the target distribution is a member of the family and can be combined with other samplers as part of a valid MCMC kernel. The key to MCMC’s success is its simplicity and applicability. In practice, however, it sometimes needs a lot of non-trivial tuning to work well (Haario et al., 2005).

To deal with this problem, many adaptive MCMC algorithms have been proposed (Gilks et al., 1998; Haario et al., 2001; Andrieu and Robert, 2001; Sahu et al., 2003). These allow parameters of the MCMC kernel to be automatically tuned based on previous samples. This breaks the Markovian property of the chain so has required special schemes and proofs that the resulting chain will converge to the target distribution (Andrieu and Moulines, 2003; Atchadé et al., 2005; Andrieu and Atchadé, 2006). Under some weaker and easily verifiable conditions, namely “diminishing adaptation” and “containment”, Roberts and Rosenthal (2007) proved ergodicity of adaptive MCMC and proposed many useful samplers.

It is important to realize, however, that such adaptive MCMC samplers address only a small aspect of a much larger problem. A typical adaptive MCMC sampler will approximately optimize performance given the kind of sampler chosen in the first place, but it will not optimize among the variety of samplers that could have been chosen. For example, an adaptive random walk Metropolis-Hastings sampler will adapt the scale of its proposal distribution, but that adaptation won’t reveal whether an altogether different kind of sampler would be more efficient. In many cases it would, and the exploration of different sampling strategies often remains a human-driven trial-and-error affair.

Here we present a method for a higher level of MCMC adaptation. The adaptation explores a potentially large space of valid MCMC kernels composed of different samplers. One starts with an arbitrary set of candidate samplers for each dimension or block of dimensions in the target distribution. The main idea is to iteratively try different candidates that compose a valid MCMC kernel, run them for a relatively short time, generate the next set of candidates based on the results thus far, and so on. Since relative performance of different samplers is specific to each model and even to each computing environment, it is doubtful whether there is a universally optimal kind of sampler. Hence we view the choice of efficient samplers for a particular problem as well-suited to empirical determination via computation.

The goal of computationally exploring valid sampler combinations in search of an efficient model-specific MCMC kernel raises a number of challenges. First, one must prove that the samples collected as the algorithm proceeds indeed converge to the target distribution, even when some of the candidate samplers are internally adaptive, such as conventional adaptive Metropolis-Hastings samplers. We provide such a proof for a general framework.

Second, one must determine efficient methods for exploring the very large, discrete space of valid sampler combinations. This is complicated by a combinatorial explosion, which is exacerbated by the fact that any multivariate samplers can potentially be used for arbitrary blocks of model dimensions. Here we take a practical approach to this problem, setting as our goal only to show basic schemes that can yield substantial improvements in useful time frames. Future work can aim to develop improvements within the general framework presented here. We also limit ourselves to relatively simple candidate samplers, but the framework can accommodate many more choices.

Third, one must determine how to measure the efficiency of a particular MCMC kernel for each dimension and for the entire model, in order to have a metric to seek to optimize. As a first step, it is vital to realize that there can be a trade-off between good mixing and computational speed. When considering adaptation within one kind of sampler, say adaptive Metropolis-Hastings, one can roughly assume that computational cost does not depend on the proposal scale, and hence mixing measured by integrated autocorrelation time, or the related effective sample size, is a sensible measure of efficiency. But when comparing two samplers with very different computational costs, say adaptive Metropolis-Hastings and slice samplers, good mixing may or may not be worth its computational cost. Metropolis-Hastings samplers may mix more slowly than slice samplers on a per iteration basis, but they do so at higher computational speed because slice samplers can require many evaluations of model density functions. Thus the greater number of random walk iterations per unit time could outperform the slice sampler. An additional issue is that different dimensions of the model may mix at different rates, and often the slowest-mixing dimensions limit the validity of all results (Turek et al., 2017). In view of these considerations, we define MCMC efficiency as the effective sample size per computation time and use that as a metric of performance per dimension. Performance of an MCMC kernel across all dimensions is defined as the minimum efficiency among all dimensions.

The rest of the paper is organized as follows. Section 2 begins with a general theoretical framework for Nested Adaptation MCMC, then extends these methods to a specific Nested Adaptation algorithm involving block MCMC updating. Section 3 presents an example algorithm that fits within the framework, and provides some explanations on its details. Section 4 then outlines some numerical examples comparing the example algorithm with existing algorithms for a variety of benchmark models. Finally, section 5 concludes and discusses some future research directions.

In this section, we present a general Nested Adaptation MCMC algorithm and give theoretical results establishing its correctness.

Let $\mathbf{\mathcal{X}}$ be a state space and $\pi$ the probability distribution on $\mathbf{\mathcal{X}}$ that we wish to sample from. Let $\mathbf{\mathcal{I}}$ be a countable set (this set indexes the discrete set of MCMC kernels we wish to choose from). For $\iota\in\mathbf{\mathcal{I}}$, let $\Theta_{\iota}$ be a parameter space in some subset space $\mathbb{R}^{m}$. For $\iota\in\mathbf{\mathcal{I}}$ and $\theta\in\Theta_{\iota}$, let $P_{\iota,\theta}$ denote a Markov kernel on $\mathbf{\mathcal{X}}$ with invariant distribution $\pi$. We set ${\displaystyle\bar{\Theta}=\bigcup_{\iota\in\mathbf{\mathcal{I}}}\{\iota\}\times\Theta_{\iota}}$ the adaptive MCMC parameter space. We want to build a stochastic process (an adaptive Markov chain) $\{(X_{n},\ \iota_{n},\ \theta_{n}),\ n\ \geq 0\}$ on $\mathbf{\mathcal{X}}\times\bar{\Theta}$ such that as $n\rightarrow\infty$, the distribution of $X_{n}$ converges to $\pi$, and a law of large numbers holds. We call $\iota$ the external adaptation parameter and $\theta$ the internal adaptation parameter.

We will follow the general adaptive MCMC recipe of Roberts and Rosenthal (2009). Assume that any internal adaptation on $\Theta_{\iota}$ is done using a function $H_{\iota}$ : $\Theta_{\iota}\times\mathbf{\mathcal{X}}\rightarrow\Theta_{\iota}$, and an “internal clock” sequence $\{\gamma_{n},\ n\ \geq 0\}$ such that $\lim_{n\rightarrow\infty}\gamma_{n}=0$. The function $H_{\iota}$ depends on $\gamma_{n}$ and updates parameters for internal adaptation. Also, let $\{p_{k},\ k\geq 1\}$ be a sequence of numbers $p_{k}\in(0,1)$ such that $\lim_{k\rightarrow\infty}p_{k}=0$. $p_{k}$ will be the probability of performing external adaptation at external iteration $k$. During the algorithm we will also keep track of two variables: $\kappa_{n}$, the number of external adaptations performed up to step $n$; and $\tau_{n}$, the number of iterations between $n$ and the last time an external adaptation is performed. These two variables are used to manage the internal clock based on external iterations, which in most situations can simply be the number of adaptation steps. We build the stochastic process $\{(X_{n},\ \iota_{n},\ \theta_{n}),\ n\ \geq 0\}$ on $\mathbf{\mathcal{X}}\times\bar{\Theta}$ as follows.

1. 1.

   We start with $\kappa_{0}=\tau_{0}=0$. We start also with some $X_{0}\in\mathbf{\mathcal{X}},\ \iota_{0}\in\mathbf{\mathcal{I}}$, and $\theta_{0}\in\Theta_{\iota_{0}}$.

2. 2.

   At the n-th iteration, given $\mathcal{F}_{n}\overset{def}{=}\sigma\{(X_{k},\ \iota_{k},\ \theta_{k}),\ k\leq n\}$, and given $\kappa_{n},\ \tau_{n}$:

   1. (a)

      Draw $X_{n+1}\sim P_{\iota_{n},\theta_{n}}(X_{n},\ \cdot)$.

   2. (b)

      Independently of $\mathcal{F}_{n}$ and $X_{n+1}$, draw $B_{n+1}\sim$ Bern$(p_{n+1})\in\{0,1\}$.

      1. i.

         If $B_{n+1}=0$, there is no external adaptation: $\iota_{n+1}=\iota_{n}$. We update $\kappa_{n}$ and $\tau_{n}$:

         $$\kappa_{n+1}=\kappa_{n},\ \tau_{n+1}=\tau_{n}+1.$$

         (2.1)

         Then we perform an internal adaptation: set $c_{n+1}=\kappa_{n+1}+\tau_{n+1}$, and compute

         $$\theta_{n+1}=\theta_{n}+\gamma_{c_{n+1}}H_{\iota_{n}}(\theta_{n},\ X_{n+1}).$$

         (2.2)

         Note that the internal adaptation interval could vary between iterations.

      2. ii.

         If $B_{n+1}=1$, then we do an external adaptation: we choose a new $\iota_{n+1}$. And we choose a new value $\theta_{n+1}\in\Theta_{\iota_{n+1}}$ based on $\mathcal{F}_{n}$ and $X_{n+1}$. Then we update $\kappa_{n}$ and $\tau_{n}$.

         $$\kappa_{n+1}=\kappa_{n}+1,\ \tau_{n+1}=0.$$

         (2.3)

For this Nested Adaptation MCMC algorithm to be valid we must show that it satisfies three assumptions:

1. 1.

   For each $(\iota,\ \theta)\in\bar{\Theta},\ P_{\iota,\theta}$ has invariant distribution $\pi$.

2. 2.

   (diminishing adaptation):

   $$\triangle_{n+1}\overset{{}_{def}}{=}\sup_{x\in\mathcal{X}}\|P_{\iota_{n},\theta_{n}}(x,\ \cdot)-P_{\iota_{n+1},\theta_{n+1}}(x,\ \cdot)\|_{\mathrm{TV}}$$

   converges in probability to zero, as $n\rightarrow\infty$,

3. 3.

   (containment): For all $\epsilon>0$, the sequence $\{M_{\epsilon}(\iota_{n},\ \theta_{n},\ X_{n})\}$ is bounded in probability, where

   $$M_{\epsilon}(\iota,\ \theta,\ x)\overset{{}_{def}}{=}\inf\{n\ \geq 1:\|P_{\iota,\theta}^{n}(x,\ \cdot)-\pi\|_{\mathrm{TV}}\leq\epsilon\}.$$

For $\iota\in\mathbf{\mathcal{I}},\ \theta,\ \theta^{\prime}\in\Theta_{\iota}$, define

We present one specific approach as an example of a Nested Adaptation algorithm. Our approach to “outer adaptation” will be to identify the “worst-mixing dimension” (i.e., some parameter or latent state of the statistical model) and update the kernel by assigning different sampler(s) for that dimension. To explain the method, we will give some terminology for describing our algorithm. In particular, we will define a valid kernel, MCMC efficiency, and worst-mixing dimension. We will define a set of candidate samplers for a given dimension, which could include scalar samplers or block samplers. In either case, a sampler may also have internal adaptation for each parameter or combination. To implement the internal clock of each sampler ($c_{n}$ of the general algorithm), we need to formulate all internal adaptation in the framework using equation 2.2. We use $P$ (without subscripts) in this section to represent $P_{\iota,\theta}$ of the general theory, so the kernel and parameters are implicit.

In this section, we evaluate our algorithm on some benchmark examples and compare them to different MCMC algorithms. In particular, we compare our approach to the following MCMC algorithms.

• •

  All Scalar algorithm: Every dimension is sampled using an adaptive scalar normal random walk sampler.

• •

  All Blocked algorithm: All dimensions are sampled in one adaptive multivariate normal random walk sampler.

• •

  Default algorithm: Samplers are assigned as follows. When possible, a conjugate (Gibbs) sampler is used. Otherwise, adaptive random-walk Metropolis-Hastings samplers are used. These will be block samplers for parameters arising from multivariate distributions and scalar samplers otherwise.

• •

  Auto Block algorithm: The Auto Block method (Turek et al., 2017) searches blocking schemes based on hierarchical clustering from posterior correlations to determine a highly efficient (but not necessarily optimal) set of blocks that are sampled with multivariate normal random-walk samplers. Thus, Auto Block uses only either scalar or multivariate adaptive random walk, concentrating more on partitioning the correlation matrix than trying different sampling methods. Note that the initial sampler of both the Auto Block algorithm and our proposed algorithm is the All Scalar algorithm.

All experiments were carried out using the NIMBLE package (de Valpine et al., 2017) for R (R Core Team, 2013) on a cluster using $32$ cores of Intel Xeon E5-2680 $2.7$ Ghz with $256$ GB memory. Models are coded using NIMBLE’s version of the BUGS model declaration language (Lunn et al., 2000, 2012). All MCMC algorithms are written in NIMBLE, which provides user-friendly interfaces in R and efficient execution in custom-generated C++, including matrix operations in the C++ Eigen library (Guennebaud et al., 2010).

To measure the performance of an MCMC algorithm, we use MCMC efficiency. MCMC efficiency depends on ESS, estimates of which can have a high variance for a short Markov chain. This presents a tuning-parameter trade-off for the Nested Adaptation method: Is it better to move cautiously (in sampler space) by running long chains for each outer adaptation in order to gain an accurate measure of efficiency, or is it better to move adventurously by running short chains, knowing that some algorithm decisions about samplers will be based on noisy efficiency comparisons? In the latter case, the final samplers may be less optimal, but that may be compensated by the saved computation time. To explore this trade-off, we try our Nested Adaptation algorithm with different sample sizes in each outer adaptation and label results accordingly. For example, Nested Adaptation 10K will refer to the Nested Adaptation method with samples of 10,000 per outer iteration.

We present algorithm comparisons in terms of time spent in an adaptation phase, final MCMC efficiency achieved, and the time required to obtain a fixed effective sample size (e.g., 10,000). Only Auto Block and Nested Adaptation have adaptation phases. An important difference is that Auto Block did not come with a proof of valid adaptive MCMC convergence (it could be modified to work in the current framework, but we compare to the published version). Therefore, samples from its adaptation phase are not normally included in the final samples, while the adaptation samples of Nested Adaptation can be included.

To measure final MCMC efficiency, we conducted a single long run of length $N$ with the final kernel of each method solely for the purpose of obtaining an accurate ESS estimate. One would not normally do such a run in a real application. The calculation of time to obtain a fixed effective sample size incorporates both adaptation time and efficiency of the final samplers. For both Nested Adaptation and Auto Block, we placed them on a similar playing field by assuming for this calculation that samples are not retained from the adaptation phase, making the results conservative.

For all comparisons, we used $20$ independent runs of each method and present the average results from these runs. To show the variation in runs, we present box-plots of efficiency in relation to computation time from the $20$ runs of Nested Adaptation. The final (right-most) box-plot in each such figures shows the $20$ final efficiency estimates from larger runs. Not surprisingly, these can be lower than obtained by shorter runs. These final estimates are reported in the tables.

A public Github repository containing scripts for reproducing our results may be found at https://github.com/nxdao2000/AutoAdaptMCMC. Some additional experiments are also provided there.

We have proposed a general Nested Adaptation MCMC algorithm. Our algorithm traverses a space of valid MCMC kernels to find an efficient algorithm automatically. There is only one previous approach, namely Auto Block sampling, of this kind that we are aware of. We have shown that our approach can substantially outperform Auto Block in some cases, and that both outperform simple static approaches. Using some benchmark models, we observe that our approach can yield improvements, which can be substantial compared to the Default method.

The comparisons presented have deliberately used fairly simple samplers as options for Nested Adaptation in order to avoid comparisons among vastly different computational implementations. A major feature of our framework is that it can incorporate almost any sampler as a candidate and almost any strategy for choosing new kernels from compositions of samplers based on results so far. Samplers to be explored in the future could include auxiliary variable algorithms such as slice sampling or derivative-based sampling algorithms such as Hamiltonian Monte Carlo (Duane et al., 1987). Now that the basic framework is established and shown to be useful in simple cases, it merits extension to more advanced cases.

The Nested Adaptation method can be viewed as a generalization of the Auto Block method. It is more general in the sense that it can use more kinds of samplers and explore the space of samplers more generally than the correlation-clustering of Auto Block. Thus, our framework can be considered to provide a broad class of automated kernel construction algorithms that use a wide range of sampling algorithm as components.

If block sampling is included in the space of the candidate samplers, choosing optimal blocks is important and can greatly increase the efficiency of the algorithm. For this reason, we extended the cutting of a hierarchical cluster tree to allow different cut heights on different branches (different parts of the model). This differs from Auto Block, which forms all blocks by cutting the entire tree at the same height. We also have different multivariate adaptive sampling other than random walk normal distribution such as automated factor slice sampler and automated factor random walk sampler. With these extensions, the final efficiency achieved by our algorithm specifically among blocking schemes is often substantially better and is found in a shorter time.

Beyond hierarchical clustering, there are other approaches one might consider to find efficient blocking schemes. One such approach would be to use the structure of the graph instead of posterior correlations to form blocks. This would allow conservation of calculations that are shared by some parts of the graph, whether or not they are correlated. Another future direction could be to improve how a new kernel is determined from previous results, essentially to determine an effective strategy for exploring the very high-dimensional kernel space. Finally, the trade-off between computational cost and the accuracy of effective sample size estimates is worth further exploration.

Development of Nested Adaptation MCMC has also raised several theoretical questions for future work. First, similarly to other adaptive MCMC methods, can more general conditions of validity be proven? Second, the estimation of effective sample size from a chain generated using Nested Adaptation is difficult because the properties of the kernel will have changed during the course of sampling. Methods for estimating effective sample size typically don’t consider such a situation. Third, what are good choices for the external and internal adaptation schedules (e.g. $\gamma_{n}$ and $p_{k}$ in Section 2)? Fourth, how can we disentangle the contribution of each sampler to mixing achieved by a kernel comprising multiple samplers? Doing so could enable better moves in kernel space during outer adaptation. Finally, what are good strategies for parallelization of Nested Adaptation?