Momentum Conserving Lagrangian Neural Networks

Realistic models of physical world rely on differentiable symmetries that , in turn , correspond to conservation laws . Recent works on Lagrangian and Hamiltonian neural networks show that the underlying symmetries of a system can be easily learned by a neural network when provided with an appropriate inductive bias . However , these models still suffer from issues such as inability to generalize to arbitrary system sizes , poor interpretability , and most importantly , inability to learn translational and rotational symmetries , which lead to the conservation laws of linear and angular momentum , respectively . Here , we present a momentum conserving Lagrangian neural network ( MCLNN ) that learns the Lagrangian of a system , while also preserving the translational and rotational symmetries . We test our approach on linear and non-linear spring systems , and a gravitational system , demonstrating the energy and momentum conservation . We also show that the model developed can generalize to systems of any arbitrary size . Finally , we discuss the interpretability of the MCLNN , which directly provides physical insights into the interactions of multi-particle systems . 1 INTRODUCTION AND RELATED WORK . Realistic modeling of the time-evolution of multi-particle physical systems lie at the heart of several fields of science and engineering ( Goldstein , 2011 ; Park et al. , 2021 ; Roehrl et al. , 2020 ; Lutter et al. , 2019 ) . Examples include multiple bodies connected by springs , bodies interacting under gravitational forces , or even systems at the atomic and mesoscale such as proteins or colloidal gels , which range several orders ( ≈ 1020 ) of length and timescales from atomic systems to planetary systems ( Park et al. , 2021 ; Cranmer et al. , 2020b ) . Traditionally , the trajectory of these systems are obtained by solving the associated differential equations numerically . These differential equations , in turn , are derived based on the invariant or conserved quantities of a system . Recent studies on Lagrangian ( LNN ) and Hamiltonian neural networks show that one of these invariant quantities , energy , can be learned directly from the data enabling realistic simulation of systems ( Cranmer et al. , 2020a ; Finzi et al. , 2020 ; Lutter et al. , 2019 ; Greydanus et al. , 2019 ; Zhong et al. , 2021 ; 2019 ; Zhong & Leonard , 2020 ) . However , these approaches fail to model multi-particle interactions effectively due to the following limitations . • Conservation of momentum : Invoking Noether ’ s theorem ( Noether , 1971 ) , “ a differentiable symmetry in action results in a conserved quantity ” . Conversely , “ every conserved quantity can be derived from an underlying symmetry in action ” . It can be shown that an interacting multi-particle system , that is closed , respects three conservation laws , namely , energy , linear momentum and angular momentum ( Noether , 1971 ; Goldstein , 2011 ) . While the first is a consequence of symmetry of the system with respect to time , the second and third is a consequence of symmetry with respect to translation and rotation in space , respectively . Thus , the Lagrangian of such a system remains invariant under the translation and rotation operations on the Cartesian coordinates . For example , three balls connected by a spring will have the same Lagrangian ( and hence the interaction forces ) when the system as a whole is subjected to rotation or translation . That is , as long as the relative positions of the balls are preserved , translation or rotation by any magnitude will not have an effect on the system dynamics . While LNNs have been demonstrated to learn the symmetry in time by conserving the energy ( Cranmer et al. , 2020a ) , they have not been designed to respect the other two conservation laws with respect to momentum . Learning these symmetries directly from the data is an extremely hard problem and may require large amounts of data . Furthermore , as we will show in Sec . 4 , even when trained on large volumes of training data , the LNNs fail to generalize well on unseen data . This questions the realistic nature of the dynamics simulated by LNNs . • Generalizability to unseen system sizes : LNNs lack the ability to generalize to system sizes beyond the training set . For instance , an LNN trained on three balls connected by a spring can not model a system of four balls connected by a spring despite having the same conserved quantities and interactions . This limits the broader applicability of LNNs to realistic physical systems . • Interpretability : LNN ( Cranmer et al. , 2020a ) directly learns ( and predicts ) the Lagrangian of a system as a function of the position and velocities of all the particles in the system . Owing to this design , it is not possible to recover the inter-particle dynamics within a multi-particle system in the form of pairwise potential energy or forces . This limits the interpretability of the LNN as the contribution of positions and velocities towards the total Lagrangian of the system is represented using a black-box deep neural network . In this paper , we address the above limitations . Specifically , our contributions are as follows : • Problem Formulation and Architecture Design : We reformulate LNN ( Cranmer et al. , 2020a ) with a relational inductive bias by applying a transformation on the Cartesian coordinates of the system and decoupling the terms of the Lagrangian . Armed with this reformulation , we develop a momentum conserving LNN called MCLNN . MCLNN introduces several key innovations . ( 1 ) First , MCLNN preserves the rotational and translation symmetries . ( 2 ) Second , MCLNN generalizes to unseen system sizes . ( 3 ) Third , MCLNN generates interpretable models with the output characterizing pair-wise interactions of the bodies in the system , and thereby providing physical insights into the system dynamics . • Theoretical Characterization : We rigorously establish that MCLNN conserves energy , linear and angular momentum . Furthermore , in contrast to LNN ( Cranmer et al. , 2020a ) , we show that MCLNN can directly learn trajectories ( positions ) of systems without requiring training on acceleration . This is a desirable property . In many cases including experimental systems , access to acceleration of each entity at every time instant may not be possible . While position of particles is a direct observable , acceleration is a derived quantity . • Empirical Evaluation : We perform in-depth evaluation across three multi-particle systems namely , balls connected by linear and non-linear springs , and gravitational system . Our experiments establish that MCLNN performs significantly better than LNN with long-term stability , generalizes to arbitrary-sized multi-particle systems , and is capable of learning interaction dynamics with 3-4 orders of magnitude lower volume of training data . 2 PRELIMINARIES AND PROBLEM FORMULATION . In this section , we introduce the preliminary concepts central to our work and formulate the problem . As notational convention , vectors are represented with an overhead arrow ( Ex . ~v ) and higher-order tensors , such as matrices , in bold . A summary notations used is provided in the appendix A.7 . Definition 1 ( Multi-particle System ) An N -body multi-particle system contains N particles P = { n1 , · · · , nN } . At any given time t , particle ni is characterized by its position ~q ti and velocity ~̇q ti . The positions of particles in anN -body system are not static . They change due to various interaction forces at play ( Ex : a gravitational system ) . These positional changes are captured in the form of trajectories . Definition 2 ( Multi-particle Trajectory ) The trajectory of an N -body multi-particle system P over a time horizon T = [ ts , te ] is the sequence of position and velocity vectors ( q t , q̇ t ) | t ∈ T } , where q t = { ~q ti | ni ∈ P } and q̇t = { ~̇q ti | ni ∈ P } . In physics , the time-evolution or trajectory of interacting particles is obtained by solving the differential equations of motion . The numerical solution to these equations provide the acceleration of the particle , which can then be used to obtain the updated velocity and positions . Indeed , the acceleration of particles can be directly learned by neural networks as a function of its position and velocities . However , it fails to learn the underlying symmetries , which in turn leads to the laws of conservation of energy and momenta . Recent approaches to predict trajectories by learning the Lagrangian through an LNN ( Cranmer et al. , 2020a ; Finzi et al. , 2020 ; Zhong et al. , 2021 ; Zhong & Leonard , 2020 ) has been shown to be effective in learning the symmetry in time ( Cranmer et al. , 2020a ; Zhong et al. , 2021 ) . Lagrangian , defined as L ( q , q̇ ) = T ( q̇ ) −V ( q ) , is a scalar functional of kinetic energy ( T ( q̇ ) ) and potential energy ( V ( q ) ) , that depends on both the set of positions q , and velocities q̇ of the particles in the system ( Goldstein , 2011 ) . LNN bypasses the necessity to learn accelerations for each particle and directly predicts the scalar Lagrangian for a system . Additionally , the trajectory predicted using Euler-Lagrange ( EL ) equation conserves energy resulting in an overall better prediction of the system dynamics . Theoretically , although learning the Lagrangian of a system is enough to predict its dynamics , the Lagrangian itself exhibits some symmetries that are not imposed by the EL equation . Lagrangian of a system exhibits translational and rotational symmetry . Specifically , the Lagrangian of a system L ( q , q̇ ) and L ( q′ , q̇ ) , where q′ = { ~qi + ~ | ni ∈ P } for any finite value of ~ are equivalent . Similarly , L ( q , q̇ ) and L ( Qq , q̇ ) , where Q is an orthogonal tensor representing a rotation , are equivalent . These symmetries are not provided as an inductive bias in the LNN framework . Further , LNNs are trained by giving the coordinates as input and minimizing the loss between the actual acceleration and acceleration predicted by LNN through EL equations . This requires a priori access to the actual acceleration , which may not be available in many cases , for example , experimental trajectory of colloidal gels visualized by video camera . Here , we aim to address these open challenges by reformulating the LNN with a relational inductive bias , resulting in Momentum Conserving LNN ( MCLNN ) . Specifically , we aim to learn the dynamics of multi-particle interacting systems purely from the trajectory of the constituent particles . Problem 1 ( Momentum Conserving LNN for Trajectory Prediction ) Let T be a set of trajectories ofN -body systems . Furthermore , let there be a hidden joint distribution of physics-constrained configurational and temporal space from which T has been sampled . Our goal is to learn this hidden distribution . Towards that end , we want to learn : i . A generative model p ( T ) that maximizes the likelihood of generating T. ii . The generative model must respect the laws of physics such as conservation of energy and momenta . Once learned , this generative model can be used to predict trajectories of unseen N -body systems .

Summary: Many recent studies have attempted to learn Lagrangian and Hamiltonian systems by using neural networks, but most of them still suffer from instability to learn the translational and rotational symmetries. These symmetries are associated with the conservation laws of linear momentum and angular momentum, respectively. In contrast to Lagrangian neural network, which learns a Lagrangian as a black-box, the proposed method learns the Lagrangian as the kinetic energy predefined by (5) minus the sum of pair-wise potential energy $V_{ij}$ modeled by a neural network. Thanks to this explicit formulation, the proposed method conserves not only the system energy but also linear/angular momentum, as confirmed by the experiments.