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: The paper proposes a novel Hamiltonian to define the sampling dynamics in the state space. The main difference with the classical Newtonian Hamiltonian is a different kinetic energy. The proposed kinetic energy is $K(v) = d/2\log(v^2/d)$, where $v$ is the momentum, and $d$ is the dimensionality of the state space. The peculiar property (which is the central point of study) of this Hamiltonian is that the velocity $\dot{x} = v/(v^2/d)$ decreases with the increase of the kinetic energy. Intuitively, this energy defines very reasonable dynamics: the particle slows down in the high-density regions and accelerates with the decrease of the density. Note that the classical HMC behaves oppositely.  For the proposed dynamics, the authors carefully develop technical details. They hypothesize that, unlike HMC, this dynamics could actually be ergodic in the state space and the resampling of the momentum is superfluous then. They also develop a specific integration scheme (based on Leap-Frog) that facilitates an efficient simulation of the dynamics. Throughout the paper, the authors claim that the accept/reject step is not necessary for the correct sampling if the integration is precise enough.  For the empirical evaluation, the authors consider several synthetic distributions and compare the algorithm against similar "physical" samplers. For the synthetic distributions, the authors report ESS and MMD. There is also a brief comparison against other samplers on the energy-based models, where the authors report the average energy of the samples and provide visual examples.