I’m on the hook to do a SciML webinar next week:
- Bob Carpenter. Multiscale generalized Hamiltonian Monte Carlo with delayed rejection. SciML Webinar, 28 September, 2023. 11 am (Eastern US time).
Slides: carpenter-sciml-webinar-2023.pdf
These are organized by Keith Phuthi (who is at CMU) through University of Michigan’s Institute for Computational Discovery and Engineering.
Sam Livingstone is moderating.This is presenting joint work with Alex Barnett, Chirag Modi, Edward Roualdes, and Gilad Turok.
I’m very excited about this project as it combines a number of threads I’ve been working on with collaborators. When I did my job talk here, Leslie Greengard, our center director, asked me why we didn’t use variable stepwise integrators when doing Hamiltonian Monte Carlo. I told him we’d love to do it, but didn’t know how to do it in such a way as to preserve the stationary target distribution.
Delayed rejection HMC
Then we found Antonietta Mira’s work on delayed rejection. It lets you retry a second Metropolis proposal if the first one is rejected. The key here is that we can use a smaller step size for the second proposal, thus recovering from proposals that are rejected because the Hamiltonian diverged (i.e., the first-order gradient based algorithm can’t handle regions of high curvature in the target density). There’s a bit of bookkeeping (which is frustratingly hard to write down) for the Hastings condition to ensure detailed balance. Chirag Modi, Alex Barnett and I worked out the details, and Chirag figured out a novel twist on delayed rejection that only retries if the original acceptance probability was low. You can read about it in our paper:
- Modi, Barnett, Carpenter. 2023. Delayed rejection Hamiltonian Monte Carlo for sampling multiscale distributions. Bayesian Analysis.
This works really well and is enough that we can get proper draws from Neal’s funnel (vanilla HMC fails on this example in either the tails in either the mouth or neck of the funnel, depending on the step size). But it’s inefficient in that it retries an entire Hamiltonian trajectory. Which means if we cut the step size in half, we double the number of steps to keep the integration time constant.
Radford Neal to the rescue
As we were doing this, the irrepressible Radford Neal published a breakthrough algorithm:
- Radford Neal. 2020. Non-reversibly updating a uniform [0,1] value for Metropolis accept/reject decisions.. arXiv.
What he managed to do was use generalized Hamiltonian Monte Carlo (G-HMC) to build an algorithm that takes one step of HMC (like Metropolis-adjusted Langevin, but over the coupled position/momentum variables) and manages to maintain directed progress. Instead of fully resampling momentum each iteration, G-HMC resamples a new momentum value then performs a weighted average with the existing momentum with most of the weight on the existing momentum. Neal shows that with a series of accepted one-step HMC iterations, we can make directed progress just like HMC with longer trajectories. The trick is getting sequences of acceptances together. Usually this doesn’t work because we have to flip momentum each iteration. We can re-flip it when regenerating, to keep going in the same direction on acceptances, but with rejections we reverse momentum (this isn’t an issue with HMC because it fully regenerates each time). So to get directed movement, we need steps that are too small. What Radford figured out is that we can solve this problem by replacing the way we generate uniform(0, 1)-distributed probabilities for the Metropolis accept/reject step (we compare the variate generated to the ratio of the density at the proposal to the density at the previous point and accept if it’s lower). Radford realized that if we instead generate them in a sawtooth pattern (with micro-jitter for ergodicity), then when we’re at the bottom of the sawtooth generating a sequence of values near zero, the acceptances will cluster together.
Replacing Neal’s trick with delayed rejection
Enter Chirag’s and my intern, Gilad Turok (who came to us as an undergrad in applied math at Columbia). Over the summer, working with me and Chirag and Edward Roualdes (who was here as a visitor), he built and evaluated a system that replaces Neal’s trick (sawtooth pattern of acceptance probability) with the Mira’s trick (delayed rejection). It indeed solves the multi scale problem. It exceeded our expectations in terms of efficiency—it’s about twice as fast as our delayed rejection HMC. Going one HMC step at a time, it is able to adjust its stepsize within what would be a single Hamiltonian trajectory. That is, we finally have something that works roughly like a typical ODE integrator in applied math.
Matt Hoffman to the rescue
But wait, that’s not all. There’s room for another one of the great MCMC researchers to weigh in. Matt Hoffman, along with Pavel Sountsov, figured out how to take Radford’s algorithm and provide automatic adaptation for it.
- Hoffman and Sountsov. 2022. Tuning-Free Generalized Hamiltonian Monte Carlo. AIStats.
What Hoffman and Sountsov manage to do is run a whole lot of parallel chains, then use information in the other chains to set tuning parameters for a given chain. In that way it’s like the Goodman and Weare affine-invariant sampler that’s used in the Python package emcee. This involves estimating the metric (posterior covariance or just variance in the diagonal case) and also estimating steps size, which they do through a heuristic largest-eigenvalue estimate. Among the pleasant properties of their approach is that the entire setup produces a Markov chain from the very first iteration. That means we only have to do what people call “burn in” (sorry Andrew, but notice how I say other people call it that, not that they should), not set aside some number of iterations for adaptation.
Edward Roualdes has coded up Hoffman and Sountsov’s adaptation and it appears to work with delayed rejection replacing Neal’s trick.
Next for Stan?
I’m pretty optimistic that this will wind up being more efficient than NUTS and also make things like parallel adaptation and automatic stopping a whole lot easier. It should be more efficient because it doesn’t waste work—NUTS goes forward and backward in time and then subsamples along the final doubling (usually—it’s stochastic with a strong bias toward doing that). This means we “waste” the work going the wrong way in time and beyond where we finally sample. But we still have a lot of eval to do before we can replace Stan’s longstanding sampler or even provide an alternative.
My talk
The plan’s basically to expand this blog post with details and show you some results. Hope to see you there!
I thought about this when was watching these videos of metronomes automatically synchronizing (transferring energy between each other via motion of the table) the other day:
https://www.youtube.com/watch?v=5v5eBf2KwF8
On an earlier post I mentioned the similar “walking droplet” synchronization phenomenon, but the metronomes may be an easier to understand example. I think there is some lesson to be learned from nature here about how to optimally explore whatever parameter space using many parallel trajectories. But I haven’t been able to put my finger on it.
The key here is trying to find a method that will work when things are way out of synch, like in a random initialization. We’ve been trying to tackle the problem of going from random init in the tail to the bulk of the probability mass separately with VI techniques like Pathfinder. Once you’re near the true answer (which is when Stan would ideally transition from Phase 1 to Phase 2 of its 3-phase warmup), it’s beneficial to start combining information on metric adaptation, step size, etc. It might be beneficial beforehand, but the geometries can be wildly varying out in the tails, causing things like stiffness in ODEs. Ben Bales worked on this in his dissertation and then he and Yi Zhang evaluated it some more around Stan. But we’ve been thinking of going to these really nice complementary methods because they don’t require a separate warmup phase.
How about chain A influences chain B with some lag and “force” that are both a function of distance. Ie, the lag increases and the influence decreases the farther apart the chains are in the parameter space.
Obviously this requires doing many more calculations per step. But perhaps in the end we would need fewer steps.
Do you know of anyone working on something like this?
> Going one HMC step at a time, it is able to adjust its stepsize within what would be a single Hamiltonian trajectory.
Probably an ignorant question, but is there any chance that these new algorithms could do away with the necessity of non-centered parameterizations for multilevel model implementations?
Yes and no. Yes in the sense that the algorithm we already published combining delayed rejection and HMC can sample from non-centered parameterizations (HMC and NUTS actually fail and get biased answers on the pure funnel with any reasonable compute time bound as we show in the paper). No in the sense that it’s still more efficient to provide a better-conditioned parameterization. What we really need is something like a low-rank Riemannian HMC that does second-order corrections for geometry (though in practice, we can’t usually compute Fisher info and the Hessian-based approximations need to be conditioned for positive-definiteness).
One thing you can do in Stan, at least for unconstrained parameters, is use offset and multiplier. So that’d be
parameters { int<lower=0> N; real mu; real<lower=0> sigma; vector<offset=mu, multiplier=sigma>[N] alpha; } model { alpha ~ normal(mu, sigma);By multiplying by sigma then adding mu, the result is effectively a non-centered parameterization of alpha.