Entropy-based adaptive Hamiltonian Monte Carlo

1 Introduction . Consider the problem of sampling from a target density π on Rd of the form π ( q ) ∝ e−U ( q ) , with a potential energy U : Rd → R being twice continuously differentiable . HMC methods [ 20 , 46 , 9 ] sample from a Boltzmann-Gibbs distribution µ ( q , p ) ∝ e−H ( q , p ) on the phase-space R2d based on the ( separable ) Hamiltonian function H ( q , p ) = U ( q ) +K ( p ) with K ( p ) = 1 2 p > M−1p . The Hamiltonian represents the total energy that is split into a potential energy term U and a kinetic energy K which we assume is Gaussian for some symmetric positive definite mass matrix M . Suppose that ( q ( t ) , p ( t ) ) t∈R evolve according to the differential equations dq ( t ) dt = ∂H ( q ( t ) , p ( t ) ) ∂p = M−1p ( t ) and dp ( t ) dt = −∂H ( q ( t ) , p ( t ) ) ∂q = −∇U ( q ( t ) ) . ( 1 ) Let ( ϕt ) t > 0 denote the flow of the Hamiltonian system , that is for fixed t , ϕt maps each ( q , p ) to the solution of ( 1 ) that takes value ( q , p ) at time t = 0 . The exact HMC flow ϕ preserves 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . volume and conserves the total energy i.e . H ◦ ϕt = H . Consequently , the Boltzmann-Gibbs distribution µ is invariant under the Hamiltonian flow , that is µ ( ϕt ( E ) ) = µ ( E ) for any Borel set E ⊂ R2d . Furthermore , the flow satisfies the generalized reversibility condition F ◦ ϕt = ϕ−t ◦ F with the flip operator F ( q , p ) = ( q , −p ) . Put differently , the Hamiltonian dynamics go backward in time by negating the velocity . If an analytical expression for the exact flow were available , one could sample from µ using the invariant Markov chain that at state ( q , p ) first draws a new velocity p′ ∼ N ( 0 , M ) with the next state set to ϕT ( q , p′ ) for some integration time T > 0 . Such a velocity refreshment is necessary as the HMC dynamics preserve the energy and so can not be ergodic . However , the Hamiltonian flow can not be computed exactly , except for very special potential functions . Numerical approximations to the exact solution of Hamiltonian ’ s equations are thus routinely used , most commonly the leapfrog method , also known as ( velocity ) Verlet integrator [ 28 , 10 ] . For a step size h > 0 and L steps , such an algorithm updates the previous state q0 and a new velocity p0 ∼ N ( 0 , M ) by setting , for 0 6 ` 6 L− 1 , p ` + 12 = p ` − h 2 ∇U ( q ` ) ; q ` +1 = q ` + hM−1p ` + 12 ; p ` +1 = p ` + 12 − h 2 ∇U ( q ` +1 ) . ( 2 ) This scheme can be motivated by splitting the Hamiltonian wherein the kick mappings in the first and third step update only the momentum , while the drift mapping in the second step advances only the position q with constant speed . For T = Lh , the leapfrog integrator approximates ϕT ( q0 , p0 ) by ( qL , pL ) while also preserving some geometric properties of ϕ , namely volume preservation and generalized reversibility . The leapfrog method is a second-order integrator , making anO ( h2 ) energy error H ( qL , pL ) −H ( q0 , p0 ) . A µ-invariant Markov chain can be constructed using a MetropolisHastings acceptance step . More concretely , the proposed state ( qL , pL ) is accepted with the acceptance rate a ( q0 , p0 ) = min { 1 , exp [ − ( H ( qL , pL ) −H ( q0 , p0 ) ) ] } , while the next state is set to F ( q0 , p0 ) in case of rejection , although the velocity flip is inconsequential for full refreshment strategies . We want to explore here further the generalised speed measure introduced in [ 54 ] for adapting RWM or MALA that aim to achieve fast convergence by constructing proposals that ( i ) have a high average log-acceptance rate and ( ii ) have a high entropy . Whereas the entropy of the proposal in RWM or MALA algorithms can be evaluated efficiently , the multi-step nature of the HMC trajectories makes this computation less tractable . The recent work in [ 41 ] consider the same adaptation objective by learning a normalising flow that is inspired by a leapfrog proposal with a more tractable entropy by masking components in a leapfrog-style update via an affine coupling layer as used for RealNVPs [ 19 ] . [ 60 ] sets the integration time by maximizing the proposal entropy for the exact HMC flow in Gaussian targets , while choosing the mass matrix to be the inverse of the sample covariance matrix . 2 Related work . The choice of the hyperparameters h , L and M can have a large impact on the efficiency of the sampler . For fixed L and M , a popular approach for adapting h is to target an acceptance rate of around 0.65 which is optimal for iid Gaussian targets in the limit d→∞ [ 8 ] for a given integration time . HMC hyperparameters have been tuned using some form of expected squared jumping distance ( ESJD ) [ 49 ] , using for instance Bayesian optimization [ 56 ] or a gradient-based approach [ 40 ] . A popular approach suggested in [ 32 ] tunes L based on the ESJD by doubling L until the path makes a U-turn and retraces back towards the starting point , that is by stopping to increase L when the distance to the proposed state reaches a stationary point [ 4 ] ; see also [ 57 ] for a variation and [ 48 ] for a version using sequential proposals . Modern probabilistic programming languages such as Stan [ 12 ] , PyMC3 [ 51 ] , Turing [ 23 , 58 ] or TFP [ 39 ] furthermore allow for an adaptation of a diagonal or dense mass-matrix within NUTS based on the sample covariance matrix . The Riemann manifold HMC algorithm from [ 25 ] has been suggested that uses a position dependent mass matrix M ( x ) based on a non-separable Hamiltonian , but can be computationally expensive , requiring O ( d3 ) operations in general . An alternative to choose M or more generally the kinetic energy K was proposed in [ 43 ] by analysing the behaviour of x 7→ ∇K ( ∇U ( x ) ) . Different pre-conditioning approaches have been compared for Gaussian targets in [ 38 ] . A popular route has also been to first transform the target using tools from variational inference as in [ 31 ] and then run a HMC sampler with unit mass matrix on the transformed density with a more favourable geometry . A common setting to study the convergence of HMC assumes a log-concave target . In the case that U is m1-strongly convex and m2-smooth , [ 45 , 15 ] analyse the ideal HMC algorithm with unit mass matrix where a higher condition number κ = m2/m1 implies slower mixing : The relaxation time , i.e . the inverse of the spectral gap , grows linear in κ , assuming the integration time is set to T = 1 2 √ m2 . [ 14 ] establish non-asymptotic upper bounds on the mixing time using a leap-frog integrator where the step size h and the number L of steps depends explicitly on m1 and m2 . Convergence guarantees are established using conductance profiles by obtaining ( i ) a high probability lower bound on the acceptance rate and ( ii ) an overlap bound , that is a lower bound on the KL-divergence between the HMC proposal densities at the starting positions q0 and q′0 , whenever q0 is close to q ′ 0 . While such bounds for controlling the mixing time might share some similarity with the generalised speed measure ( GSM ) considered here , they do not lend themselves easily to a gradient-based adaptation . 3 Entropy-based adaptation scheme . We derive a novel method to approximate the entropy of the proposed position after L leapfrog steps . Our approximation is based on the assumption that the Hessian of the target is locally constant around the mid-point of the HMC trajectory , which allows for a stochastic trace estimator of the marginal proposal entropy . We develop a penalised loss function that can be minimized using stochastic gradient descent while sampling from the Markov chain in order to optimize a GSM . 3.1 Marginal proposal entropy . Suppose that CC > = M−1 , where C is defined by some parameters θ and can be a diagonal matrix , a full Cholesky factor , etc . Without loss of generality , the step size h > 0 can be fixed . We can reparameterize the momentum resampling step p0 ∼ N ( 0 , M ) by sampling v ∼ N ( 0 , I ) and setting p0 = C− > v . One can show by induction , cf . Appendix E for details , that the L-th step position qL and momentum pL of the leapfrog integrator can be represented as a function of v via qL = TL ( v ) = q0 − Lh2 2 M−1∇U ( q0 ) + LhCv − h2M−1ΞL ( v ) , ( 3 ) and pL =WL ( v ) = C− > v − h 2 [ ∇U ( q0 ) +∇U ◦ TL ( v ) ] − h L−1∑ i=1 ∇U ◦ Ti ( v ) ( 4 ) where ΞL ( v ) = L−1∑ i=1 ( L− i ) ∇U ◦ Ti ( v ) , ( 5 ) see also [ 42 , 21 , 14 ] for the special case with an identity mass matrix . Observe that for L = 1 leap-frog steps , this reduces to a MALA proposal with preconditioning matrix M−1 . Under regularity conditions , see for instance [ 21 ] , the transformation TL : Rd → Rd is a C1diffeomorphism . With ν denoting the standard Gaussian density , the density rL of the HMC proposal for the position qL after L leapfrog steps is the pushforward density of ν via the map TL so that1 log rL ( TL ( v ) ) = log ν ( v ) − log |detDTL ( v ) | . ( 6 ) Observe that the density depends on the Jacobian of the transformation TL : v 7→ qL . We would like to avoid computing log |detDTL ( v ) | exactly . Define the residual transformation SL : Rd → Rd , v 7→ 1 Lh C−1TL ( v ) − v. ( 7 ) Then DTL ( v ) = LhC ( I +DSL ( v ) ) and consequently log |detDTL ( v ) | = d log ( Lh ) + log |detC|+ log |det ( I +DSL ( v ) ) | . ( 8 ) Combining ( 6 ) and ( 8 ) yields the log-probability of the HMC proposal log rL ( TL ( v ) ) = log ν ( v ) − d log ( Lh ) − log |detC| − log |det ( I +DSL ( v ) ) | . ( 9 ) 1We denote the Jacobian matrix of a function f : Rd → Rd at the point x as Df ( x ) . Comparing the equations ( 3 ) and ( 7 ) , one sees that SL ( v ) = c − hLC > ΞL ( v ) for some constant c ∈ Rd that depends on θ but is independent of v and consequently , DSL ( v ) = − hLC > DΞL ( v ) . We next show a recursive expression for DSL with a proof given in Appendix B. Lemma 1 ( Jacobian representation ) . It holds that DS1 = 0 and for any ` ∈ { 2 , . . . , L } , v ∈ Rd , DS ` ( v ) = −h2 ` −1∑ i=1 ( ` − i ) i ` C > ∇2U ( Ti ( v ) ) C ( I +DSi ( v ) ) . ( 10 ) In particular , DS ` ( v ) is a symmetric matrix . Suppose further that L2h2 < supq∈Rd 14‖C > ∇2U ( q ) C‖ 2 . Then for any ` ∈ { 1 , . . . , L } and v ∈ Rd , we have ‖DS ` ( v ) ‖2 < 1 8 . Notice that the recursive formula ( 10 ) requires computing 12L ( L−1 ) terms , each involving the Hessian , in order to compute the Jacobian after L leapfrog steps . Consider for the moment a Gaussian target with potential function U ( q ) = 12 ( q − q ? ) > Σ−1 ( q − q ? ) for q ? ∈ Rd and positive definite Σ ∈ Rd×d . Then , due to ( 10 ) , for any q ∈ Rd , v ∈ Rd , DSL ( v ) = −h2 L−1∑ i=1 ( L− i ) i L C > Σ−1C ( I +DSi ( v ) ) = DL +RL ( v ) , where DL = −h2C > Σ−1C ( L−1∑ i=1 ( L− i ) i L ) = −h2L 2 − 1 6 C > Σ−1C ( 11 ) and a remainder term RL ( v ) = −h2C > Σ−1C ( ∑L−1 i=1 ( L− i ) i LDSi ( v ) ) . From Lemma 1 , we see that if ∥∥C > Σ−1C∥∥ 2 6 14h2L2 , then I +DSL ( v ) and −DSL ( v ) are positive definite . Then RL is also positive definite and log det ( I +DL ) 6 log |det ( I +DSL ( v ) ) | and we can maximize the lower bound instead . Put differently , for Gaussian targets , DSL can be decomposed into a component DL that contains all terms that are linear in h2C > Σ−1C and that does not require a recursion ; plus a component RL that contains terms that are higher than linear in h2C > Σ−1C and that needs to be solved recursively . Our suggestion is to ignore this second term . Note that R2 = 0 and an extension can be to include higher order terms O ( [ h2C > Σ−1C ] k ) , k > 1 , in the approximation DL . For an arbitrary potential energy U , equation ( 10 ) shows that evaluating DSL leads to a non-linear function of the Hessians evaluated along the different points of the leapfrog-trajectory . We suggest to replace it with a first order term with one Hessian evaluation which is however scaled accordingly . Concretely , we maximize L ( θ ) = log |det ( I +DL ) | with DL = −h2 L2 − 1 6 C > ∇2U ( qbL/2c ) C ( 12 ) as an approximation of log |det ( I +DSL ) | . The intuition is that we assume that the target density can be approximated locally by a Gaussian one with precision matrix Σ−1 in ( 11 ) given by the Hessian of U at the mid-point qbL/2c of the trajectory . We want to optimize L ( θ ) given in ( 12 ) even if we do not have access to the Hessian ∇2U explicitly , but only through Hessian-vector products ∇2U ( q ) w for some vector w ∈ Rd . Vector-Jacobian products vjp ( f , x , w ) = w > Df ( x ) for differentiable f : Rd → Rd can be computed efficiently via reverse-mode automatic differentiation , so that∇2U ( q ) w = vjp ( ∇U , q , w ) > can be evaluated with complexity linear in d. Suppose the multiplication with DL is a contraction so that all eigenvalues of DL have absolute values smaller than one . Then one can apply a Hutchinson stochastic trace estimator of log |det ( Id +D , L ) |with a Taylor approximation , truncated and re-weighted using a Russian-roulette estimator [ 44 ] , see also [ 29 , 5 , 13 ] for similar approaches in different settings . More concretely , let N be a positive random variable with support on N and let pk = P ( N > k ) . Then , L ( θ ) = log det ( I +DL ) = EN , ε [ N∑ k=1 ( −1 ) k+1 kpk ε > ( DL ) k ε ] , ( 13 ) where ε is drawn from a Rademacher distribution . While this yields an unbiased estimator for L ( θ ) and its gradient as shown in Appendix A.1 if DL is contractive , it can be computationally expensive if N has a large mean or have a high variance if DL has an eigenvalue that is close to 1 or −1 , see [ 44 , 17 ] . Since both the first order Gaussian approximation as well as the Russian Roulette estimator hinges on DL having small absolute eigenvalues , we consider a constrained optimisation approach that penalises such large eigenvalues . For the random variableN that determines the truncation level in the Taylor series , we compute bN = ( DL ) Nε/ ∥∥ ( DL ) Nε∥∥2 and µN = b > NDLbN . Note that this corresponds to applying N times the power iteration algorithm and with |λ1| > |λ2| > . . . > |λd| denoting the eigenvalues of the symmetric matrix DL , almost surely µn → λ1 for n → ∞ , see [ 26 ] . For some δ ∈ ( 0 , 1 ) , we choose some differentiable monotone increasing penalty function h : R→ R such that h ( x ) > 0 for x > δ and h ( x ) = 0 for x 6 δ and we add the term γh ( |µN | ) for γ > 0 to the loss function that we introduce below , see Appendix A.2 for an example of h .

Summary: In this paper, the authors propose an adaptive Hamiltonian Monte Carlo method. In their method, they adapt the hyperparameters of HMC by maximizing an approximation of the proposal entropy. And they show the efficiency of their method by experiments.