Adding the WALNUTS sampler

[sethaxen]

The Within-Orbit Adaptive Leapfrog No-U-Turn Sampler (WALNUTS), presented in https://arxiv.org/abs/2506.18746, is an extension of NUTS that adapts to local curvature by dividing each integrator step into a certain number of sub-steps, where the number of sub-steps is chosen to keep the energy difference between the endpoints below a user-specified (or tuned) threshold. When the energy difference is low (i.e. curvature is low), it does no more work than NUTS; but when some regions have high curvature, where NUTS might require tuning a very low step size to make divergent transitions rare, WALNUTS can permit tuning a larger step size for exploring low-curvature regions while still exploring high-curvature regions well. They show that WALNUTS performs well for high-curvature models like Neal's funnel and hierarchical models.

The paper also proposes a strategy for tuning the global energy error and a different different step size tuning strategy than the one used in Stan.

Implementation

There are reference implementations at https://github.com/bob-carpenter/walnuts and https://github.com/flatironinstitute/walnuts.

From an initial skim of the paper, I suspect WALNUTS can be implemented by introducing a new AbstractIntegrator subtype that wraps another leapfrog integrator and using this with the existing NUTS sampler, but it's quite likely I'm missing something where information about the "micro steps" is needed elsewhere in the algorithm.

[nsiccha]

They (including @bob-carpenter) write (emphasis mine):

Remark 1. In the optimized implementation of WALNUTS, the energy error used to assess the validity of a micro trajectory is computed as the absolute difference of the Hamiltonian evaluated at the endpoints, i.e., | H (θ (0) , ρ(0) ) − H (θ (ℓ) , ρ(ℓ) )|. This choice improves practical performance by reducing mismatches between forward and backward micro steps, and ensures that the energy error threshold δ has a consistent interpretation across different values of the step size reduction factor ℓ. While this simplification deviates slightly from the theoretical formulation, the core theoretical guarantees are expected to remain intact.

Should this worry us/potential users?

[sethaxen]

Should this worry us/potential users?

I don't think so. That being said, computing the maximum energy error over the whole macro-step would be pretty cheap (with each micro-step, track the minimum and maximum energy encountered so far during the macro-step). If I understand their "practical performance" argument correctly, they're claiming that using the endpoints only makes it easier to pass the reversibility check. I don't see why this would be the case, but if true this could significantly decrease the amount of compute used. So far as I can tell, only their Lemma 1 result would be impacted by this decision.

FWIW, I have a local implementation of their optimized method (using p_micro as a delta function), and it definitely diverges less than NUTS on Neal's funnel but still has abysmally low ESS values. Certainly doesn't seem to resolve all of the issues that arise when a model is poorly conditioned (which is consistent with their results).

[bob-carpenter]

Certainly doesn't seem to resolve all of the issues that arise when a model is poorly conditioned

To that end, we're finding WALNUTS is performing worse than NUTS when there is a lot of correlation. Digging deeper, @torekleppe pointed out that the implementation I currently have is missing a key tuning parameter---the minimum number of micro steps per macro step. If you start with a larger baseline number of micro steps per macro step, it is easier to balance, and sampling highly correlated distributions stabilizes. With reasonable choices here, @torekleppe found WALNUTS outperformed NUTS by a factor of 2 in ESS per iteration and underperformed NUTS by a factor of two in ESS per gradient. So it does what we want, but it's too expensive. It turns out NUTS is using the divergence as a kind of early stopping criterion to not waste work on trajectories that are too expensive.

My original C++ implementation did just as @sethaxen suggests and kept a min and max over the whole trajectory rather than just testing the endpoints. That's a stricter condition---just the endpoints allows more variation within the trajectory and @torekleppe found. Testing along the whole trajectory has the advantage of being easy to early stop if the error is too big after a couple of steps. We're going to have to evaluate which one is best once we get a bigger evaluation harness instrumented.

My C++ implementation also takes a deterministic step size---the max in the schedule that satisfied the error criteria. In the paper, we show that you can select this according to any distribution you want conditioned on the state. Expanding this to be a bit more forgiving might introduce some inefficiencies in number of steps in some cases but make up for it elsewhere. I don't know the answer yet.

To conclude, what we really need is local mass matrix adaptation. WALNUTS lets you deal with varying scales, but it doesn't help with bad conditioning on a per gradient basis, at least as we've been able to tune it so far.

[bob-carpenter]

Also, I should point out that my C++ implementation uses an online version of the warmup from Nutpie. Either what I'm doing in the C++ repo or what @aseyboldt is doing in Nutpie in Rust is way better than the warmup schedule we have been using in NUTS. That's clear after some simple evaluations. The geometric motivation is to minimize Fisher divergence to keep the gradients as close as possible. The use of (co)variance of scores and of draws seems to work really well for mass matrix tuning, not only in terms of tuning but in terms of how fast @aseyboldt's method is to converge in adaptation.