NUTS tree build - Time reversibility

Hi, I’ve been trying to dig in to HMC theory and its variations and had a hard time understanding how the forward/backwards 2^n leapfrog integration tie in to allowing the NUTS to be reversible. Specifically, I’ve understood how a naive approach to a stopping rule for no U-Turns can mess with reversibility but I haven’t understood the intuition behind the solution proposed by Hoffman, and although the notation is succinct and elegant, it has quite given me that final push to understand.
If anyone has some intuitive way of seeing this I’d be very grateful.

Also I’d love to know how far I am from getting to the actual HMC used by Stan currently.

Thank you :)

2 Likes

@nhuurre @betanalpha

1 Like

The usual reference is A Conceptual Introduction to HMC. I think appendix A in particular may be what you’re looking for.

The transition is composed of two steps:

  1. starting from the initial point, build a trajectory
  2. draw a multinomial sample from all the points on the trajectory

For this to be reversible every point on the trajectory must have equal “step 1” likelihood of reaching the trajectory.

This is straightforward when the depth is fixed. Repeating multiplicative forward/backward expansion, you end up with a trajectory of 2^{D} points with the initial point being at \left(k+1\right) th position from the back end, (i.e. there are k points before it and 2^{D}-k-1 points after.)

k=\sum_{d=0}^{D}x_{d}\times 2^{d}

where x_{d}=1 if the expansion at depth d was backwards and 0 otherwise.
This is just the binary representation of a uniform random number in the range 0 to 2^{D}; conditional on the trajectory, all points on it have the same likelihood of having been the initial point, as required.

A dynamic termination criterion could be thought of as three-step process

  1. build the full trajectory at fixed maximum depth
  2. break the trajectory into maximal subtrees that do not have internal u-turns
  3. select the subtree that contains the initial point

Of course, in practice you want to stop building the trajectory as soon as you encounter the first u-turn but it may be more intuitive to consider the full binary tree and how it breaks into subtrees.

As an example, I made this picture based on figure 38 of the aforementioned paper.
trajectory
U-turn detected in the \left\{ z_{4},z_{5}\right\} subtree breaks the tree into three subtrees: \left\{ z_{4},z_{5}\right\} , \left\{ z_{3},z_{2}\right\} , and \left\{ z_{-2},z_{-1},z_{0},z_{1}\right\} . The initial point was z_{0} so the final point will be chosen from the four-element subtree; but the same trajectory and partitioning could be reached starting from any of the points.

5 Likes

Great explanation thank you! I think I get it now. I was following that same paper, but it is so complete that I believe I got overwhelmed from your explanation was very clear.