Adaptive Metropolis-coupled MCMC for BEAST 2

Background

Metropolis-coupled MCMC makes use of running n different chains i = 1,…,n at different temperatures (Geyer, 1991; Gilks & Roberts, 1996; Altekar et al., 2004). Each of the different chains works similar to a regular MCMC chain. In regular MCMC, a parameter space is explored as follows: Given that the MCMC is currently at state x, we propose a new state x′ from a proposal distribution g(x′|x) given the current state. At this new state, we calculate the likelihood P(D|x′) of the data D given the state and the prior probability of the new state P(x′) and compare it the to old state. The probability of accepting this new state is then calculated as follows: (1) $$R=min[1,{\displaystyle \frac{P(D|x^{\mathrm{\prime }})P(x^{\mathrm{\prime }})}{P(D|x)P(x)}{\displaystyle \frac{g(x|x^{\mathrm{\prime }})}{g(x^{\mathrm{\prime }}|x)}]}}$$

If R is greater than a randomly drawn value between [0,1], the new state x′ is accepted as the current state, otherwise it is rejected and we remain in the same state. If we keep proposing new states x′ and accept these using Eq. (1), we eventually explore parameter space with the frequency at which values of a parameter are visited being its marginal probability (Geyer, 1991).

One of the issues of using this approach is that acceptance probabilities can be quite low, which makes it hard to move between different states in parameter space. Alternatively, an MCMC chain can be heated by using a temperature scaler ${\mathrm{\beta }}_i={\displaystyle \frac{1}{1+(i-1)\mathrm{\Delta }t}}$, with i being the number of the chain (Altekar et al., 2004). Heating of an MCMC chain changes its acceptance probability R_heated to: $$R_{\mathrm{h}\mathrm{e}\mathrm{a}\mathrm{t}\mathrm{e}\mathrm{d}}=\mathrm{m}\mathrm{i}\mathrm{n}[1,({\displaystyle \frac{P(D|x^{\mathrm{\prime }})P(x^{\mathrm{\prime }})}{P(D|x)P(x)}{)}^{{\mathrm{\beta }}_i}{\displaystyle \frac{g(x|x^{\mathrm{\prime }})}{g(x^{\mathrm{\prime }}|x)}]}}$$

For a heated chain, the frequency at which a value of a parameter is visited does not correspond to its marginal probability any more. However, heated chains can be used as a proposal to update the cold chain by performing what is essentially an MCMC move. This move proposes to swap the current states of two random chains i and j with the temperature β_i and β_j such that β_i < β_j. Exchanging the states of chains i and j is accepted with an acceptance probability R_ij of: $$R_{ij}=\mathrm{m}\mathrm{i}\mathrm{n}[1,{\displaystyle \frac{P{(x_i|D)}^{{\mathrm{\beta }}_j}P{(x_j|D)}^{{\mathrm{\beta }}_i}}{P{(x_i|D)}^{{\mathrm{\beta }}_i}P{(x_j|D)}^{{\mathrm{\beta }}_j}}]}$$

As for a regular MCMC move, swapping the states of the two chains is accepted when a randomly drawn uniformly distribution value in [0,1] is smaller than R_ij.

Additional to randomly swapping states between chains, we also implemented the possibility to only swap the states of neighbouring chains. That means that we condition on i = j + 1 instead of randomly sampling both i and j.

Locally aware adaptive tuning of the temperature of heated chains

Choosing an optimal temperature of the different heated chains can be a tedious task, requiring running an analysis, updating temperatures of the analysis and re-running everything. Instead, the temperatures of chains can be tuned automatically during the run itself to achieve a targeted average acceptance probability. Ideally, we would like to adjust the temperature such that effective sample size (ESS) of parameters of interest is maximised per unit of time, but ESSs are hard to estimate while running an analysis. Therefore, optimising for average acceptance probability balances the need for moving through the MCMC’s state space (at higher acceptance probability), and making bold moves (at lower acceptance probability), which are two requirements for getting good ESSs per unit of time. As stated above, we consider the temperatures difference between the n different chains to be a constant value ∆t, which we tune during the analysis.

When updating the temperature based on the global acceptance probability, we compute p_global based on all proposed exchanges of states from the start of a run to the current state. We then iteratively tune the temperature to achieve the target average acceptance probability p_target over the course of an analysis as follows. Given p_global and p_target, we update the difference in temperature between chains ∆t as follows: (2) $$\mathrm{\Delta }{t}_{\mathrm{n}\mathrm{e}\mathrm{w}}=\mathrm{m}\mathrm{a}\mathrm{x}[0,\mathrm{\Delta }{t}_{\mathrm{c}\mathrm{u}\mathrm{r}\mathrm{r}\mathrm{e}\mathrm{n}\mathrm{t}}+{\displaystyle \frac{p_{\mathrm{g}\mathrm{l}\mathrm{o}\mathrm{b}\mathrm{a}\mathrm{l}}-p_{\mathrm{t}\mathrm{a}\mathrm{r}\mathrm{g}\mathrm{e}\mathrm{t}}}{\mathrm{\#}\mathrm{e}\mathrm{x}\mathrm{c}\mathrm{h}\mathrm{a}\mathrm{n}\mathrm{g}\mathrm{e}\mathrm{s}}]}$$

With # exchanges denoting the total number of proposed exchanges, which increases throughout the BEAST run. This means that updating the temperature as in Eq. (2), leads the tuning of the temperature to become smaller and smaller and eventually approaches zero.

Tuning ∆t is only performed after an initial burn-in period of (by default) 100 proposed exchanges. By default, the target acceptance probability is set to 0.234, which for many MCMC proposals can be shown to be an optimal trade-off between as many accepted moves as possible and as large of a move as possible (Kone & Kofke, 2005; Atchadé, Roberts & Rosenthal, 2011). Datasets where unfavourable intermediate states are of particular issue may, however, require higher temperatures and therefore lower acceptance probabilities to overcome these intermediate states.

Changing the temperature of a heated chain changes the equilibrium distribution of that chain. There can be a significant time lag between changing the temperature of a chain and that chain moving to its new equilibrium state. If the temperature is updated too fast, heated chains may not have reached this new equilibrium yet which in turn can lead to over-adaptation. This is particularly problematic at the beginning of an analysis where $\sum \mathrm{e}\mathrm{x}\mathrm{c}\mathrm{h}\mathrm{a}\mathrm{n}\mathrm{g}\mathrm{e}\mathrm{s}$ is relatively small and where large changes in the temperature could occur. In order to reduce the risk of that, we maximise the difference between ∆t_current and ∆t_new, that is by how much the temperature can be changed, to be 0.001.

Another issue can arise when the global acceptance probability strongly differs from the current acceptance probability. In order to avoid that, we made the adaptation procedure aware of the local acceptance probability. To do so, we compute a local acceptance probability p_local of the last 100 proposed exchanges. We only update the temperature if the global and the local acceptance are on the same side of the target acceptance probability, that is if p_local > p_target & p_global > p_target or p_local < p_target & p_global < p_target.

Implementation

In this implementation of MC³, we run n different MCMC chains, with each chain i ∈ [1,…,n] running at a temperature ${\mathrm{\beta }}_i={\displaystyle \frac{1}{1+(i-1)\mathrm{\Delta }t}}$ (Altekar et al., 2004). We additionally implemented a scenario where the values for β_i are given by the quantiles of a beta distribution, such that ${\mathrm{\beta }}_i=1-\mathrm{c}\mathrm{d}\mathrm{f}({\displaystyle \frac{i-1}{nrchains})}$. With cdf being the cumulative density function of a beta distribution with α = 1 and β being the tuning parameter.

Upon initialisation, we first sample at random at which iteration the states of two chains with which number are proposed to be exchanged. We then initialise each chain to be run in its own Java thread using multiple CPU cores, if available. Each chain is then run until it reaches the time when an exchange of states with another chain will be proposed. This means than every chain runs independently of each other until an iteration at which it actually participates in a proposed exchange, minimising the crosstalk between threads (Altekar et al., 2004).

This is, however, only true for swapping between random chains. When restricting swaps to only occur between neighbouring chains, we run each chain until the next possible swap. We then randomly choose between which two chains, a swapping of states is proposed.

If the exchange of states between different chains is accepted, we exchange the temperature of the two chains instead of the states themselves (Altekar et al., 2004; Ronquist et al., 2012; Aberer, Kobert & Stamatakis, 2014; Höhna et al., 2016). The states can be quite large and exchanging them across different chains is potentially quite time consuming. Instead of exchanging the states themselves, we exchange the operators and loggers, which are the objects that produce the log files. Exchanging the operator specifications is done such that the individual tuning parameters of operators of a chain can be optimised to run at specific temperatures. The loggers are exchanged such that each heated chain logs its states to the log file that corresponds to its temperature and not the number of the chain.

The temperature is adapted at any potential exchange of states between chains, after an initial phase of 100 potential exchanges without any adaption. The temperature is updated simultaneously on all chains, not just the ones participating in the exchange of states, independent of which iterations they are in.

Adaptive MC³ is implemented, such that runs that were prematurely stopped or didn’t reach sufficient convergence yet can be resumed. Usually, a graphical user interface called BEAUti is used to set up BEAST 2 analyses. Setting up analyses with MC³ works differently depending on whether a BEAUTi template is needed to set up an analysis as required for some packages. If no such template is needed, an analysis can be set up to run with MC³ directly in BEAUTi and we provide a tutorial on how to do this on https://taming-the-beast.org/tutorials/CoupledMCMC-Tutorial/(Barido-Barido-Sottani et al., 2017). Alternatively, we provide an interface that converts BEAST 2 XMLs set up to run with MCMC into such that run with adaptive MC³.

Data availability and software

The BEAST 2 package coupledMCMC can be downloaded by using the package manager in BEAUti. The source code for the software package can be found here: https://github.com/nicfel/CoupledMCMC. The XML files used for the analysis performed here can be found in https://github.com/nicfel/CoupledMCMC-Material. All plots were done using ggplot2 (Wickham, 2016) in R (R Development Core Team, 2013).

Validation

Similar to the validation of MCMC operators, we can sample under the prior to validate the implementation of the MC³ approach. To do so, we sampled typed trees with five taxa and two different states under the structured coalescent using the MultiTypeTree (Vaughan et al., 2014) package for BEAST 2. We did this sampling once using MCMC and once using MC³. If the implementation of the MC³ algorithm explores the same parameter space as MCMC, marginal parameter distributions sampled using both approaches should be equal. In Fig. S1, we compare the distribution of different summary statistics of typed trees between MCMC and MC³, which shows both methods are in agreement.

Ergodicity of the adaptive Metropolis-coupled MCMC algorithm

First, we test if the distribution of posterior probability values using adaptive MC³ algorithm are consistent over time, that is ergodic. To do so, we ran 100 skyline (Drummond et al., 2005) analyses of Hepatitis C in Egypt (Ray et al., 2000), with three different target acceptance probabilities, 0.234 (Kone & Kofke, 2005; Atchadé, Roberts & Rosenthal, 2011), 0.468 (= 2 * 0.234) and 0.117 $(={\displaystyle \frac{0.234}{2})}$. The temperature difference between chains ∆t is being adapted during the analyses, particularly during the initial phase (see Fig. 1B).

We then computed the distribution of posterior probability estimates of the 100 different runs using the posterior probability estimates at different iterations. The distribution of posterior probability estimates stays constant over the different iterations (see Fig. 1A), despite the temperature difference between chains being adapted. This is true for all three different target acceptance probabilities.

Automatic tuning of the temperature of heated chains

Next, we tested how well the adaptive tuning of the temperature of heated chains over the course of an analysis works starting from different initial values. To do so, we ran two different datasets, the Hepatitis C dataset (Ray et al., 2000) as well an influenza A/H3N2 analysis using MASCOT as analysed previously (Müller, Rasmussen & Stadler, 2018). We ran each dataset with four different initial temperatures (0.0001, 0.001, 0.01 and 0.1), each targeting three different acceptance probabilities, 0.234, 0.468 and 0.117. Additionally, we used two different frequencies to propose swaps between chains, once proposing swaps every 100 iterations and once every 1,000. Since the temperature is adapted at every possible swap, this means that the runs with swaps every 100 iterations adapt ∆t 10 times more frequently than the ones proposing swaps every 1,000 iterations. We kept the temperature scaler constant for the first 100 potential swaps of states between chains.

As shown in Fig. S2, for any of the here considered initial values of the temperature scaler, the target acceptance probability is reached quite early in the run and very well approximated at the end of the run using the Hepatitis C example. The same applies to the analysis of the influenza A/H3N2 dataset (see Fig. S4).

After an initial phase where the adaption of the temperature difference can overshoot the optimal value, ∆t is adapted such that it approximates the target value better and better during the run (see Fig. 2 and Fig. S3 for the MASCOT analysis).

The effect of heating on exploring the posterior

In order to explore how heating affects exploring the posterior probability space, we next compared ESS values between regular MCMC and MC³ at different temperatures on a dataset where we do not expect any problems in exploring the posterior space caused by several local optimas. ESS values denote the number of effective samples if all samples would be drawn randomly from a distribution and are estimate here using Tracer (Rambaut et al., 2018).

To compare ESS values, we ran the Bayesian coalescent skyline (Drummond et al., 2005) analysis of Hepatitis C in Egypt (Ray et al., 2000) for 4 * 10⁷ iterations using MCMC in 100 replicates. We then compare these ESS values to those received when performing the same analysis using MC³ with four different chains for 1 * 10⁷ iterations using three different target acceptance probabilities, 0.468, 0.234 and 0.117. We also ran four times 100 additional analysis using different settings for the adaptive MC³ algorithm, all with a target acceptance probability of 0.234. First, we assume the temperature differences between chains to be distributed according to the quantiles of a beta distribution. We next allowed only swapping of states between chains with neighbouring temperature. Additionally, we estimate ESS values when running the same analysis using 8 and 16 chains for 5 * 10⁶ respectively 2.5 * 10⁶ iterations.

The different chain lengths between MCMC and MC³ are chosen such that the overall number of iterations over the cold and heated chains is the same for MC³ as for MCMC. After running all eight times 100 analyses, we computed the ESS values of the posterior probability estimates using loganalyser in BEAST 2 (Bouckaert et al., 2014).

As shown in Fig. 3A, the average ESS values are highest for the cold scenario when using MC³ and decrease with lower target acceptance probabilities. Lower target acceptance probabilities mean higher temperatures of heated chains in those analyses. With an increasing number of chains, but proportionally less iterations per chain, the ESS values decreases. This is particularly pronounced when using 16 chains.

We next tested if higher ESS values actually correspond to a run approximating the distribution of posterior probability values better. To do so, we compared Kolmogorov–Smirnov (KS) distances between individual runs and the true distribution of posterior values. The KS distance denotes the maximal distance between two cumulative density distributions, which is smaller the better two distributions match. Since we can not directly calculate the true distribution of posterior values, we concatenated the 800 regular and MC³ runs and used the concatenated distribution of posterior values as the true distribution. While this is technically not an independent run to compare to, each individual run contributes relatively little to the reference run.

Figure 3B shows the distribution of KS distances between individual runs using regular and MC³ to what we assume to be the true distribution. In contrast to the comparison of ESS values, we find that the distribution of KS distances is fairly comparable across all methods. This indicates that in this analysis, MC³ with four individual chains performs equally well as MCMC run for four times as long. With an increasing number of chains, however, this relationship hold less and less true. While the analysis with 8 chains still leads to a similar distribution of KS values, using 16 chains leads to a higher KS values.

We additionally tested how well the true tree distribution is recovered. To do so, we computed for each individual run the posterior clade support and compared it to a reference run consisting of all 100 runs combined. We then compare the maximal difference between clade support for each individual run to the reference run and show the estimated values in Fig. S5. Overall, the same patterns as for the KS distance holds true, with the analysis with 16 chains performing the worst, while the other analyses performed comparably.

It also shows that the differences in ESS values between the MC³ runs with different target acceptance probabilities are indicative of more swaps, rather than a better approximation of the true posterior probability distribution. Using the quantiles of a beta distribution instead of incremental heating as spacing between adjecent chains did not seem to impact the ESS values nor the KS distance. Only swapping states of chains with neighbouring temperatures performs equal to randomly swapping chains, but with a lower target acceptance probability. Swaps between neighbouring chains leads to, on average, hotter chains at the same acceptance probability.

We next compared the inference of trees on a dataset (typically referred to as DS1) that has proved problematic for tree inference using MCMC (Lakner et al., 2008; Höhna & Drummond, 2011; Whidden & Matsen, 2015; Maturana Russel et al., 2018). This dataset is essentially made up of different tree islands (Whidden & Matsen, 2015). Transitioning between the different tree island is highly unlikely due to very unfavourable intermediate states, making heating necessary to travel between local optima (Höhna & Drummond, 2011; Whidden & Matsen, 2015).

We ran the dataset using MCMC for 5 * 10⁷ iteration and MC³ for 5 * 10⁷ with 4 different chains. We ran MC³ targeting three different acceptance probabilities, that is 0.117, 0.234 and 0.468. As shown previously (Lakner et al., 2008; Whidden & Matsen, 2015) MCMC gets stuck in different local optimas, resulting in differences between inferred clade probabilities across different runs (see Fig. 4). As above, we additionally analysed this dataset using the quantiles of a beta distribution as spacing between chains or restricted swaps to only occur between neighbouring chains. We also ran two analyses with 8 and 16 chains, but with half respectively one quarter of the iterations per individual chain, such that the overall computations remained constant.

The clade probabilities are more comparable when targeting an acceptance probability of 0.468 and become more consistent between the different runs with acceptance probabilities of 0.234 and 0.117. At higher target acceptance probabilities (i.e. lower temperatures), the heating of chains is not sufficient to efficiently travel between local optimas.

We additionally compared how well the different runs approximate the posterior probability distribution compared to how long they ran. Consistent with for example Lakner et al. (2008), several MCMC runs sample from a different posterior probability distribution compared to MC³ with a low target acceptance probability and a high temperature (see Fig. S6). When running MC³ with a relatively high target acceptance probability of 0.468, the KS distance to the reference distribution decreases relatively slowly with the number of iterations compared to lower acceptance probabilities. This suggests that at lower temperatures (i.e. higher acceptance probabilities), some of the chains get stuck in local optimas (Brown & Thomson, 2018).

In all other scenarios using MC³, the KS values steadily decrease, indicating convergence (see Fig. S6). This suggest that for this dataset, the most important thing is that the temperature of at least some of the heated chains is high enough to overcome the unfavourable intermediate states. Once this is achieved there does not seem to be a big difference between the settings to explore the tree space.