GIST: Gibbs self-tuning for locally adaptive Hamiltonian Monte Carlo
Nawaf Bou-Rabee Bob Carpenter Milo Marsden


UniMERNet: A Universal Network for Real-World Mathematical Expression Recognition
    
    
    
   Bin Wang^*, 
   Zhuangcheng Gu^*, 

   Chao Xu,
   Bo Zhang,
   Botian Shi,
   Conghui He^†

    Received xx; accepted xx
==============================================================================================================


We present a novel and flexible framework for localized tuning of Hamiltonian Monte Carlo samplers by sampling the algorithm's tuning parameters conditionally based on the position and momentum at each step. For adaptively sampling path lengths, we show that randomized Hamiltonian Monte Carlo, the No-U-Turn Sampler, and the Apogee-to-Apogee Path Sampler all fit within this unified framework as special cases.  The framework is illustrated with a simple alternative to the No-U-Turn Sampler for locally adapting path lengths.







§ INTRODUCTION


Tuning the parameters of Markov chain Monte Carlo (MCMC) algorithms is critical to their  performance, but notoriously difficult in practice.  This challenge is particularly acute for Hamiltonian Monte Carlo (HMC), where the tuning of path length (integration time) <cit.>, step length (time discretization parameter) <cit.>, and pre-conditioner (mass matrix) <cit.> frequently presents a complex trade-off between computational cost and mixing.   The successful self-tuning of path length provided by the No-U-Turn Sampler (NUTS) has led to its widespread adoption as the default sampler in probabilistic programming languages <cit.>.

In NUTS, the path length is adaptively sampled according to U-turn avoiding conditions, which roughly speaking stop the underlying Hamiltonian trajectory whenever it doubles back <cit.>. The main idea is to numerically integrate Hamilton's equations forward and backward in time until a U-turn occurs.[A U-turn must occur in a proper density, but in practice, a cap such as 1024 iterations is imposed.]  A point along this path is then randomly selected such that: (i) the resulting sampler has the correct invariant distribution; and (ii) points far from the starting point are more likely to be selected.  

Can the algorithm's other tuning parameters  be similarly self-tuned?   Motivated by the benefits of NUTS, this paper addresses this question by presenting a new framework for adaptively sampling HMC tuning parameters such as  path length, step length, pre-conditioner, etc.[The approach presented here could also be applied to Metropolis samplers other than HMC.]  This framework is termed Gibbs self-tuning HMC or GIST for short.  In the context of adaptively sampling path lengths,  we demonstrate that GIST includes as special cases almost all locally adaptive HMC samplers in current use including randomized HMC <cit.>, NUTS <cit.>, and the Apogee-to-Apogee Path Sampler <cit.>.  Owing to its simplicity and generality, GIST can be naturally extended to adapting the algorithm's remaining tuning parameters including step size and mass matrix.

To be sure, the primary contribution of this paper is the GIST sampler, which admits a relatively simple proof of correctness.  Moreover, we demonstrate the utility of the GIST sampling framework as a theoretical tool by unifying the proofs of correctness for existing locally adaptive HMC samplers and as a practical tool by developing a novel alternative for locally adapting path lengths.

The remainder of the paper is organized as follows.  
Section <ref> presents the GIST sampling framework for  self-tuning HMC samplers.  In the context of adaptively sampling path lengths, Section <ref> considers some fundamental special cases of  GIST samplers including randomized HMC (Section <ref>), NUTS (Section <ref>), and the Apogee-to-Apogee Path Sampler (Section <ref>).  Section <ref> compares GIST samplers for path-length adaptation in the continuous-time context applied to a truncation of an infinite-dimensional Gaussian target measure.  Section <ref> considers path-length adaptation in the discrete-time context, and presents several concrete proposals for sampling the number of leapfrog steps based on the current position and momentum given a fixed step size and mass matrix.  Section <ref> describes a novel sampler for adapting path lengths along with experimental evaluation and an empirical comparison to NUTS.





§ A FRAMEWORK FOR SELF-TUNING HAMILTONIAN MONTE CARLO




Hamiltonian Monte Carlo (HMC) is a class of MCMC methods for sampling absolutely continuous probability target distributions of the form 

    μ(dθ, dρ) ∝ e^-H(θ, ρ)  m( dθ  dρ)  ,

which have a density relative to Lebesgue measure m on phase space ℝ^2d, with position θ∈ℝ^d and momentum ρ∈ℝ^d.  For simplicity, we assume  the case of a separable Hamiltonian function with unit mass matrix, i.e., 
    H(θ,ρ) = U(θ) + 1/2 ρ^⊤ ρ ,

for a continuously differentiable potential energy function U: ℝ^d →ℝ such that ∫ e^-U(θ) dθ < ∞.   The non-normalized target density of interest is e^-U(θ) for model parameters θ.

A defining ingredient of any HMC algorithm is a reversible, volume-preserving map F(α): ℝ^2d→ℝ^2d where α encapsulates the algorithm tuning parameters.  This map is typically obtained from discretizing Hamilton's equations for the Hamiltonian function H(θ,ρ), i.e.,  
    d/dtθ_t = ρ_t  ,   d/dtρ_t = - ∇ U(θ_t)  .
 Reversibility and volume-preservation ensure that the map F(α) is Metropolis adjustable <cit.>.  Note that 𝒮∘ F(α) is a volume-preserving involution where 𝒮: ℝ^2d→ℝ^2d is the momentum flip map defined by 𝒮(θ, ρ) = (θ, -ρ) for all (θ, ρ) ∈ℝ^2d <cit.>.  We suppose that the tuning parameter α
takes values in a set 𝒜 where (𝒜, ℬ, η) is a measure space with σ-algebra ℬ and background measure η.  

To “self tune”  the parameter α, the state space ℝ^2d is enlarged to a product space ℝ^2d×𝒜. On this enlarged space, an enlarged target measure is defined by specifying a conditional distribution of the tuning parameter α given the position and momentum (θ, ρ), i.e.,  
    μ_e(dθ, dρ, dα) ∝  e^-H(θ, ρ)  p( α  |  θ, ρ)   (m ⊗η) (dθ  dρ  dα)  .
  Note that the (θ, ρ)-marginal of μ_e is the desired target measure μ.  Explicit forms of the conditional distribution p( α  |  θ, ρ) are specified in subsequent sections for the case of path length adaptivity.  


For any position and momentum (θ, ρ) ∈ℝ^2d, let π(θ, ρ) : 𝒜→𝒜 be a measurable map acting on the  tuning parameter α∈𝒜 such that 
    G: (θ, ρ, α) ↦( S ∘ F(α)(θ, ρ),  π(θ, ρ)(α) )
 is a (m ⊗η)-preserving involution on ℝ^2d×𝒜.   This clearly holds when the map π(θ, ρ) is taken to be the identity map on 𝒜.

In terms of this notation, Figure <ref> describes a transition step of the “GIST” (Gibbs self-tuning HMC) sampler.  The following theorem indicates that the GIST sampler is correct.  
 
The GIST sampler in Figure <ref> generates a Markov Chain θ_0, θ_1, … on ℝ^d that is reversible with respect to e^-U( θ)  m(d θ).



A detailed, self-contained proof of Theorem <ref> is provided in Appendix <ref>. By interpreting the GIST sampler as a Gibbs sampler, the proof shows that each update in Figure <ref> leaves invariant the enlarged target (<ref>).  In particular, the accept/reject step is interpreted as a Metropolis-within-Gibbs step on the enlarged space ℝ^2d×𝒜 with target measure defined in (<ref>), proposal given by G defined in (<ref>), and corresponding acceptance probability 
    a_e(θ_0, ρ_0, α) = 1 ∧(  e^ -Δ H(θ_0, ρ_0) p(π(θ_0, ρ_0)(α)   |  𝒮∘ F(α)(θ_0, ρ_0) ) p (α  |  θ_0, ρ_0 ))  ,
 
where we used the shorthand Δ H(θ, ρ) := H ∘ F(α)(θ,ρ) - H(θ, ρ) for (θ, ρ, α) ∈ℝ^2d×𝒜.
Note, this Metropolis-within-Gibbs step is correct because we assume that the map G on ℝ^2d×𝒜 is a (m ⊗η)-preserving involution; for a detailed justification see Lemma <ref> in Appendix <ref>.


Just as the Markov chain for standard HMC can be formulated over phase space variables (θ, ρ) ∈ℝ^2d, the Markov chain for GIST can also be formulated over triples (θ, ρ, α) ∈ℝ^2d×𝒜.  However, in the case treated here of full momentum and tuning parameter refreshment, the sequence of θ values by itself forms a Markov chain and including the other variables is superfluous.



In addition to adaptively sampling HMC tuning parameters, the GIST sampler can also be applied to randomize the time integrator for the Hamiltonian flow, as in <cit.>.  In this case, the tuning parameter would specify a particular time integrator within a parametric family of time integrators that are each reversible and volume-preserving.  In certain representative models, randomized time integrators have provably better complexity for Hamiltonian MCMC than the frequently used leapfrog integrator <cit.>.  




§ ADAPTIVELY SAMPLING PATH LENGTHS


To make the GIST sampling framework more concrete and demonstrate its breadth, this section shows that several fundamental HMC samplers that locally adapt path lengths fit within this unified framework.  



 §.§ Sampling path lengths in randomized HMC


In randomized forms of HMC, the path length or the step length is generated randomly at each iteration independently of the current position and momentum <cit.>; using, e.g., an exponential distribution over path length <cit.>, an empirically learned distribution over path length <cit.>, or a uniform distribution over path and step lengths <cit.>.  These can all be analyzed as instances of the GIST sampler given in Figure <ref>.  

As an example, consider F(α)(θ, ρ) = φ_α(θ, ρ) where we have introduced the exact Hamiltonian flow φ_α: ℝ^2d→ℝ^2d at time α≥ 0. Take 𝒜 = [0, ∞) and η to be the Lebesgue measure on ℝ. Define p(α|θ, ρ) = λ e^- λα where λ > 0; that is, α is an exponential random variable with parameter λ.  In this case, the map G in (<ref>) is given explicitly by G : (θ, ρ, α) ↦ ( 𝒮∘ F(α)(θ, ρ),  α), for which a_e ≡ 1 (the proposal is always accepted), and Figure <ref> reduces to a draw from randomized HMC at the first jump time <cit.>.

The other randomized forms of HMC that generate randomly the tuning parameter independently of the current position and/or momentum can be analyzed in exactly the same way.  In Section <ref>, in the setting of a truncation of an infinite-dimensional Gaussian measure, the performance of randomized HMC is compared with a GIST sampler that we describe next. 



 §.§ Adapting path lengths in exact HMC


Consider again F(α) = φ_α where  φ_α: ℝ^2d→ℝ^2d is the exact Hamiltonian flow map at time α≥ 0.  Take 𝒜 = [0, ∞) and η to be the Lebesgue measure on ℝ.  Let τ(θ, ρ):ℝ^2d→ (0, ∞) be any measurable function.   Define 
    p(α|θ, ρ)   =  1τ(θ, ρ)1_[0,  τ(θ,  ρ)]( α )  .
 That is, conditioned on the position and momentum (θ, ρ), α is a uniform random variable over the interval  [0,τ(θ, ρ)].  Here 1_A denotes the standard indicator function of a set A.  As a shorthand, let τ_1 = τ(θ_0, ρ_0) and τ_2 = τ(𝒮∘φ_α(θ_0, ρ_0)), i.e., the function τ evaluated at the current state of the chain (θ_0, ρ_0) and the proposed state but with momentum reversed 𝒮∘φ_α(θ_0, ρ_0).
In this case, the map G in (<ref>) is given explicitly by 
    G : (θ, ρ, α) ↦ ( 𝒮∘ F(α)(θ, ρ),  α)  ,
 and the corresponding acceptance probability in (<ref>) reduces to 
    a_e(θ_0, ρ_0, α)   =   1 ∧( τ_1/τ_21_{τ_2 ≥α})  ,
 because Δ H(θ_0, ρ_0) = 0 for the exact Hamiltonian flow.  The indicator in this Metropolis ratio indicates that  p(α|𝒮∘φ_α(θ_0, ρ_0))  0.

Conditions avoiding U-turns in the exact Hamiltonian flow can be used to specify the function τ(θ, ρ).  There are several ways to characterize such U-turn conditions. For example, here is a condition based on when the angle between the initial velocity ρ_0 ∈ℝ^d and the velocity ρ_t ∈ℝ^d at time t ≥ 0 first exceeds π/2 
    τ(θ_0, ρ_0)  :=  inf{ t > 0  : ρ_0 ·ρ_t ≤ 0 } ,
 where we have introduced (θ_t, ρ_t) := φ_t(θ_0, ρ_0) for t ≥ 0.  Another condition is based on when the squared distance Γ(t) := | θ_0 - θ_t|^2 between the initial configuration θ_0 ∈ℝ^d and the configuration θ_t ∈ℝ^d at time t ≥ 0 first decreases 
    τ(θ_0, ρ_0)  :=  inf{ t > 0  : Γ'(t) < 0 } .

These continuous-time U-turn conditions have discrete-time analogs, which are discussed in the next example and Section <ref>.  






 §.§ Adapting path lengths in numerical HMC


Fix h>0.  Let Φ_h: ℝ^2d→ℝ^2d denote one step of the leapfrog integrator with step length h.  Consider F(α) = Φ_h^α where α is the number of leapfrog steps. Take 𝒜 = ℕ and let η be the counting measure. Similarly to the above, let τ: ℝ^2d→ℕ be a measurable function and define 
    p(α|θ, ρ) = 1/τ(θ, ρ)1_{0, …,  τ(θ,  ρ)-1}(α)  .

Similarly to the above, the map G in (<ref>) is given explicitly by 
    G : (θ, ρ, α) ↦ ( 𝒮∘ F(α)(θ, ρ),  α)  ,
 and the corresponding acceptance probability in (<ref>) simplifies to 
    a_e(θ_0, ρ_0, α)   =   
    1 ∧( e^-Δ H(θ_0, ρ_0)τ_1/τ_21_{τ_2 ≥α})  ,

where τ_1 = τ(θ_0, ρ_0) and τ_2 = τ(𝒮∘Φ^α_h(θ_0, ρ_0)).




 §.§ Adapting path lengths in NUTS
 

Here we show that NUTS is a special case of the GIST sampler presented in Figure <ref>. 

In order to write NUTS as a GIST sampler, we first explain NUTS more precisely following the terminology and framework introduced in <cit.>.   Given the current position and momentum (θ_0, ρ_0) ∈ℝ^2d, the idea behind NUTS is to first randomly generate a finite subset of integers J ⊂ℤ, and in turn, sample the next state of the chain from among the leapfrog iterates {Φ_h^i(θ_0, ρ_0) }_i ∈ J. Following the notation of <cit.>, let P_h(J |θ_0, ρ_0) be the probability of selecting J and let Q_h(i | J, θ_0, ρ_0) be the probability  of selecting the i-th leapfrog iterate Φ_h^i(θ_0, ρ_0). In <cit.>, these are referred to as the “orbit” and “index” selection kernels, respectively. 

More precisely, the random procedure to generate the set J can be described by induction. First, take J_0 = { 0 }. Then, given J_k, we first check whether the leapfrog iterates {Φ_h^i(θ_0, ρ_0) }_i ∈ J_k  satisfy the U-turn condition:

    ρ_+ · (θ_+ - θ_-) < 0   and  ρ_- · (θ_- - θ_+) < 0

where (θ_+, ρ_+) = Φ_h^maxJ_k(θ_0, ρ_0) and (θ_-, ρ_-) = Φ_h^min J_k(θ_0, ρ_0). If this condition is met, we output J = J_k. If instead no U-turn is encountered, then sample I_k+1 uniformly from { J_k + |J_k|, J_k - |J_k| }. This set is the next proposed extension to the set J_k. 

This proposed extension will either be accepted, yielding J_k+1 = J_k ∪ I_k+1, or rejected, returning J = J_k. To ensure the correctness and efficiency of the algorithm the following symmetry condition is imposed: P_h(J |θ_0, ρ_0) = P_h(J - i |Φ_h^i(θ_0, ρ_0)) for all i ∈ J. This condition ensures that the probability of obtaining a particular set of leapfrog iterates {Φ_h^i(θ_0, ρ_0)}_i ∈ J is uniform with respect to the starting point Φ_h^i(θ_0, ρ_0) within this set for all i ∈ J. 

Imposing this constraint will require that some proposed extensions I_k+1 are rejected. For every i ∈ J_k, corresponds an alternate sequence of extensions I'_1, …, I'_k which would have produced the set J_k - i had the orbit generating procedure instead been started at Φ_h^i(θ_0, ρ_0).  To enforce P_h(J |θ_0, ρ_0) = P_h(J-i |Φ_h^i(θ_0, ρ_0)), it is sufficient to require that J_k satisfies the No-Intermediate-U-Turn property. This property states that at no stage in this alternate sequence of extensions I_1', …, I_k' would the U-turn condition  have been triggered.








To ensure that every J_k satisfies the No-Intermediate-U-Turn property, we simply need to reject any proposal I_k+1 where J_k+1 = J_k ∪ I_k+1 does not satisfy the No-Intermediate-U-Turn property. While naively this appears to involve checking over exponentially many sequences of extensions, this step can be implemented efficiently. Indeed, over different choices of i ∈ J_k many of the alternate sequences will refer to the same set of leapfrog iterates. Since the U-turn condition depends only on the leapfrog iterates, we therefore need only check that all possible intermediate sets of leapfrog iterates have no U-turns. As the number of possible intermediate sets of leapfrog iterates is relatively small, this observation greatly reduces the complexity of checking the No-Intermediate-U-Turn property. See Remark <ref> for more details about the computational complexity of the orbit generation procedure.

Having generated J using this randomized orbit selection procedure, NUTS then generates i ∈ J according to the index selection kernel Q_h(i |θ_0, ρ_0, J) and returns Φ_h^i(θ_0, ρ_0) as the next state. 
There are two index selection kernels commonly used in practice. The first is given by the Boltzmann weight of the corresponding leapfrog iterate, i.e., Q_h(i |θ_0, ρ_0, J) ∝ e^-H(Φ_h^i(θ_0, ρ_0)). The second is biased progressive sampling, as detailed in <cit.>. 

We are now in position to describe NUTS as a GIST sampler. Let 𝒜 = 𝒫_F ×ℤ where 
𝒫_F = {J ⊂ℤ| |J|<∞}, let η be the counting measure on 𝒫_F ×ℤ, and write the tuning parameter in components as α = (J, i) ∈𝒜. For all (θ, ρ, α) ∈ℝ^2d×𝒜, define

    p(J, i |θ, ρ) = P_h(J |θ, ρ) Q_h(i |θ, ρ, J)  .

As in Section <ref>, let F(α)(θ, ρ) = Φ_h^i(θ, ρ).
The map G given in (<ref>) can be written explicitly as

    G(θ, ρ, J, i) = (𝒮∘Φ_h^i(θ, ρ), -(J- i), i)  .

The resulting GIST sampler in Figure <ref> corresponds to NUTS. Notably, the Metropolis-within-Gibbs step for NUTS is always accepted. Indeed, since {Φ_h^i(S(θ, ρ))}_i ∈ -J = {Φ_h^-i(S(θ, ρ))}_i ∈ J ={S(Φ_h^i(θ, ρ))}_i ∈ J a short coupling construction shows P_h(-J |𝒮(θ, ρ)) = P_h(J |θ, ρ). Combining this with P_h(J-i |Φ_h^i(θ, ρ)) = P_h(J |θ, ρ) gives P_h(-(J-i) |𝒮∘Φ_h^i(θ, ρ)) = P_h(J |θ, ρ). Since the index selection kernels described above satisfy the following detailed balance condition

    e^-H(Φ_h^i(θ_0, ρ_0)) Q_h(i |𝒮∘Φ_h^i(θ_0, ρ_0), -(J - i)) = e^-H(θ_0, ρ_0) Q_h(i |θ_0, ρ_0, J)  ,

it follows that the acceptance probability in Figure  <ref> is always equal to 1. For the index selection kernel based on the Boltzmann weights this detailed balance condition can be seen by a direct computation, and for biased progressive sampling this property follows from Proposition 6 in <cit.>.

  As noted above, one can efficiently check whether J_k+1 satisfies the No-Intermediate-U-Turn condition by checking whether all possible intermediate sets of leapfrog iterates are free of U-turns. If we define random variables β_1, …, β_k+1 as β_l = 1_I_l = J_l - |J_l| and B = ∑_l=1^k+1 2^l-1β_l then one can prove by induction that J_k+1 = {-B, …, 2^k -1-B } = [2^k] - (B+1). The possible sets of intermediate leapfrog iterates corresponding to J_k+1 then correspond exactly to the sets {Φ_h^i(θ, ρ)}_i ∈ [m· 2^l : (m+1)2^l -1] - (B+1), for l=1, …, k and m = 0, 1, …, 2^k+1-l - 1,  where [m : n ] = {m, m+1, …, n-1, n }. 

Since the U-turn condition involves only the endpoints of the set of leapfrog iterates, the cost of checking whether or not a given set of leapfrog iterates satisfies the U-turn condition is constant in the size of the set. Therefore, when generating an orbit starting from (θ_0, ρ_0) the total cost of checking whether J_k+1 = I_k+1∪ J_k satisfies the No-Intermediate-U-Turn condition is upper bounded by the number of sets we need to check. 



Since each choice of l above corresponds to 2^k-l sets the total cost for checking that J_k+1 satisfies the No-Intermediate-U-Turn property is upper bounded by ∑_l=1^k 2^k-l≤ 2^k+1. Consequently, returning a set J with |J| = M = 2^k has computational cost O(M). In practice, by induction we need only check those intermediate sets of leapfrog iterates which correspond to subsets of I_k+1 which gives further speedup.

Compare this to the simpler strategy of extending J_k one point at a time by sampling I_k+1 uniformly from {min(J_k) -1, max(J_k)+1 }. For such a strategy, the cost of appending I_k+1 is again of order |J_k| and therefore to return a set J with |J| = M has computational cost O(M^2). Reducing the cost of producing large orbits motivates the strategy of having the orbits J_k double in size with each iteration. For more details on the computational cost of the orbit selection procedure in NUTS as well as other aspects of the implementation, see Appendix A in <cit.>.




We can consider I_k+1 in the same way. That is, to every point i ∈ I_k+1 also corresponds a hypothetical sequence of binary choices U_1”, …, U_k” which represents a sequence of intermediate sets I'_1, …, I'_k which would have produced from Φ_h^i(θ_0, ρ_0) the same set of leapfrog iterates L((θ_0, ρ_0), I_k+1) as are represented by I_k+1 starting at (θ_0, ρ_0). In order to preserve the property P_h(J |θ_0, ρ_0) = P_h(J-i |Φ_h^i(θ_0, ρ_0)) we need only check that none of these hypothetical intermediate sets would have triggered the No-U-Turn condition. If I_k+1 satisfies this property, we set J_k+1 = J_k ∪ I_k+1. Otherwise, we return J_θ_0, ρ_0 = J_k. We continue this procedure until J_k satisfies the U-Turn condition, we encounter a rejected proposal, or J_k reaches some user-specified maximum size. As a consequence of this careful construction, the orbit selection kernel satisfies P_h(J |θ_0, ρ_0) = P_h(J-i |Φ_h^i(θ_0, ρ_0)).





 §.§ Adapting path lengths in the Apogee-to-Apogee Path Sampler


The Apogee-to-Apogee Path Sampler (AAPS) introduced in <cit.> is a special case of the GIST sampler presented in Figure <ref>.  The idea behind AAPS is to replace the No-U-Turn condition for adapting the number of leapfrog steps with an alternative condition based on partitioning the discrete path {Φ_h^i (θ, ρ) }_i ∈ℤ into `segments' {S_i(θ, ρ) }_i ∈ℤ. The i-th segment S_i(θ, ρ) consists of points along the discrete path originating from the current point (θ, ρ) that lie between consecutive local maxima (termed “apogees”) of the underlying  potential energy U(x). 

In order to write AAPS as a GIST sampler, we first explain AAPS more precisely starting with the partitioning mentioned above. Our description follows <cit.>.  Proceeding both forwards and backwards in time from (θ, ρ)  via leapfrog steps, consecutive outputs belong to the same segment if the potential U is either instantaneously increasing at both points or instantaneously decreasing at both points. Otherwise, an apogee occurs between this pair of points, and hence, a new segment is defined. Iterating this procedure produces a two-sided sequence of segments: 
    …, S_-1(θ, ρ), S_0(θ, ρ), S_1(θ, ρ), …
 
where S_0(θ, ρ) is the initial segment.  Following the notation of <cit.>, let S_#((θ, ρ), (θ', ρ')) be the index of the segment started at (θ, ρ) containing the leapfrog iterate (θ', ρ'). For instance, S_#((θ, ρ), (θ, ρ)) = 0 since the starting point is always in the initial segment S_0(θ, ρ) and S_#((θ, ρ), (θ', ρ')) = k for every (θ', ρ') ∈ S_k(θ, ρ) and k ∈ℤ. Additionally, let S_a:b(θ, ρ) = ⋃_i = a, …, b S_i(θ, ρ) for a, b ∈ℤ such that a ≤ b.

A positive integer K and a user-defined weight function w: ℝ^4d→ [0, ∞) are required in AAPS.  The latter assigns weights to points within selected sets of segments. From the current point (θ, ρ), AAPS uniformly randomly chooses among the K+1 sets of segments: S_0:K, S_-1:K-1, …, S_-K:0. These sets of segments are all possible unions of K+1 consecutive segments whose union contains the current point. Having selected such a set of segments S_-c:K-c, AAPS then randomly selects a leapfrog iterate (θ', ρ') ∈ S_-c:K-c as a proposal with weight w((θ, ρ), (θ', ρ')). Finally, this proposal is Metropolized with Metropolis-Hastings acceptance probability

    1 ∧( e^-Δ H(θ,  ρ)  w((θ', ρ'), (θ, ρ))  ∑_(θ,  ρ) w((θ, ρ), (θ, ρ))/ w((θ, ρ), (θ',ρ'))  ∑_(θ,  ρ) w((θ', ρ'), (θ, ρ)) )  .

The weight function can be chosen to preferentially select proposals according to various desiderata. For instance, one might take w((θ,ρ), (θ',ρ')) = e^-H(θ', ρ') to select points which will always be accepted under the Metropolis-Hastings step, or w((θ, ρ), (θ',ρ')) = e^-H(θ', ρ') || θ - θ'||^2 to bias toward proposals which are farther from the current point. In addition to these, several other choices of weight functions are suggested in <cit.>. 

We are now in position to write AAPS as a GIST sampler. Let 𝒜 = ℤ^2, let η be the counting measure on 𝒜, and write the algorithm tuning parameter in components as α = (c, i) ∈𝒜. The component c represents the choice of set of segments S_-c:K-c while the component i represents the choice of leapfrog iterate Φ_h^i(θ, ρ) within this set of segments.  Define 

    p(c,i |θ, ρ) 
    = 1/K+1 w((θ, ρ), Φ_h^i(θ, ρ)) 
             1_S_-c:K-c(θ, ρ)(Φ_h^i(θ, ρ))/∑_j w((θ, ρ), Φ_h^j(θ, ρ)) 
              1_S_-c:K-c(θ, ρ)(Φ_h^j(θ, ρ)) ,
 
where c ∈{0, …, K }.  The map G in (<ref>) is given explicitly by 

    G : (θ, ρ, c, i) 
    ↦(𝒮∘Φ_h^i(θ, ρ),
     c + S_#((θ, ρ),  𝒮∘Φ_h^i(θ, ρ)), 
     i )  ,
 
and the corresponding acceptance probability in (<ref>) simplifies to (<ref>).  Thus, the GIST sampler in Figure <ref> with the above specification of the conditional distribution of the algorithm tuning parameter α = (c,i) given the state (θ, ρ) correspond to AAPS. 







§ RELATION TO DYNAMIC HMC SAMPLER




Since they both use U-turn avoiding conditions, it is natural to compare the path length adaptivity provided by the self-tuning HMC sampler to the one found in the well-established dynamic HMC sampler implemented in NUTS <cit.>. As we review below, dynamic HMC can be viewed as an instance of a generalized Hit-and-Run sampler in the sense of <cit.>. On the other hand, as discussed in Section <ref>, the self-tuning HMC sampler is a Metropolis-within-Gibbs sampler aimed at the enlarged target measure (<ref>).  As we highlight below, a key benefit of incorporating a Metropolis step in self-tuning HMC is that it substantially relaxes certain restrictive symmetry conditions required for correctness of dynamic HMC. 


Here we follow the terminology and framework introduced in the recent paper <cit.>.
From a given position and momentum (θ, ρ), dynamic HMC generates an orbit J ⊂ℤ according to an orbit selection kernel P_h(J |θ, ρ).  This set is an `orbit' insofar as it represents the indices of forward and backward leapfrog iterates from (θ, ρ).  In turn, dynamic HMC generates an index i ∈ J according to an index selection kernel Q_h(i | J, θ, ρ).  The updated position and momentum is then set to Φ_h^i(θ, ρ).  Under certain conditions on the orbit and index selection kernels, this sampler leaves invariant the target measure. In practice, an explicit formula for the orbit selection kernel is not needed to update the chain. Rather, the orbit J is generated as the output of an implementable randomized algorithm, which implicitly defines the orbit selection kernel.  

To leave invariant the target distribution, it is sufficient that the orbit and index selection kernels satisfy a `semi-global' balance condition in (4) of <cit.>.
Let 𝔽⊂ℤ be a finite subset of the integers containing 0 that specifies the number of forward/backward leapfrog iterates and ℘(𝔽) denote the corresponding power set.   Heuristically, the semi-global balance condition can be motivated by viewing dynamic HMC as a form of Gibbs sampling on the augmented space ℝ^2d×℘(𝔽) with distribution  
    μ_e(dθ, dρ, dJ, di) ∝ e^-H(θ, ρ) Q_h(i | J, θ, ρ) P_h(J |θ, ρ)   (m ⊗η) (dθ  dρ  dJ   di)
 
where P_h( ·|θ, ρ) is a distribution on ℘(𝔽), Q_h( ·| J, θ, ρ) is a distribution on J, and η is counting measure over the corresponding finite product space. From the state (θ, ρ) the algorithm chooses a set J ⊂𝔽 with probability P_h( J |θ, ρ) then selects (θ', ρ')  = Φ_h^i(θ, ρ) with probability Q_h(i | J, θ, ρ).  Selecting the index i in this way gives a Markov transition kernel on the set J and the semi global balance condition essentially provides a sufficient condition for a dynamic HMC sampler to preserve the enlarged target distribution.  

In particular, it suffices if one step of the index transition kernel Q_h(i | J, θ, ρ)  preserves the conditional stationary distribution of i given J. Ignoring normalizing constants, this is the same as requiring the index selection kernel preserve the stationary probability distribution on J induced by the weights  

    e^-H(Φ_h^i(θ, ρ)) P_h(J - i |Φ_h^i(θ, ρ))  ,   ∀ i ∈ J  .

This gives the semi-global balance condition presented in (4) of <cit.>. 

Designing index selection kernels which are both efficiently implementable and preserve the induced measure on orbits is not easy in general. Because the orbit selection kernel P_h(J |θ, ρ) is not known, computing the weights in (<ref>) is often neither straightforward nor efficient. This severely limits the otherwise straightforward process of producing a sample from a probability distribution supported on a finite set.

To remove this difficulty, all dynamic HMC samplers in current use impose the restriction P_h(J |θ, ρ) = P_h(J-i |Φ_h^i(θ', ρ')). If the orbit selection kernels satisfy this symmetry condition, then one can ignore the P_h(J-i |Φ_h^i(θ, ρ)) term in the weights in equation (<ref>). That is from a position θ, ρ we only need that the index selection kernel preserves the probability measure on J induced by the weights e^-H(Φ_h^i(θ, ρ))) - which is a much more straightforward task.

Consider how this symmetry condition arises in the orbit selection procedure implemented in NUTS <cit.>. At each stage, if the current orbit of length 2^k does not yet trigger the No-U-Turn condition the algorithm proposes extending the orbit either forward another 2^k points or backward another 2^k points.

Hence, at each stage the orbit corresponds to a sequence of binary choices selected uniformly at random. For any other point in the orbit, there is an equivalent set of binary choices which would have produced J-i. In order to maintain this symmetry property after extending the orbit, the algorithm must check that none of the sub-orbits in the proposed extension which correspond to alternative sequences of binary choices would have prematurely triggered the No-U-Turn condition. If a proposed extension contains any such undesirable sub-orbits, it must be rejected. 

The selected orbit J thus still corresponds to a sequence of binary choices. As a consequence of this construction, we still have the property that for every point i ∈ J in the selected orbit there is another equiprobable sequence of binary choices which would have given the orbit J-i. In other words, the orbit selection kernel satisfies the property P_h(J |θ, ρ) = P_h(J-i |θ, ρ).  (Compare this condition with a similar condition appearing in the proof of correctness of the slice sampler based on the stepping out procedure <cit.>.)

The construction of the dynamic HMC sampler is both elegant and delicate. By contrast, the Metropolis-within-Gibbs procedure in the Gibbs self-tuning HMC sampler is much simpler and does not require such a careful construction in order to be correct.  Owing to its simplicity, it offers significant flexibility to adapt other algorithm tuning parameters. 



[Dynamic HMC Sampler]
One can view the dynamic HMC sampler presented in <cit.> in terms of a self-tuning HMC algorithm (as in Figure  <ref>) but, with a slight modification.  As before, fix a step length h>0.  Let Φ_h: ℝ^2d→ℝ^2d denote one step of the leapfrog integrator operated at step length h. Let 𝒜 = ℤ, η be counting measure, and α be the # of steps. Let P_h(J |θ, ρ) and Q_h(i | J,θ, ρ) denote respectively the orbit and index selection kernels as defined in <cit.>. Define 

    p(α|θ, ρ)
    = 
    ∑_J ⊂ℤ1_J(α) 
     P_h(J |θ, ρ) 
     Q_h(α| J,θ, ρ)
 
and assume something similar to the symmetry condition (5) in <cit.>, e.g., 
    ∑_α∈ℤ e^-H(θ, ρ) p( α|θ, ρ) = 
    ∑_α∈ℤ  e^-H( 𝒮∘Φ_h^α  (θ, ρ)) p( α|𝒮∘Φ_h^α  (θ, ρ) )
 
for all (θ, ρ) ∈ℝ^2d.  In other words, the assumption is made that the phase space update satisfies a semi-global balance condition with respect to the enlarged target density.
Under this assumption, referring to the algorithm presented in Figure <ref>, the Metropolis accept/reject step is no longer needed for correctness.  The corresponding transition kernel (without the Metropolis accept/reject step) is that of the dynamic HMC sampler. To be sure, the self-tuning HMC algorithm relaxes the semi-global balance condition to a detailed balance condition that is enforced by a Metropolis accept/reject step.  This is a major simplification that offers significant flexibility for devising other adaptation strategies and adapting other algorithm tuning parameters.







§ TRUNCATION OF AN INFINITE-DIMENSIONAL GAUSSIAN
 

This section demonstrates that the exact self-tuned HMC algorithm described in Section <ref> can achieve a similar mean-squared jump distance (MSJD) as  randomized HMC from Section <ref> on a target distribution that can be interpreted as a truncation of an infinite-dimensional Gaussian measure <cit.>.  This is a worst-case example since the corresponding Hamiltonian dynamics is highly oscillatory <cit.>. More precisely, the target distribution is a d-dimensional centered normal distribution with covariance matrix given by 

    Σ = diag(σ_1, …, σ_d)  , 
      where  σ_i 
    = i/d  for i ∈{1, …, d } .

This example has traditionally been used to illustrate the importance of path length randomization to avoid slow mixing due to periodicities or near-periodicities in the  underlying Hamiltonian dynamics <cit.>.  Here we use this worst-case example to demonstrate the efficacy of the self-tuned HMC algorithms based on the U-turn conditions defined by (<ref>) and (<ref>). 

For this model, the exact solution to the corresponding Hamiltonian dynamics (θ_t, ρ_t) at time t ≥ 0 from initial condition (θ_0, ρ_0) is given in components by 

    θ_t^i = cos( t/σ_i) θ_0^i + σ_i sin( t/σ_i) ρ_0^i 
    
      and  ρ_t^i = - 1/σ_isin( t/σ_i) θ_0^i + cos( t/σ_i) ρ_0^i,
 
for i ∈{1, …, d }.  Using these solutions, the evaluation of the U-turn path lengths in (<ref>) or (<ref>) reduces to finding the first positive root of a scalar function of time, which itself crucially relies on a good initialization.  

To be precise, we provide a complete description of the initialization procedure for the case of the U-turn path length given in (<ref>); a similar initialization is used for (<ref>) and therefore omitted.  Our goal is to find the first positive root of f, i.e., 

    τ = inf{ t ≥ 0  :  f(t) = 0 } .
 
where we have defined the following function of time  

    f(t) = ∑_i=1^d f_i(t)  ,

where

    f_i(t) = -1/σ_isin( t/σ_i)  θ_0^i ρ_0^i + cos( t/σ_i) (ρ_0^i)^2  .

The first positive root of f_i can be computed analytically, i.e., 

    τ_i = arctan( ρ^i_0/θ^i_0)  + k^⋆π ,
 
where

    k^⋆ = min{ k ∈ℤ : arctan( ρ^i_0/θ^i_0)  + k π≥ 0  } .

The mean of {τ_i }_i=1^d is used to seed the root solver for finding τ.     
Using the above root finding procedure to evaluate τ per transition step, the numerically estimated mean-acceptance probability (mean a_e),  mean-squared jump distance (MSJD), and mean path length (mean τ) for d=1000 and using 10^5 transition steps is summarized in the following table.  



    mean a_e     MSD     mean τ 

randomized HMC     100%     429.81     1.00 

exact self-tuned HMC (<ref>)     97.4%     174.85      0.44 

exact self-tuned HMC (<ref>)     94.4%      573.15      1.16



Randomized HMC is operated using the optimal choice of mean path length for maximizing expected square jump distance in this example, which corresponds to the standard deviation of the least constrained coordinate <cit.>.  The self-tuned HMC algorithm based on (<ref>) slightly outperformed exact randomized HMC in terms of MSD, while the one based on (<ref>) performed worse.  Remarkably, neither degenerated in high dimension.  This high-dimensional example demonstrates both the efficacy of the U-turn conditions and the leniency of the corresponding acceptance probability in the self-tuned HMC algorithm.  



These numerical findings are of independent interest, since they motivate using GIST samplers based on U-turn avoiding conditions to sample from perturbed Gaussian measures on Hilbert spaces.   This class of target measures is relevant to several important applications including Path Integral Molecular Dynamics <cit.>, Transition Path
Sampling  <cit.>, and Bayesian statistical inverse problems for Hilbert spaces <cit.>. 
Since the corresponding Hamiltonian dynamics is potentially highly oscillatory in high modes <cit.>, in actual numerical implementations, preconditioning <cit.> or strongly stable integrators <cit.> are necessary to be able to choose the step size
independently of the dimension.  








§ PATH LENGTH SAMPLING TO AVOID U-TURNS
  

We now turn to an extended example of a novel U-turn avoiding sampler based on Gibbs self tuning.  We will focus on locally adapting the number of steps.  Thus for each HMC step, the self-tuned HMC algorithm will generate the number of steps probabilistically according to a distribution p(L |θ, ρ, ϵ, Σ).  We consider several such distributions based on the no U-turn condition in the next section and evaluate them in the following sections.




Figure <ref> provides pseudocode for a general algorithm adapting the number of leapfrog steps.  The ratio of the proposal joint density to the staring point joint density can be factored as

    p(θ^*, ρ^*,  L^* |ϵ, Σ)/ p(θ0, ρ0, L^* |ϵ, Σ) =   p(θ^*)/ p(θ0)· p(ρ^* |Σ)/ p(ρ0|Σ)· p(L^* |θ^*, ρ^*, ϵ, Σ)/ p(L^* |θ0, ρ0, ϵ, Σ),

where p(θ) = p(θ| y) is the target probability density function, p(ρ|Σ) = normal(0, Σ) is the momentum probability density function, and p(L |θ, ρ, ϵ, Σ) is the conditional step size probability mass function.




 §.§ Step distributions avoiding U-turns
 



To make our sampler concrete, a specific distribution over the number of steps must be defined.  We evaluate a few related choices, all of which are motivated by the observation driving the No U-Turn Sampler (NUTS) <cit.>, which is that it's wasteful to evaluate the Hamiltonian dynamics beyond the point at which the trajectory has made a U-turn and is heading back toward the starting position.  Figure <ref> illustrates the U-turn condition, which is made precise in Equation (<ref>).  

Let U(θ, ρ, ϵ, Σ) be the maximum number of leapfrog steps that can be taken before a U-turn, starting from (θ, ρ) and using step size ϵ and mass matrix Σ.  The point just before a U-turn can be defined for the discrete steps of the leapfrog integrator by

    U(θ0, ρ0, ϵ, Σ)
    = arg min_n ∈ℕ (θn - θ0)^⊤·ρn < 0,

where θn, ρn is the location in phase space after n leapfrog steps from θ0, ρ0 with step size ϵ and positive-definite mass matrix Σ.  

  
Figure <ref> shows a single step of the algorithm.  The number of steps is sampled conditionally based on the number of steps possible before a U-turn.  To preserve detailed balance, the trajectory from the selected point backward in time is evaluated until it makes a U-turn and the probability of selecting the initial state (i.e., the same number of steps) is used to balance the selection.  This requires sampling zero or more steps backward in time before the initial position.



 §.§ Conditional distribution of steps
 

In this section, we consider a few closely related approaches to generating the number of steps given the number of steps to a U-turn.  



  §.§.§ Steps generated uniformly


The most obvious choice for a conditional distribution over the number of steps is uniform between 1 and the number of steps before a U-turn,

    p(L |θ, ρ, ϵ, Σ)
    = uniform(L | 1,  U(θ, ρ, ϵ, Σ)),

where the bounds are read inclusively.



  §.§.§ Steps generated uniformly from later states


One of the techniques NUTS uses to take such long jumps on average is to bias the selection of a candidate toward the end of the Hamiltonian trajectory <cit.>.  To create a simple approximation to this, the uniform distribution can be restricted to the latter part of the trajectory by changing the lower bound from 1 to something greater.  Any number between 0 and M = U(θ, ρ, ϵ, Σ) is valid as a lower bound.  The choice of 1 corresponds to the uniform distribution of the previous section.  We will also evaluate lower bounds of ⌊1/2· M ⌋ and ⌊3/4· M ⌋, where ⌊ x ⌋ is the floor of x, which is largest integer less than or equal to x.  With smaller intervals, the proposed trajectory lengths will be longer, but the reversibility balance condition will reduce the acceptance probability.



  §.§.§ Binomial step generation


Non-uniform distributions may also be used.  For example, a binomial distribution for the number of steps could be defined for a fixed ψ∈ (0, 1) as 

    p(L |θ, ρ, ϵ, Σ)
        = binomial(L |ψ,  U(θ, ρ, ϵ, Σ)).

The expected value of L starting from (θ, ρ) is thus ψ·U(θ, ρ, ϵ, Σ).  We evaluated the binomial method and it winded up generating proposals that were too concentrated and thus hard to balance and accept, so we do not include results here.  We did not evaluate a more dispersed beta-binomial distribution due to the cost of the beta and gamma functions required for normalization.





§ EMPIRICAL EVALUATION


In this section, we evaluate the performance of our various proposals for step size adaptation and compare their performance to the state-of-the-art No U-Turn Sampler as coded in Stan <cit.>.  




 §.§ Models evaluated


The models considered are multivariate standard normal, ill-conditioned multivariate normal, correlated multivariate normal, eight schools meta-analysis, item-response theory 2 parameter logistic model, mixed effects Poisson regression, normal mixture, hidden Markov model, one-compartment pharmacokinetic / pharmacodynamic model (PK/PD), Lotka-Volterra population dynamics, autoregressive time series (AR), autoregressive moving average time series (ARMA), and generalized autoregressive conditional heteroskedasticity time series (GARCH).  All but the normal models are drawn from the  package <cit.>, but reparameterized according to Stan best practices so that NUTS can sample them.





 §.§ Learning curve


As a simple validation that the sampler is targeting the correct distribution, Figure <ref> plots a learning curve (expected absolute error) versus iteration for the uniform sampler with the full path for the eight schools model. The plot shows that error decreases as expected at a rate of 1 / √(n) and the efficiency is greater than that of independent draws for the parameters and slightly worse for the parameters squared in this simple case.




 §.§ Effect of step size and path fraction




In Figure <ref>, the performance of the uniform self-tuning sampler is shown for two step sizes, 0.36 (blue lines) and 0.18 (red lines), across a range of lower bound fractions (x axis) for uniform sampling.  The step size 0.36 is what NUTS adapted for an 80% average Metropolis acceptance probability (Stan's default); 0.18 is the step size for roughly 95% Metropolis acceptance.  Halving the step size roughly doubles the number of leapfrog steps taken, as shown in the upper left of the plot.  The remaining plots show that performance is better with a smaller step size.  

Mean square jump distance (MSJD) is also shown in Figure <ref>; it is defined for M steps of sampling by

    MSJD
    = 1/M∑_m = 1^M ||θm+1 - θm||_2^2.

The MSJD is a reasonable estimate of the expected squared jump distance, 

    ESJD
    = 𝔼[||θt+1 - θt||^2_2].

As the path fraction goes up, the MSJD goes up until the no-return rejection rate overtakes it and it begins to decrease. 

The rejection rate is broken down into total rejection rate, and then the number of rejections due to not being able to return to the origin before hitting a U-turn (see Figure <ref>).  Smaller step sizes do a better job at preserving the Hamiltonian and thus have lower rejection rates.  

Root mean square error (RMSE) is defined by

    RMSE = √(1/M∑_m=1^M ||θm - θ||_2^2),

where θ is the value of the parameter and the θ^(m) are MCMC draws.  If the draws were independent, the expected RMSE for parameters is 0.1 (the standard deviation, 1, divided by the square root of the number of i.i.d. draws, √(100)).  The expected RMSE for parameters squared is √(2) / 10 because standard normal variates squared follow a chi-squared distribution with 1 degree of freedom, the standard deviation of which is √(2).  The RMSE for parameters is below that of independent draws, going down to roughly half at lower bound fraction 0.6.  Not surprisingly, because MSJD tracks inverse lag-one autocorrelation, the best RMSE for parameters is achieved where the MSJD is maximized.  The RMSE for parameters squared is mazimized at around a 0.35 lower-bound fraction, though values are very close between 0 and 1/2.  With a single lower bound fraction, a trade off must be made, as it must be for HMC and NUTS.  Here a fraction of 0.6 appears reasonable.  The RMSE is lower for the smaller step sizes, but doubling the number of iterations with the larger step size would reduce expected RMSE by a factor of 1 / √(2) (about 70%).




 §.§ Evaluations for multiple models




We run all of our reference models for 10,000 warmup and 40,000 iterations of NUTS using Stan's defaults and use the sampling draws to define reference means for parameters and parameters squared.  The results for the uniform sampler are shown in Figure <ref>.  There are four plots, which show RMSE for parameters, RMSE for squared parameters, MSJD, and leapfrog steps.  

We evaluate NUTS (in red) and three settings of uniform sampling, S(hort), which uses uniform(1, L) sampling, where L is the number of steps to a U-turn, M(edium), which uses uniform(⌊ 0.3 · L ⌋, L), and L(ong), which uses uniform(⌊ 0.5 · L ⌋, L), and XL (extra long), which uses uniform(⌊ 0.7 · L, L) sampling.  Two step sizes are considered, (1) the step size adapted by NUTS with its default acceptance target (80%), and (1/2), which is half of that step size, which is approximately what NUTS would adapt for a target of 95% step size, which is recommended for the numerical stability of the leapfrog integrator in more difficult problems.

Starting with leapfrog steps (the bottom plot), it is clear that the shorter expected paths require more steps for self-tuning HMC, because they are likely to have to go further past the origin backward in time before hitting a U-turn.  There is never more than a factor of two difference among systems in terms of leapfrog steps.  The steps required for NUTS and ST-HMC are similar; for some models, ST-HMC requires fewer steps than NUTS and for others it requires more steps. For the three normal cases, ST-HMC requires fewer steps for the standard normal and correlated normal, but about the same number for the ill-conditioned normal.

Smaller step sizes lead to roughly twice as many leapfrog steps, but not quite that in complicated models (or models that run up against the 1024 maximum step limit imposed on both NUTS and ST-HMC).

The second plot from the bottom shows the MSJD for all samplers, which is inversely related to autocorrelation at lag 1.  The clear trend is that larger steps lead to lower error, but this is not always borne out.  For instance, in the standard normal, the ST-HMC sampler takes longer steps and improves the estimate of parameters, but NUTS has lower error estimating parameters squared.

The first two plots show RMSE for parameters and parameters squared.  All plots have a y-axis starting at 0 so that relative performance may be assessed visually.  Many of the more complicated models are neck and neck in RMSE for parameters and parameters squared, with NUTS consistently slightly outperforming ST-HMC.  

In conclusion, NUTS slightly outperforms ST-HMC overall.  This is highly encouraging in that this is a relatively simple first attempt at self-tuned HMC, whereas NUTS is represented by Stan's current implementation, which has been implemented robustly and improved considerably over the first NUTS paper (and first version in Stan) <cit.>.




 §.§ Open-source code and reproduciblity


The results and plots can be reproduced from our code distribution, which is available from GitHub under a permissive open-source license.[The code is distributed under the MIT License at <https://github.com/bob-carpenter/adaptive-hmc>.]



§ OTHER RELATED PRIOR WORK


The GIST sampler may be viewed as a dynamic version of the empirical HMC sampler <cit.>. This empirical HMC sampler learns an empirical distribution of the underlying path lengths to U-turns during a warmup phase and then fixes this empirical distribution, which is then used to sample path lengths at each step during sampling. Like randomized HMC in Section <ref>, the empirical HMC sampler is an instance of the GIST sampler where the conditional distribution of path lengths does not depend on the current position and/or momentum.

Both NUTS and the Apogee-to-Apogee Path Sampler reviewed in Section  <ref> can be formulated within the dynamic HMC framework introduced in <cit.>. Dynamic HMC can be viewed as an instance of a generalized Hit-and-Run sampler in the sense of <cit.>. On the other hand, as discussed in Section <ref>, the GIST sampler is a Gibbs sampler (with a Metropolis-within-Gibbs step) aimed at the enlarged target measure (<ref>).  A key benefit of incorporating a Metropolis-within-Gibbs step in GIST is that it relaxes the restrictive symmetry conditions required for correctness of dynamic HMC in practice, and hence, offers significant flexibility to adapt other algorithm tuning parameters.

The autoMALA sampler introduced a self-tuning version of the Metropolis-adjusted Langevin (MALA) sampler <cit.>.  It uses a forward and reverse non-deterministic scheme to choose adaptation parameters in a way that satisfies detailed balance. MALA is equivalent to one-step Hamiltonian Monte Carlo, but is simpler in not needing to evolve momentum.  The GIST sampler can be viewed as a probabilistic generalization of the autoMALA adaptation selection criteria.

For step size and mass matrix tuning, two approaches have been popular in practice. In the first, an adaptation phase is used to estimate algorithm tuning parameters. These tuning parameters are then fixed so that the resulting chain is Markovian. This is the strategy used by NUTS for step size and mass matrix adaptation for HMC <cit.>. In the second approach, the adaptation phase is never turned off, but the amount of adaptation is decreased so that asymptotically the results are valid.  This is the strategy used by delayed rejection Metropolis (DRAM) <cit.>.

Another approach for step size adaptivity is the recently introduced delayed rejection HMC sampler of <cit.>.  The delayed rejection algorithm <cit.> is a generalization of Metropolis-Hastings to a sequence of proposal moves.  This proposal sequence can be tuned to start with larger scale moves and then scaled down for subsequent proposals <cit.>. In the same vein, the delayed rejection HMC sampler automatically tries smaller step sizes if proposals with larger step sizes are rejected, allowing it to sample from target densities with varying scale. Like other delayed rejection methods, it generates “ghost points” using reversed proposals, which are then used as part of the acceptance probability to ensure detailed balance <cit.>.




 §.§ Acknowledgements


We would like to thank Edward Roualdes, Tore Kleppe, Andreas Eberle, Sam Livingstone, and Chirag Modi for feedback on the Gibbs self-tuning idea.

N. Bou-Rabee has been partially supported by NSF Grant No. DMS-2111224.








§ PROOF OF CORRECTNESS


The correctness of the GIST sampler given in Figure <ref>  follows from the observation that the GIST sampler is a Gibbs sampler which interleaves Gibbs resampling of the velocity variable ρ and tuning parameter α with a Metropolis-within-Gibbs step that uses a proposal given by a  measure-preserving involution. In fact, this idea is also the essence of the standard HMC sampler, and the GIST sampler is a natural extension of this idea to the selection of tuning parameters as well.

The proof of correctness relies on the following general lemma. This lemma is not new; see, e.g.,  special case 2 of Theorem 2 in <cit.>.  However, a complete proof is given for the reader's convenience.

  
Suppose that (𝕊, ℱ, λ) is a measure space, ν is a probability measure with strictly positive density γ with respect to λ, and F: 𝕊→𝕊 is a measurable involution which preserves λ. Given a state Z ∈𝕊, define a Markov transition kernel by the following procedure:

    
  * Draw 𝒰∼uniform([0,1]).
    
  * Set Z' = F(Z) if  𝒰≤min(1, γ(F(Z))/γ(Z)) and Z' = Z otherwise.

Then, this transition step is reversible with respect to ν.




Let Z ∼ν.  The proof shows that 
    (Z  ∈ A , Z' ∈ B) =  (Z'  ∈ A , Z ∈ B)  ,
 
for all measurable sets A,B ∈ℱ.
Let β(x,F(x)) = min(1, γ(F(x))/γ(x)) for all x ∈𝕊. For all x ∈𝕊 and for any measureable set A ∈ℱ, the  transition kernel associated to the transition step described in Lemma <ref> is given by 
    p(x, A) = β(x,F(x)) δ_F(x)(A) + (1 - β(x, F(x))) δ_x(A)  .
 Hence,

    (     Z  ∈ A , Z' ∈ B) = ∫_A p(x,B) ν(dx) 
       = ∫_A [ β(x, F(x)) 1_B(F(x)) + (1 - β(x,F(x))) 1_B(x) ] ν(dx) 
       = ∫_𝕊[ β(x, F(x)) 1_B(F(x)) 1_A(x) + (1 - β(x,F(x)))1_A(x)1_B(x) ] ν(dx)  .

The second term in (<ref>) is already symmetric in the sets A,B — so we turn our attention to the first term. For this term, we use the elementary identity a min(1, b/a) = b min(1, a/b) valid for all a, b  0, as follows  
    ∫_𝕊   β(x, F(x)) 1_B(F(x)) 1_A(x) ν(dx)  
       = ∫_𝕊β(x, F(x)) 1_B(F(x)) 1_A(x) γ(x) λ(dx) 
       = ∫_𝕊1_B(F(x)) 1_A(x) γ(x) min(1, γ(F(x))/γ(x)) λ (dx) 
       = ∫_X 1_B(F(x)) 1_A(x) γ(F(x)) min(1, γ(x)/γ(F(x))) λ (dx)  .

But, since the map F preserves λ and is an involution we have, by change of variables under F,

    ∫_𝕊   1_B(F(x)) 1_A(x) γ(F(x)) min(1, γ(x)/γ(F(x))) λ (dx) 
        = ∫_𝕊1_B(F(F(x))) 1_A(F(x)) γ(F(F(x))) min(1, γ(F(x))/γ(F(F(x)))) λ (dx) 
       =∫_𝕊1_B(x) 1_A(F(x)) γ(x) min(1, γ(F(x))/γ(x)) λ (dx) 
       =∫_𝕊β(x, F(x)) 1_B(x) 1_A(F(x))  ν (dx) =∫_B β(x, F(x))  1_A(F(x))  ν (dx)  .


Inserting (<ref>) into (<ref>), and comparing with (<ref>), we observe that

    (Z ∈ A , Z' ∈ B) = (Z ∈ B , Z' ∈ A)

as required.


With Lemma <ref> in hand, we are now in position to prove correctness of the GIST sampler. 



It is clear that the iterates θ_0, θ_1, … generated by the GIST sampler in Figure <ref> form a Markov chain, so we need only establish reversiblity. 
In the notation of Figure  <ref>, suppose  θ_0  = θ∼μ_θ. After selecting ρ_0 ∼normal(0, I_d × d)  and α∼ p( ·|θ_0, ρ_0), then (θ_0, ρ_0, α) ∼μ_e. 

Let u ∼uniform([0,1]) and set 

    (θ^*, ρ^*, α^*) =
    G(θ_0, ρ_0, α)    if   u ≤ e^-Δ H(θ_0, ρ_0)p( π(θ_0, ρ_0)(α) | S ∘ F(α)(θ_0, ρ_0) )p(α|θ_0, ρ_0) ,  
    
    (θ_0, ρ_0, α)    else .




Since G is a measure-preserving involution by assumption, by Lemma <ref> we have 
    ℙ((θ_0, ρ_0, α) ∈ A, (θ^*, ρ^*, α^*) ∈ B)  = ℙ( (θ^*, ρ^*, α^*) ∈ A, (θ_0, ρ_0, α) ∈ B)

for any measurable A,B ⊂ℝ^2d×𝒜. In particular, for Borel sets A, B⊂ℝ^d,

    ℙ( θ∈A, θ^* ∈B) = ℙ(θ^* ∈A, θ∈B)

by taking A = A×ℝ^d×𝒜 and B = B×ℝ^d ×𝒜 in (<ref>).  Hence, the GIST sampler produces a  Markov Chain that is reversible with respect to  e^-U(x) m(d θ).






As a corollary, Theorem <ref> implies correctness of the GIST samplers presented in Sections <ref> and <ref>, as well as the GIST sampler in Section <ref>.


If π(θ, ρ)(α) = α for all (θ, ρ, α) ∈ℝ^2d×𝒜, the corresponding GIST sampler is correct. 



We need only verify the map G: ℝ^2d×𝒜→ℝ^2d×𝒜 defined by G(θ, ρ, α) = (𝒮∘ F(α)(θ, ρ), α) is a measure-preserving involution. As 𝒮∘ F(α) is an involution G is automatically an involution on ℝ^2d×𝒜. Additionally 𝒮∘ F(α) preserves Lebesgue measure m on ℝ^2d for every fixed α.  Hence, by Fubini's theorem, G preserves (m ⊗η) on ℝ^2d×𝒜.

Thus, iterating the transition step in Figure <ref> produces a Markov chain reversible with respect to μ_θ by Theorem <ref>.


The next two corollaries of Theorem <ref> show the correctness of the GIST samplers in Section <ref> and Section <ref>. 


The No-U-Turn Sampler presented in Section <ref> is correct. 



In this case, the map G is of the form G(θ, ρ, J, i) = (S ∘Φ_h^i(θ, ρ), -(J-i), i). This is clearly an involution. Moreover, for θ, ρ fixed the map (J,i) → (-(J-i), i) is a bijection and thus preserves the counting measure on 𝒫_F ×ℤ. By Fubini's theorem G then preserves (m ⊗η) on ℝ^2d×𝒫_F ×ℤ.

Thus, by Theorem <ref>, we get correctness of NUTS as presented in Section <ref>. 



The Apogee-to-Apogee Path Sampler presented in Section <ref> is correct.



Here, we take G(θ, ρ, c, i)= (𝒮∘Φ_h^i(θ, ρ), c + S_#((θ, ρ), 𝒮∘Φ_h^i(θ, ρ)), i ). Since by definition S_#((θ, ρ), (θ', ρ')) = - S_#((θ', ρ'), (θ, ρ)), we observe by direct computation that G is an involution. 

For fixed (θ, ρ, i), the map c ↦ c + S_#((θ, ρ), 𝒮∘Φ_h^i(θ, ρ)) is a bijection and hence preserves counting measure. Additionally, for i fixed, θ, ρ↦𝒮∘Φ_h^i(θ, ρ) preserves Lebesgue measure.

Applying Fubini's theorem then implies that the map G is measure preserving. Hence, by Theorem <ref>, the Apogee-to-Apogee Path Sampler presented in Section <ref> is correct. 


In the description above, the proposal G(θ_0, ρ_0, α) has phase space state given by 𝒮∘ F(α)(θ_0, ρ_0) whereas the proposal in the original algorithm is F(α)(θ_0, ρ_0). This difference is inconsequential since the acceptance probabilities are the same and the map 𝒮 affects only the velocity variable. Since this variable is immediately discarded in the next step of the chain, we disregard this difference. Equivalently, one could adapt the description above so that 𝒮 is applied to the phase space coordinates immediately after Metropolizing the proposal. Since 𝒮 preserves the extended target distribution μ_e, the whole procedure will still preserve the extended target.





 §.§ Reduction to standard Metropolis-Hastings


It is worth noting that the above yields the original Metropolis-Hastings algorithm as a special case. Let (X, ℱ, λ) be a measure space with σ-finite measure λ. Suppose μ(dx) = p_0(x) λ(dx) is an absolutely continuous target probability measure and P(x, dy) = p(x,y) λ(dy) is an absolutely continuous transition kernel with strictly positive  densities. Analogously to (<ref>), define the following enlarged target measure on X × X: 
    μ_e(dx   dy) = γ(x,y) λ(dx   dy) = p_0(x) p(x,y) (λ⊗λ)(dx   dy)  .
 
Note that the marginal of μ_e in the first component gives the desired target.  



On X × X, and from a state with first component x ∈ X, first generate a sample from the second component y  conditional on the first, i.e.,  

    y ∼ p(x, ·)  .

Then, apply the Metropolis procedure to the involutive proposal ϕ: (x,y) ↦ (y,x) which clearly preserves λ⊗λ.
According to Lemma <ref>, the corresponding acceptance probability is given by

    a_e(x,y) = 1 ∧(γ(ϕ(x,y))/γ(x,y)) = 1 ∧(p_0(y) p(y,x)/p_0(x) p(x,y))   ,

which we recognize as the standard Metropolis-Hastings acceptance probability.   





§ TEST MODELS


In this section, we briefly describe the models used for evaluations.  Greek letters are used for parameters, roman letters for constants, predictors and modeled data, and italics for indexes.  Where not otherwise specified, parameters have weakly informative priors concentrated on their rough scale.



  Standard normal
A 500-dimensional standard normal, with α∼normal(0, I), where I is a 500 × 500 identity matrix.


  Correlated normal
A 250-dimensional correlated normal α∼normal(0, S), where the covariance is that of a unit scale first-order random walk, S_i, j = r^|i - j|, with correlation r = 0.9.


  Ill-conditioned normal
A 250-dimensional ill-conditioned normal, with α∼normal(0, diag-matrix(σ)), where σ = [ 1/250  2/250⋯250/250]^⊤.


  Poisson generalized linear mixed model
The expected count for observation i is
λ_i =  exp( α + β_1 ·t_i + β_2 ·t_i^2 + β_3 ·t_i^3 + ε_i),
with a hierarchical prior
ε_i ∼normal(0, σ),
and a likelihood
y_i ∼Poisson(λ_i),
for i < 40.


  Eight-schools meta-analysis Observations of mean effects y_j and standard errors σ_j in school j are used to estimate a school effect θ_j with likelihood y_j ∼normal(θ_j, σ_j) and a non-centered parameterization of the prior θ_j ∼normal(μ, τ).


  Order K autoregressive A K-th order autoregressive model with likelihood y_t ∼normal(α + [y_t-1 y_t-2⋯ y_t-K] ·β, σ).  The prior does not restrict coefficient vectors β to ensure stationary.  A history of K=5 and T=200 time points is used in the evaluations.


  Autoregressive moving average (order 1, 1)
An autoregressive time series with first element y_1 ∼normal(μ + ϕ·μ, σ), and subsequent elements y_t ∼normal(μ + ϕ· y_t-1 + θ·ϵ_t-1, σ) for t > 1.


  Generalized autoregressive conditional heteroskedasticity 
An autoregressive time series model with stochastic volatility, with y_t ∼normal(μ, σ_t), where σ_1 = s is given as data and σ_t = √(α_0 + α_1 · (y_t-1 - μ)^2 + β_1 ·σ_t-1^2).


  Hidden Markov model
A hidden Markov model (HMM) with normal observations.  The data generating process is Markovian in hidden state, z_t ∼categorical(ϕ_z_t-1), and then normal in observation, y_t ∼normal(μ_z_t, σ_z_t).  In the implementation, the forward algorithm is used to marginalize the z_t to calculate the likelihood.  The vector μ is constrained to ascending order for identifiability.


  Hierarchical item-response theory 2 parameter logistic
For test item i. a student j's correctness is generated by y_i, j∼bernoulli(logit^-1(α_i · (θ_j - β_i))), for ability θ, difficulty β, and discrimination α.  For identifiability, β has a standard hierarchical prior, β_i ∼normal(μ^β, σ^β), whereas the prior for α is centered, α_i ∼normal(0, σ^α), and the prior for θ is fixed to unit scale, θ∼normal(0, 1).  


  Normal mixture
A normal mixture model with z_n ∼bernoulli(θ), and y_n ∼normal(μ_z_n, σ_z_n), with μ constrained to ascending order for identifiability.  The responsiblity parameter z_n is marginalized out.





  Lotka-Volterra population dynamics

A lognormal model of population time series for prey (y_t,1) and predator (y_t,2).  The population dynamics is modeled by a system of ordinary differential equations, d/dt u(t) = (α - β· v(t)) · u(t) and d/dt v(t) = (-γ + δ· u(t)) · v(t) with unknown starting point (u(0), v(0)) and discrete observations modeled by y_t, 1∼lognormal(log u(t), σ_1) and y_t, 2∼lognormal(log v(t), σ_2),  where u(t) and v(t) are solutions to the ODE.


  Pharmacometric/pharmacodynamic model

A one-compartment PK/PD elimination model with Michaelis-Mentin dynamics defined by a differential equation for concentration with dosing (D) of 
d/dt C
= ν/V·C(t)/μ + C(t)
- exp(-δ· t) · D ·δ/V, with lognormal likelihood for discrete observations
y_t ∼lognormmal(log C(t), σ), with all parameters constrained to be positive.



