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: The paper proposes to adapt the mass matrix used in HMC by maximising the speed measure introduced in [1]. Estimating the speed measure is non-trivial and it is achieved by firstly finding a lower bound of the density induced by a HMC kernel and then estimating the bound using Russian-roulette sampling. Experiments performed on multiple targets show that the method is working in practice.  [1] Michalis Titsias and Petros Dellaportas. Gradient-based adaptive Markov chain Monte Carlo. NeurIPS, 2019.