Imagine a No-U-Turn Sampler that can adapt both its path-length and step-size on the fly, responding to the local geometry of the target distribution, while still preserving detailed balance. In our latest work, we have developed this novel capability, offering an important leap forward in MCMC sampling of continuously differentiable target densities. You can explore the full details in our paper on arXiv here. Let’s break down how it works.
The key innovation lies in our formulation of NUTS using the Gibbs self-tuning framework. Specifically, we treat all random variables within NUTS as auxiliary variables, including the symmetric Bernoulli process that determines the endpoints of leapfrog orbits. By conditioning on this Bernoulli process, the leapfrog orbit sampling becomes deterministic. This crucial step makes it possible to introduce adaptive step-size selection, where we monitor the energy envelope across the leapfrog orbit and dynamically select the largest step-size that keeps the energy envelope within a user-specified threshold, ensuring efficient sampling.
One of the main challenges in this adaptive approach is dealing with discretization artifacts that can break reversibility. Preserving detailed balance requires careful tracking of the acceptance probability. Our new formulation of NUTS within the Gibbs self-tuning framework enables this precise tracking, ensuring detailed balance even in the presence of such artifacts.
An important feature of our framework is that the acceptance probability depends solely on the conditional probability of the step-size. By avoiding any dependence on the conditional distribution of the auxiliary variables used in NUTS’ orbit and index selection kernels, we obtain a simpler and more transparent acceptance probability.
Step-size adaptive NUTS has shown great promise in overcoming mixing bottlenecks in notoriously difficult cases such as Neal’s funnel and high-dimensional problems. For those interested in testing it out, our implementation is publicly available on Bob Carpenter’s GitHub here.
I had the opportunity to present this work on August 16th during an invited session organized by Yuansi Chen (ETH Zürich) at the IMS-Bernoulli World Congress for Probability and Statistics, held in Bochum, Germany. The session provided a great platform to discuss the potential of this new approach in advancing NUTS. While there is still much work to be done to refine and expand this approach, we think this represents an important step forward in improving the efficiency of NUTS.