Stan and its implementation of dynamic Hamiltonian Monte Carlo is an extremely powerful tool for specifying and then fitting complex Bayesian models. In order to ensure a robust analysis, however, that power must be complemented with responsibility.

In particular, while dynamic implementations of Hamiltonian Monte Carlo, i.e.Â implementations where the integration time is dynamic, do perform well over a large class of models their success is not guaranteed. When they do fail, however, their failures manifest in diagnostics that are readily checked.

By acknowledging and respecting these diagnostics you can ensure that Stan is accurately fitting the Bayesian posterior and hence accurately characterizing your model. And only with an accurate characterization of your model can you properly utilize its insights.

# A Little Bit About Markov Chain Monte Carlo

Hamiltonian Monte Carlo is an implementation of Markov chain Monte Carlo, an algorithm which approximates expectations with respect to a given target distribution, $$\pi$$, $\mathbb{E}_{\pi} [ f ] = \int \mathrm{d}q \, \pi (q) \, f(q),$ using the states of a Markov chain, $$\{q_{0}, \ldots, q_{N} \}$$, $\mathbb{E}_{\pi} [ f ] \approx \hat{f}_{N} = \frac{1}{N + 1} \sum_{n = 0}^{N} f(q_{n}).$ Typically the target distribution is taken to be the posterior distribution of our specified model.

These estimators are guaranteed to be accurate only asymptotically, as the Markov chain grows to be infinitely long, $\lim_{N \rightarrow \infty} \hat{f}_{N} = \mathbb{E}_{\pi} [ f ].$ To be useful in applied analyses, we need these Markov chain Monte Carlo estimators to converge to the true expectation values sufficiently quickly that they are reasonably accurate before we exhaust our finite computational resources. This fast convergence requires strong ergodicity conditions to hold, typically a condition called geometric ergodicity between the Markov transition and target distribution. In particular, geometric ergodicity is a sufficient condition for Markov chain Monte Carlo estimators to follow a central limit theorem, which ensures not only that they are unbiased after only a finite number of iterations but also that we can empirically quantify their precision, $\hat{f}_{N} - \mathbb{E}_{\pi} [ f ] \sim \mathcal{N} \! \left( 0, \sqrt{ \mathrm{Var}[f] / N_{\mathrm{eff}}} \right).$

Unfortunately proving geometric ergodicity theoretically is infeasible for any nontrivial problem. Instead we must rely on empirical diagnostics that identify obstructions to geometric ergodicity, and hence well-behaved Markov chain Monte Carlo estimators. For a general Markov transition and target distribution, the best known diagnostic is the split $$\hat{R}$$ statistic over an ensemble of Markov chains initialized from diffuse points in parameter space.
To do any better we need to exploit the particular structure of a given transition or target distribution.

Hamiltonian Monte Carlo, for example, is especially powerful in this regard as its failures to be geometrically ergodic with respect to any target distribution manifest in distinct behaviors that have been developed into sensitive diagnostics. One of these behaviors is the appearance of divergences that indicate the Hamiltonian Markov chain has encountered regions of high curvature in the target distribution which it cannot adequately explore. Another is the energy Bayesian fraction of missing information, or E-BFMI, which quantifies the efficacy of the momentum resampling in between Hamiltonian trajectories.

For more details on Markov chain Monte Carlo and Hamiltonian Monte Carlo see Betancourt (2017).

In this case study I will demonstrate the recommended Stan workflow in R where we not only fit a model but also scrutinize these diagnostics and ensure an accurate fit.

# Setting Up The RStan Environment

We begin by importing the RStan module and setting some local options,

library(rstan)
Loading required package: ggplot2
Loading required package: StanHeaders
rstan (Version 2.14.1, packaged: 2016-12-28 14:55:41 UTC, GitRev: 5fa1e80eb817)
For execution on a local, multicore CPU with excess RAM we recommend calling
rstan_options(auto_write = TRUE)
options(mc.cores = parallel::detectCores())
rstan_options(auto_write = TRUE)
options(mc.cores = parallel::detectCores())

By setting rstan_options(auto_write = TRUE) we allow RStan to cache compiled models so that we can run them multiple times without the overhead of recompilation. options(mc.cores = parallel::detectCores()) enables RStan to run multiple Markov chains in parallel over any cores that your computer may have. These settings are recommended if you are running locally on your own machine, as opposed to running on a remote cluster, and your local machine which has plenty of RAM. For very large problems running multiple Markov chains in parallel may exhaust your RAM and degrade performance, in which case you would not want to utilize this option.

In order to facilitate checking diagnostics we source the attached utility script which loads some useful functions,

source('stan_utility.R')
lsf.str()
check_div : function (fit)
check_energy : function (fit)
check_treedepth : function (fit, max_depth = 10)
count_divergences : function (fit)
hist_treedepth : function (fit)
partition_div : function (fit)  

## Specifying and Fitting A Model in Stan

To demonstrate the recommended Stan workflow letâ€™s consider a hierarchical model of the eight schools dataset infamous in the statistical literature (D. B. Rubin 1981),

$\mu \sim \mathcal{N}(0, 5)$

$\tau \sim \text{Half-Cauchy}(0, 5)$

$\theta_{n} \sim \mathcal{N}(\mu, \tau)$

$y_{n} \sim \mathcal{N}(\theta_{n}, \sigma_{n}),$

where $$n \in \left\{1, \ldots, 8 \right\}$$ and the $$\left\{ y_{n}, \sigma_{n} \right\}$$ are given as data.

For more information on the eight schools dataset see Gelman et al. (2014).

## Specifying the Model with a Stan Program

In particular, letâ€™s implement a centered-parameterization of the model which is known to frustrate even sophisticated samplers like Hamiltonian Monte Carlo. In Stan the centered parameterization is specified with the Stan program

writeLines(readLines("eight_schools_cp.stan"))
data {
int<lower=0> J;
real y[J];
real<lower=0> sigma[J];
}

parameters {
real mu;
real<lower=0> tau;
real theta[J];
}

model {
mu ~ normal(0, 5);
tau ~ cauchy(0, 5);
theta ~ normal(mu, tau);
y ~ normal(theta, sigma);
}

Note that we have specified the Stan program in its own file. We strongly recommend keeping your workflow modular by separating the Stan program from the R environment in this way. Not only does it make it easier to identify and read through the Stan-specific components of your analysis, it also makes it easy to share your models Stan users exploiting workflows in environments, such as Python and the command line.

## Specifying the Data

Similarly, we strongly recommend that you specify the data in its own file.

Data specified in R lists can be immediately converted to an external Stan data file using RStanâ€™s stan_rdump function,

J <- 8
y <- c(28,  8, -3,  7, -1,  1, 18, 12)
sigma <- c(15, 10, 16, 11,  9, 11, 10, 18)

stan_rdump(c("J", "y", "sigma"), file="eight_schools.data.R")

At the same time, an existing Stan data file can be read into the R environment using the read_rdump function,

data <- read_rdump("eight_schools.data.R")

## Fitting the Model

With the model and data specified we can now turn to Stan to quantify the resulting posterior distribution with Hamiltonian Monte Carlo,

fit <- stan(file='eight_schools_cp.stan', data=data, seed=194838)
Warning: There were 126 divergent transitions after warmup. Increasing adapt_delta above 0.8 may help. See
http://mc-stan.org/misc/warnings.html#divergent-transitions-after-warmup
Warning: There were 2 chains where the estimated Bayesian Fraction of Missing Information was low. See
http://mc-stan.org/misc/warnings.html#bfmi-low
Warning: Examine the pairs() plot to diagnose sampling problems

We recommend explicitly specifying the seed of Stanâ€™s random number generator, as we have done here, so that we can reproduce these exactly results in the future, at least when using the same machine, operating system, and interface.
This is especially helpful for the more subtle pathologies that may not always be found, which results in seemingly stochastic behavior.

By default the sampling method runs 4 Markov chains of Hamiltonian Monte Carlo, each initialized from a diffuse initial condition to maximize the probability that at least one of the chains might encounter a pathological neighborhood of the posterior, if it exists. Because we set options(mc.cores = parallel::detectCores()) above, these chains will run in parallel when possible.

Each of those chains proceeds with 1000 warmup iterations and 1000 sampling iterations, totaling 4000 sampling iterations available for diagnostics and analysis.

## Validating a Fit in Stan

We are now ready to validate the fit using the information contained in the fit object.

## Checking Split $$\hat{R}$$ and Effective Sample Sizes

The first diagnostics we check are universal to Markov chain Monte Carlo and are displayed using the print method of the fit object,

print(fit)
Inference for Stan model: eight_schools_cp.
4 chains, each with iter=2000; warmup=1000; thin=1;
post-warmup draws per chain=1000, total post-warmup draws=4000.

mean se_mean   sd   2.5%    25%    50%    75% 97.5% n_eff Rhat
mu         4.50    0.18 3.45  -2.09   2.18   4.52   6.76 11.87   367 1.01
tau        4.29    0.19 3.34   0.60   1.95   3.36   5.62 13.27   321 1.01
theta[1]   6.77    0.19 5.94  -3.09   3.16   6.10   9.69 20.95   980 1.00
theta[2]   5.21    0.17 5.13  -4.58   2.05   5.41   8.07 15.43   915 1.00
theta[3]   3.89    0.21 5.88  -9.51   0.66   4.26   7.25 14.69   760 1.00
theta[4]   4.91    0.17 5.05  -5.20   1.68   4.96   7.88 15.00   922 1.00
theta[5]   3.57    0.18 5.03  -6.99   0.41   3.92   6.87 12.73   783 1.00
theta[6]   4.07    0.18 5.22  -6.87   0.86   4.43   7.36 13.56   845 1.00
theta[7]   6.83    0.18 5.34  -2.70   3.30   6.52   9.84 18.54   880 1.00
theta[8]   5.15    0.16 5.92  -6.44   1.63   5.12   8.24 17.29  1319 1.00
lp__     -16.11    0.44 5.81 -27.19 -20.05 -16.29 -12.05 -5.27   172 1.01

Samples were drawn using NUTS(diag_e) at Wed Jul 26 23:42:44 2017.
For each parameter, n_eff is a crude measure of effective sample size,
and Rhat is the potential scale reduction factor on split chains (at
convergence, Rhat=1).

Firstly we want to ensure that the split $$\hat{R}$$ for each parameter is close to 1. Empirically we have found that Rhat > 1.1 is usually indicative of problems in the fit. Here all of the parameters look good.

Then we want to consider the effective sample size, or n_eff. The issue here is related to the fact that we are estimating the effective sample size from the fit output. When n_eff / n_transitions < 0.001 the estimators that we use are often biased and can significantly overestimate the true effective sample size.

Both large split $$\hat{R}$$ and low effective samples per transition are consequences of poorly mixing Markov chains. Improving the mixing of the Markov chains almost always requires tweaking the model specification, for example with a reparameterization or stronger priors.

## Checking the Tree Depth

The dynamic implementation of Hamiltonian Monte Carlo used in Stan has a maximum trajectory length built in to avoid infinite loops that can occur for non-identified models. For sufficiently complex models, however, Stan can saturate this threshold even if the model is identified, which limits the efficacy of the sampler.

We can check whether that threshold was hit using one of our utility functions,

check_treedepth(fit)
[1] "0 of 4000 iterations saturated the maximum tree depth of 10 (0%)"

Weâ€™re good here, but if our fit had saturated the threshold then we would have wanted to rerun with a larger maximum tree depth,

fit <- stan(file='eight_schools_cp.stan', data=data, seed=194838, control=list(max_treedepth=15))

and then check if still saturated this larger threshold with

check_treedepth(fit, 15)

## Checking the E-BFMI

Hamiltonian Monte Carlo proceeds in two phases â€“ the algorithm first simulates a Hamiltonian trajectory that rapidly explores a slice of the target parameter space before resampling the auxiliary momenta to allow the next trajectory to explore another slice of the target parameter space. Unfortunately, the jumps between these slices induced by the momenta resamplings can be short, which often leads to slow exploration.

We can identify this problem by consulting the energy Bayesian Fraction of Missing Information,

check_energy(fit)

The check_energy function uses the threshold of 0.2 to diagnose problems, although this is based on preliminary empirical studies and should be taken only as a very rough recommendation. In particular, this diagnostic comes out of recent theoretical work and will be better understood as we apply it to more and more problems. For further discussion see Section 4.2 and 6.1 of Betancourt (2017).

As with split $$\hat{R}$$ and effective sample size per transition, the problems indicated by low E-BFMI are remedied by tweaking the specification of the model. Unfortunately the exact tweaks required depend on the exact structure of the model and, consequently, there are no generic solutions.

## Checking Divergences

Finally, we can check divergences which indicate pathological neighborhoods of the posterior that the simulated Hamiltonian trajectories are not able to explore sufficiently well. For this fit we have a significant number of divergences

check_div(fit)
[1] "126 of 4000 iterations ended with a divergence (3.15%)"
[1] "Try running with larger adapt_delta to remove the divergences"

indicating that the Markov chains did not completely explore the posterior and that our Markov chain Monte Carlo estimators will be biased.

Divergences, however, can sometimes be false positives. To verify that we have real fitting issues we can rerun with a larger target acceptance probability, adapt_delta, which will force more accurate simulations of Hamiltonian trajectories and reduce the false positives.

fit <- stan(file='eight_schools_cp.stan', data=data, seed=194838, control=list(adapt_delta=0.90))
Warning: There were 52 divergent transitions after warmup. Increasing adapt_delta above 0.9 may help. See
http://mc-stan.org/misc/warnings.html#divergent-transitions-after-warmup
Warning: There were 3 chains where the estimated Bayesian Fraction of Missing Information was low. See
http://mc-stan.org/misc/warnings.html#bfmi-low
Warning: Examine the pairs() plot to diagnose sampling problems

Checking again,

check_div(fit)
[1] "52 of 4000 iterations ended with a divergence (1.3%)"
[1] "Try running with larger adapt_delta to remove the divergences"

we see that while the divergences were reduced they did not completely vanish. In order to argue that divergences are only false positives, the divergences have to be completely eliminated for some adapt_delta sufficiently close to 1. Here we could continue increasing adapt_delta, where we would see that the divergences do not completely vanish, or we can analyze the existing divergences graphically.

If the divergences are not false positives then they will tend to concentrate in the pathological neighborhoods of the posterior. Falsely positive divergent iterations, however, will follow the same distribution as non-divergent iterations.

Here we will use the partition_div function of the stan_utility module to separate divergence and non-divergent iterations,

c_dark <- c("#8F272780")
green <- c("#00FF0080")

partition <- partition_div(fit)
div_params <- partition[[1]]
nondiv_params <- partition[[2]]

par(mar = c(4, 4, 0.5, 0.5))
plot(nondiv_params$'theta[1]', log(nondiv_params$tau),
col=c_dark, pch=16, cex=0.8, xlab="theta[1]", ylab="log(tau)",
xlim=c(-20, 50), ylim=c(-1,4))
points(div_params$'theta[1]', log(div_params$tau),
col=green, pch=16, cex=0.8)