Symplectic Recurrent Neural Networks

1 Introduction . Can machines learn physical laws from data ? A recent paper ( Greydanus et al. , 2019 ) , Hamiltonian Neural Networks ( HNN ) , proposes to do so representing the Hamiltonian function H ( q , p ) as a multilayer neural network . The partial derivatives of this network are then trained to match the time derivatives ṗ and q̇ observed along the trajectories in state space . The ordinary differential equations ( ODEs ) that express Hamiltonian dynamics are famous for both their mathematical elegance and their challenges to numerical integration techniques . Except maybe for the simplest Hamiltonian systems , discretization errors and measurement noise lead to quickly diverging trajectories . In other words , Hamiltonian systems can often be stiff , a concept that usually refers to differential equations where we have to take very small time-steps of integration so that the numerical solution remain stable ( Lambert , 1991 ) . A plethora of numerical integration methods , symplectic integrators , have been developed to respect the conserved quantities in Hamiltonian systems , thereby usually being more stable and structure-preserving than non-symplectic ones ( Hairer et al. , 2002 ) . For example , the simplest symplectic integrator is the well-known leapfrog method , also known as the Stömer-Verlet integrator ( Leimkuhler and Reich , 2005 ) . However , even the best integrators remain severely challenged by phenomena as intuitive as a mechanical rebound or a slingshot effect , which are more severe forms of stiffness . Such numerical issues are almost doomed to conflict with the inherently approximate nature of a learning algorithm . In the first part of this paper , we propose Symplectic Recurrent Neural Networks ( SRNNs ) , where ( i ) the partial derivatives of the neural-network-parametrized Hamiltonian are integrated with the leapfrog integrator and where ( ii ) the loss is back-propagated through the ODE integration over multiple time steps . We find that in the presence of observation noise , SRNN are far more usable than HNNs . Further improvements are achieved by simultaneously optimizing the initial state and the Hamiltionian network , presenting an interesting contrast to previous literature on the hardness of general initial state optimization ( Peifer and Timmer , 2007 ) . The optimization can be motivated from a maximum likelihood estimation perspective , and we provide heuristic arguments for why the initial state optimization is likely convex given the symplecticness of the system . Furthermore , experiments in the three-body problem show that the SRNN-trained Hamiltonian compensates for discretization errors and can even outperform numerically solving the ODE using the true Hamiltonian and the same time-step size . This could be of particular interest to researchers who study the application of machine learning to numerically solving differential equations . The second part of this paper focuses on perfect rebound as an example of the more severe form of stiffness . When a point mass rebounds without loss of energy on a perfectly rigid obstacle , the motion of the point mass is changed in ways that can be interpreted as an infinite force applied during an infinitesimal time . The precise timing of this event affects the trajectory of the point mass in ways that essentially make it impossible to merely simulate the Hamiltonian system on a predefined grid of time points . In order to address such events in learning , we augment the leapfrog integrator used in our SRNN with an additional trainable operator that models the rebound events and relates their occurrence to visual hints . Training such an augmented SRNN on observed trajectories not only learns the point mass dynamics but also learns the visual appearance of the obstacles . 2 Related work . Learning physics with neural networks A popular category of methods attempts to replicate the intuitive ways in which humans perceive simple physical interactions , identifying objects and learning how they relate to each other ( Battaglia et al. , 2016 ; Chang et al. , 2016 ) . Since such methods can not be used for more general physical systems , another category of methods seeks to learn which differential equations govern the evolution of a physical system on the basis of observed trajectories . Brunton et al . ( 2016 ) assemble a small number of predefined primitives in order to find an algebraically simple solution . Lutter et al . ( 2019 ) use a neural network to model the Lagrangian function of a robotic system . Most closely related to ours , Greydanus et al . ( 2019 ) use a neural network to learn the Hamiltonian of the dynamical system in such a way that its partial derivatives match the time derivatives of the position and momentum variables , which are both assumed to be observed . Although the authors show success on a simple pendulum system , this approach does not perform well on a more complex system such as a three-body problem . ODE-based learning and recurrent neural networks ( RNNs ) To learn an ODE that underlies some observed time series data , Chen et al . ( 2018a ) proposes to solve a neuralnetwork-parameterized ODE numerically and minimize the distance between the generated time series with the observed data . To save memory , they propose to use the adjoint ODE instead of back-propagating through the ODE solver . Using stability analysis of ODEs , Chang et al . ( 2019 ) propose the AntisymmetricRNN with better trainability . Niu et al . ( 2019 ) establish a correspondence between RNNs and ODEs , and propose an RNN architecture inspired by a universal quantum computation scheme . Summary of our main contributions : In this paper , we propose SRNN , which • learns Hamiltonian dynamics directly from position and momentum time series • performs well on noisy and complex systems such as a spring-chain system and a three-body system , and is compatible with initial state optimzation • is augmented to handle perfect rebound , an example of very stiff Hamiltonian dy- namics 3 Framework . 3.1 Hamiltonian systems . A Hamiltonian system of dimension d is described by two vectors p , q ∈ Rd . Typically , they correspond to the momentum and position variables , respectively . The evolution of the system is determined by the Hamiltonian function H : ( p , q , t ) ∈ R2d+1 7→ H ( p , q , t ) ∈ R through a system of ordinary differential equations called Hamilton ’ s equations , ṗ = −∂H ∂q , q̇ = + ∂H ∂p , ( 1 ) where we use the dot notation to compactly represent derivatives with respect to the time variable t. We are focusing in this work on Hamiltonians that are conservative,1 that is , they do not depend on the time variable t , and separable,2 that is , they can be written as a 1In our opinion , extending this work to non-conservative systems should not be done by adding a time dependency in the Hamiltonian , but by adding additional dissipation or intervention operators in the numerical integration schema , as illustrated in section 6 . 2Extending this work to non-separable Hamiltonians can be achieved by rewriting the numerical integration schema using an extended phase space ( Tao , 2016 ) . sum H ( p , q ) = K ( p ) + V ( q ) . In this case , ( 1 ) becomes ṗ = −V ′ ( q ) , q̇ = K ′ ( p ) ( 2 ) With a proper choice of the p and q variables , the evolution of essentially all physical systems can be described with the Hamiltonian framework . In other words , Hamilton ’ s equations restrict the vast space of dynamical systems to the considerably smaller space of dynamical systems that are physically plausible . Therefore , instead of modeling the dynamics of a physical system with a neural network fθ ( p , q ) whose outputs are interpreted as estimates of the time derivatives ṗ and q̇ , we can also use a neural network Hθ ( p , q ) = Kθ1 ( p ) + Vθ2 ( q ) with θ = [ θ1 , θ2 ] , whose partial derivatives −V ′θ2 ( q ) and K ′ θ1 ( p ) are interpreted as the time derivatives ṗ and q̇ . We refer to the former as ODE neural networks ( O-NET ) and the latter approach as Hamiltonian neural networks ( H-NET ) . In order to define a complete learning system , we need to explain how to determine the parameter θ of the neural networks on the basis of observed discrete trajectories . For instance , Greydanus et al . ( 2019 ) trains H-NET in a fully supervised manner using the observed tuples ( p , q , ṗ , q̇ ) . 3.2 From ODEs to discrete trajectories . A numerical integrator ( or ODE solver ) approximates the true solution of an ODE of the form ż = f ( z , t ) at discrete time steps t0 , t1. . .tT . For instance , the simplest integrator , Euler ’ s integrator , starts from the initial state z0 at time t0 and estimates the function z ( t ) at uniformly spaced time points tn = t0 + n∆t with the recursive expression zn+1 = zn + ∆t f ( zn , tn ) ( 3 ) In stiff ODE systems , however , using Euler ’ s method could easily lead to unstable solutions unless the time-step is chosen to be very small ( Lambert , 1991 ) . The development of efficient and accurate numerical integrators is the object of considerable research ( Hairer et al. , 2008 ; Hairer and Wanner , 2013 ) . Symplectic integrators3 are particularly attractive for the integration of Hamilton ’ s equations ( Leimkuhler and Reich , 2005 ) . They are able to preserve quadratic invariants , and therefore usually have desired stability properties as well as being structure-preserving ( McLachlan et al. , 2004 ) , even for certain non-Hamiltonian systems ( Chen et al. , 2018b ) . A simple and widely-used symplectic integrator is the leapfrog integrator . When the Hamiltonian is conservative and separable ( 2 ) , it computes successive estimates ( pn , qn ) with pn+1/2 = pn − 12∆t V ′ ( qn ) qn+1 = qn + ∆tK ′ ( pn+1/2 ) pn+1 = pn+1/2 − 12∆tV ′ ( qn+1 ) ( 4 ) Repeatedly executing update equations ( 4 ) is called the leapfrog algorithm , which is as computationally efficient as Euler ’ s method yet considerably more accurate when the ODE belongs to a Hamiltonian system ( Leimkuhler and Reich , 2005 ) . 3.3 Learning ODEs from discrete trajectories . Following Chen et al . ( 2018a ) , let the right hand side of the ODE be a parametric function fθ ( z , t ) and let z0 . . . zT be an observed trajectory measured at uniformly spaced time points t0 . . . tT . We can estimate the parameter θ that best represents the dynamics of the observed trajectory by minimizing the mean squared error ∑T i=1 ‖zi − ẑi ( θ ) ‖2 between the observed trajectory { zi } Ti=0 and the trajectory { ẑi ( θ ) } Ti=0 generated with our integrator of choice , { ẑi ( θ ) } Ti=0 = Integrator ( z0 , fθ , { ti } Ti=0 ) . For instance , this minimization can be achieved using stochastic gradient descent after back-propagating through the steps of our numerical integration algorithm of choice and then through each call to the functions fθ . This can be done when fθ ( z ) is a neural network ( O-NET ) , or is the concatenation [ −V ′θ2 ( q ) , K ′ θ1 ( p ) ] of the partial derivatives of an H-NET 3An integrator is symplectic if applied to Hamiltonian systems , its flow maps are symplectic for short enough time-steps . For details , see Hairer et al . ( 2002 ) and Leimkuhler and Reich ( 2005 ) . Hθ ( p , q ) = Kθ1 ( p ) + Vθ2 ( q ) , where the partial derivatives can be expressed using the same parameters θ as the Hamiltonian Hθ ( p , q ) , for instance using automatic differentiation . We can then predict trajectories at testing time using the trained fθ∗ and initial state ztest0 , { ẑtesti } T test i=0 = Integrator ( z test 0 , fθ∗ , { ti } T test i=0 ) . Note that neither the integrator , nor the number of steps , nor the step size , need to be the same at training and testing .

Summary: This paper proposes to represent a Hamiltonian model of a physical system by a neural network. The parameters are then adjusted, so that the observations are considered maximally likely under a probabilistic model. The novelty is to consider a symplectic Leapfrog integration scheme for the Hamiltonian system, which is known to conserve important quantities such as volume in the state space. The proposed approach is shown to outperform the recent work "Hamiltonian Neural Networks" by a large margin on mass-spring chain dynamics and three body systems. The approach can even handle stiff dynamical systems such as bouncing billiards.