Hamiltonian Dynamics with Non-Newtonian Momentum for Rapid Sampling

1 Introduction Langevin Dynamics : Random walk mixes slowly between modes Sampling with Deterministic , Non-Newtonian Momentum Nosé-Hoover Thermostat Deterministic , but slow thermalization HMC + NUTS Adapts to traverse long distance , but high overhead in gradient computation Figure 1 : Sampling with 50 gradient evaluations per method . A special property of our proposed deterministic dynamics using a non-Newtonian momentum , proportional to line thickness , is that it ergodically samples from the target distribution . While probabilistic reasoning is crucial for science and cognition [ 1 ] , distributions that can be directly sampled are rare . Without sampling it is difficult to measure likely outcomes , especially in high dimensional spaces . A general purpose method to sample any target distribution , the Metropolis algorithm for Markov Chain Monte Carlo ( MCMC ) , is considered one of the top algorithms of the 20th century [ 2 ] . The “ Monte Carlo ” in MCMC refers to the essential role of stochasticity , or chance , in the approach : start with any state , propose a random update , and accept it according to a weighted coin flip . The most widely used methods for sampling mainly differ in their choice of proposals for the next state of the Markov chain . The most prominent example , Hamiltonian Monte Carlo ( HMC ) , uses Hamiltonian dynamics to suggest proposals that quickly move over large distances in the target space [ 3 ] . HMC using the No U- 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Turn Sampling ( NUTS ) heuristic for choosing hyper-parameters [ 4 ] forms the backbone of modern probabilistic programming languages like Stan [ 5 ] . In the limit of one discrete HMC step per proposal , we recover the widely used discrete Langevin dynamics [ 3 ; 6 ] . Many approaches train neural networks to improve sampling either through better proposal distributions or flexible variational distributions [ 7 ; 8 ; 9 ; 10 ; 11 ; 12 ; 13 ; 14 ] , but these can be expensive if , for instance , the target distribution to be sampled is itself being optimized . We propose a new Hamiltonian dynamics , based on the introduction of a non-Newtonian momentum , which leads to a deterministic dynamics that directly samples from a target energy function . The idea that deterministic dynamics in an extended space can directly sample from a Gibbs-Boltzmann distribution goes back to Nosé [ 15 ] . This principle is widely used in molecular dynamic ( MD ) simulations [ 16 ] but has had limited impact in machine learning . One reason is that direct application of techniques developed for MD often exhibit slow mixing in machine learning applications . Our contribution stems from the recognition that many of the design considerations for molecular dynamics samplers are to enforce physical constraints that are irrelevant for machine learning problems . By returning to first principles and discarding those constraints , we discover a Hamiltonian whose dynamics rapidly sample from a target energy function by using a non-physical form for the momentum . This is in contrast to the more physical , but slower , alternative of introducing auxiliary “ thermostat ” variables that act as a thermal bath [ 17 ; 18 ; 19 ; 20 ] . The result is a fast and deterministic drop-in replacement for the MCMC sampling methods used in a wide range of machine learning applications . Alternatively , our method can be viewed as a normalizing flow in an extended state space which directly samples from an unnormalized target density without additional training . The most significant impact of our approach is for applications where minimizing memory and computation is the priority . For training energy-based models , for example , sampling appears in the inner training loop . Re-training a neural sampler after each model update is costly and entails complex design and hyper-parameter choices . While MCMC sidesteps the training of an extra model , MCMC converges slowly and has led to widespread use of a number of dubious heuristics to reduce computation [ 21 ; 22 ] . Because our energy sampling dynamics does not have the stochastic , random walk component of MCMC , it converges towards low energy states much faster , especially in the transient regime , and explores modes faster as shown in Fig . 1 and demonstrated in experiments . 2 Energy Sampling Hamiltonian ( ESH ) Dynamics . We define a separable Hamiltonian , H ( x , v ) , over position x 2 Rd and velocity v 2 Rd ( we use momentum and velocity interchangeably since we set mass to 1 ) . H ( x , v ) = E ( x ) +K ( v ) The potential energy function , E ( x ) , is the energy for the target distribution with unknown normalization , Z , that we would like to sample , p ( x ) = e E ( x ) /Z , and is defined by our problem . We consider an unusual form for the kinetic energy , K ( v ) , where v2 ⌘ v · v. KESH ( v ) = d/2 log ( v 2/d ) ( 1 ) For contrast , HMC uses Newtonian kinetic energy , K ( v ) = 1/2 v2 . Hamiltonian dynamics are defined using dot notation for time derivatives and using g ( x ) ⌘ @ xE ( x ) . We also suppress the time dependence , x ( t ) , v ( t ) , unless needed . General Hamiltonian ESH Dynamics HMC/Newtonian Dynamics ẋ = @ vH ( x , v ) ẋ = v/ ( v 2/d ) ẋ = v ( 2 ) v̇ = @ xH ( x , v ) v̇ = g ( x ) v̇ = g ( x ) Hamiltonian dynamics have a number of useful properties . Most importantly , the Hamiltonian is conserved under the dynamics . Secondly , via Liouville ’ s theorem we see that dynamics are volume preserving ( or more generally , symplectic ) . Finally , the dynamics are reversible . While the difference in form between ESH dynamics and Newtonian dynamics appears slight , the effect is significant as we show with an example . Imagine riding a roller-coaster running infinitely around a circular track shown in Fig . 2 . E ( x ) is the potential energy from gravity at a point in the loop specified by a scalar angle x 2 [ 0 , 2⇡ ] . At the bottom , potential energy is small while kinetic energy is large . Under Newtonian physics the speed is slowest at the top so most of a rider ’ s time is spent in dreadful anticipation of the fall , before quickly whipping through the bottom . In both cases , the “ momentum ” is largest at the bottom , but for the ESH dynamics we see a different outcome . Changes in the roller-coaster position , ẋ , happen in slow motion at the bottom when momentum is large and then go into fast-forward at the top of the track where the momentum is small . A spectator will see the rider spending most of their time at the bottom , with low potential energy , and the time spent in each region is exactly proportional to the Boltzmann distribution , as we show mathematically in the next section . An intriguing connection between Newtonian and ESH dynamics is discussed in Appendix A.1 , which mirrors a thermodynamic effect emerging from Tsallis statistics that leads to anomalous diffusion [ 23 ] . Note that v2 = 0 is a singularity in the dynamics in Eq . 2 . For 1-D dynamics , this could be an issue , as the velocity can never change signs . However , in higher dimensions , dynamics always avoid the singularity ( Appendix A.2 ) . 2.1 Ergodic ESH Dynamics Sample the Target Distribution . Assume that we have an unnormalized target ( or Gibbs ) distribution specified by energy E ( x ) as p ( x ) = e E ( x ) /Z with an unknown ( but finite ) normalization constant , Z . The reason for defining the non-standard Hamiltonian dynamics in Eq . 2 is that the uniform distribution over all states with fixed energy , c , gives us samples from the Gibbs distribution after marginalizing out the velocity . p ( x , v ) = E ( x ) +KESH ( v ) c /Z 0 ! p ( x ) = e E ( x ) /Z ( 3 ) The proof is as follows . p ( x ) = Z dv p ( x , v ) = 1/Z 0 Z dv E ( x ) + d/2 log v2/d c = 1/Z 0 Z J ( ) d ⇢d 1 d⇢ ( E ( x ) + d log ⇢ c d/2 log d ) = 1/Z 0 Z J ( ) d ⇢d 1 d⇢ ⇢/d ( ⇢ e E ( x ) /d+c/d+ 1/2 log d ) = e E ( x ) /Z In the second line , we switch to hyper-spherical coordinates , with ⇢ = |v| and angles combined in . In the third line , we use an identity for the Dirac delta that for any smooth function , h ( z ) , with one simple root , z⇤ , we have ( h ( z ) ) = ( z z⇤ ) /|h0 ( z⇤ ) | [ 24 ] . Finally we integrate over ⇢ using the Dirac delta , with the integral contributing constants absorbed into the normalizing constant . This proof resembles results from molecular dynamics [ 15 ] , where an auxiliary scalar “ thermostat ” variable is integrated over to recover a “ canonical ensemble ” in both the coordinates and momentum . In our case , the momentum variables are already auxiliary so we directly use the velocity variables in a role usually filled by the thermostat . The form of the energy is inspired by the log-oscillator [ 25 ] . Although log-oscillator thermostats never gained traction in molecular dynamics because they violate several physical properties [ 26 ] , this is not relevant for machine learning applications . Although Eq . 3 gives us the target distribution as the marginal of a uniform distribution over constant energy states , this is not helpful at first glance because it is not obvious how to sample the uniform distribution either . At this point , we invoke the ergodic hypothesis : the dynamics will equally occupy all states of fixed energy over time . More formally , the ( phase space ) average over accessible states of equal energy is equal to the time average under Hamiltonian dynamics for sufficiently long times . Ergodic hypothesis : 1 Z 0 Z dxdv ( H ( x , v ) c ) h ( x , v ) = lim T ! 1 1 T Z T 0 dt h ( x ( t ) , v ( t ) ) ( 4 ) The importance and frequent appearance of ergodicity in physical and mathematical systems has spurred an entire field of study [ 27 ] with celebrated results by Birkhoff and von Neumann [ 28 ; 29 ; 30 ] explaining why ergodicity typically holds . Ergodicity for a specific system depends on the details of each system through E ( x ) , and the appearance of hidden invariants of the system breaks ergodicity [ 19 ; 31 ; 24 ] . While we do not expect hidden invariants to emerge in highly nonlinear energy functions specified by neural networks , systems can be empirically probed through fixed point analysis [ 20 ] or Lyapunov exponents [ 32 ] . Standard MCMC also assumes ergodicity , but the use of stochastic steps typically suffices to ensure ergodicity [ 3 ] . The key distinction to MCMC is that our scheme is fully deterministic , and therefore we use ergodic dynamics , rather than randomness in update proposals , to ensure that the sampler explores the entire space . As an alternative to relying on ergodicity , we also give a normalizing flow interpretation .

Summary: This paper proposes a deterministic sampling dynamic, rapid energy sampling Hamiltonian dynamics (ESH). Since there is no stochasticity, there is no rejection. The kinetic energy is formulated in a way that the dynamic stays at a state for the time proportional to the Boltzmann distribution of the (potential) energy. This paper also propose a leapfrog integrator specialized for ESH. Compared to MCMC approaches, ESH opens a way of much faster sampling from unnormalized energy based models without any auxiliary generating models.