This post is by Bob

I’ve been working with Reuben Cohn-Gordon on evaluations of MCMC samplers, and we’ve come up with a scheme that I quite like for evaluation. Before we finish a more formal write up, I’d like to get some feedback from the blog readers.

Evaluating with an effective sample size target

First, let’s suppose we have an effective sample size (ESS) target, say something like ESS = 100 (or choose a dozen to make David MacKay’s ghost and Andrew happy or choose 10,000 because reviewer 2 insisted—the point is, pick a number that matters to you—I like 100). The standard error (standard deviation of the estimator) will be

se = sd / sqrt(ESS),

and thus errors will be distributed approximately as

err ~ normal(0, se),

where sd is the posterior standard deviation of the variable whose expectation is being estimated. With ESS = 100, the standard error will be 1/10th of a standard deviation and we thus expect 68% or so of our estimates to have an error of less than 0.1 sd. In this situation, our uncertainty in the value of a variable is dominated by standard deviation, not standard error. That’s why we stop at ESS = 100. You can choose a different ESS target if you like, but we have to choose one for a paper.

From a cold start

We’d like to evaluate how long a sampler takes to get to ESS = 100 from a cold start such as a random initialization. We care about cold starts because it’s usually the best we can do in applied problems. That can mean sampling from the prior or from something like a uniform(-2, 2) as we do in Stan. In some cases we can take a better guess or use something like Pathfinder variational inference to generate an init, but we’re not going to consider special cases like this.

Evaluate square parameter estimates

A problem with HMC is that it’s easy to boost ESS on parameter means (by taking roughly half an orbit where every other draw acts like a control variate) while tanking performance on parameters squared (see, e.g., this post on the Stan Forums about harmonics in HMC/jittered HMC/NUTS trajectories). So we want to evaluate estimates of both parameter posterior means and posterior variance.

Gradients are (usually) the dominant cost

We’ll evaluate using gradients as a proxy for compute, because most of the tight algorithm implementations are dominated by log density and gradient calls, which happen together with autodiff. The obstacle to just doing this is that the ESS estimators are themselves very noisy, especially for ESS around 25, which is what we need in each of four chains to hit ESS = 100 overall, for example. Of course, this doesn’t work for gradient-free samplers, such as Gibbs or random-walk Metropolis. We can measure using wall time, but that’s going to be very implementation dependent rather than algorithm dependent.

Reference moments versus effective sample size

So instead of estimating ESS, which is noisy, we propose to compare in terms of error versus reference moments. We have to be careful to evaluate those either analytically or with relatively low standard error compared to the samplers we’re evaluating so that noise in the reference moments does not dominate the evaluation. So we will take for granted that we can run millions of MCMC iterations and get reference moments that are much more precise than the estimates we will get from the shorter runs from the samplers being evaluated. This becomes very challenging if your ESS target gets much bigger than 100, say to 10,000 to cut errors a factor of 10 (we only get square root growth, so we need a sample 100 times larger to reduce errors by a factor of 10).

Our proposed metric

What we propose is running until the standardized error in each expectation of interest (D parameters, D parameters squared, 1 log density) is below 0.1. Suppose we have a random variable Z whose posterior mean is E[Z | y] and its posterior standard deviation is sd[Z | y]. Given an estimate hatZ, the standardized (not standard) error is

standardized-error = (hatZ – E[Z | y]) / sd[Z | y].

The standardized error is how many standard deviations the estimate is away from the mean. Taking Z = Theta, Z = Theta^2, and Z = log p(Theta) as our random variables of interest, we want to calculate how many gradients it takes until the standardized error in each one in each dimension is below 0.1.

Some issues

One of the practical issues that arises is that our algorithms like Stan’s implementation of the no-U-turn sampler (NUTS) is that it involves a fixed warmup period of N_warmup iterations and a fixed sampling period of N_sampling iterations. We only want to run warmup as long as necessary so that the total number of gradient evals to get all estimates below our prescribed error threshold is minimized. But we don’t know how many iterations or gradient evaluations that is a priori. So what we will do in practice is run increasingly large windows of adaptation as one might do in practice, starting with 50 iterations, then going to 100, then 200, and so on. Maybe with a base of sqrt(2) like f-stops in a camera, i.e., 50, 71, 100, 141, 200, 283, … Not perfect, but at least I won’t be making our cluster go brrr for a week.

The second practical issue is the usual one facing the evaluation of stochastic processes—there can be high variability due to random seeds. We can mitigate this somewhat by pairing the initializations in the comparisons, but it’s still very noisy. Stan can vary by an order of magnitude in the time it takes to solve these problems, so we need quite a few more evaluations than are traditional to have any confidence in a result such as the expected number of gradients required for a sampler to achieve standardized errors lower than 0.1. The noise is exacerbated by requiring the maximum error to be less than the tolerance.

I suspect we’ll hear more in the comments…