MOMENTUM CONSERVING LAGRANGIAN NEURAL NETWORKS
Year: 2021

Abstract:
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.

text: 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.

heading: 1 INTRODUCTION AND RELATED WORK
text: 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 cannot 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.

heading: 2 PRELIMINARIES AND PROBLEM FORMULATION
text: 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.

heading: 3 THEORY
text: The flowchart of MCLNN is shown in Fig. 1. In this approach, the positions and velocities at any time t is used to predict the trajectory of the system, while respecting the conservation laws of energy and momenta. We analyze the performance of MCLNN on three systems, namely, linear spring, non-linear spring, and gravitational system to predict the accurate dynamics of these multiparticle interacting systems.

heading: 3.1 EULER-LAGRANGE EQUATION
text: Consider a system of N -particles interacting with each other under a potential. Let us define functional, namely, “action”, S as:
S = ∫ t1 t0 L dt (1)
where L is the Lagrangian of the system. Then trajectory taken by this system to move from the position q0 to q1 in time t0 to t1 is the one that makes the “action” S stationary. This leads to the
equation of motion governing the dynamics of system, namely, the EL equation as: d
dt
∂L ∂q̇ = ∂L ∂q (2)
The acceleration of particles in the system, q̈ = {~̈qi | ni ∈ P} can be computed directly from the EL equation (see Appendix A.3) as:
q̈ = (∇q̇q̇L)−1 [∇qL− (∇q̇qL) q̇] (3)
This acceleration can then be used to compute the updated positions and velocities of particles. Note that EL equation leads to the conservation of Hamiltonian, defined as: H(q, q̇) = T (q̇) + V (q), which also represents the total energy of the system.

heading: 3.2 MOMENTUM CONSERVING LAGRANGIAN NEURAL NETWORK
text: As mentioned earlier, Lagrangians exhibit translational and rotational symmetry. Learning this directly from the positions of particles is extremely challenging. To address this challenge, we reformulate the LNN. Consider five balls connected by linear springs. We will use this running example to explain MCLNN. The input for the MCLNN is the position and velocities at any time t. Further, we identify the set of all the pairs of particles in the system as (i, j). For a given set, we compute the difference between the positions followed by computing the l2 norm to obtain the pairwise distances qij , which corresponds to the length of the spring. Thus,
qij = √ (~qi − ~qj) · (~qi − ~qj) (4)
The qij is given as the input to a neural network which outputs a scalar Vij . Note that the Vij may be considered to be the pair-wise potential energy of a spring connecting two particles i and j. Vij summed over all the pairs of particles gives the V (q), the total potential energy of the system. For multi-particle systems, the kinetic energy of a particle ni is the function of its velocity ~̇qi, can defined as (Goldstein, 2011; Zhong et al., 2021):
Ti(~̇qi) = 1
2 mi~̇q
2 i . (5)
Here, we invoke this expression of kinetic energy, which when summed over all the particles provide the total kinetic energy, T . Then, T and V are used to obtain the L as in Eq. 7. The L when substituted into the EL equation (Eq. 2), provides the acceleration and the updated position and velocities of the system using the velocity-Verlet integration (Rapaport, 2004). The loss function of MCLNN is on the predicted and actual positions at time t + 1, t + 2, . . . , t + n, which is then back-propagated to train the neural network. Specifically, the loss function is as follows.
L = 1 |T| ∑ ∀T ∈T  1 |T | |T |∑ t=2 1 N ( N∑ i=1 ( T .q ti − p ( T .q ti ))2) (6) Here, p(T .q ti ) is the predicted position for the ith particle in T at time t and T .q ti is the true position. |T | denotes the last time step in trajectory T . The updated expression for Lagrangian for an N -particle system can be written as:
L(q, q̇) = n∑ i=1 1 2 mi~̇q 2 i − n−1∑ i=1 n∑ j=i+1 Vij(qij) (7)
It should be noted that the specific form of kinetic energy T does not limit the applicability of the present approach (Zhong et al., 2021). Since, the focus of the work is to model multi-particle systems and it is well-known that the kinetic energy of the system generally follows Eq. 5, we choose this expression to simplify the problem. In cases where the expression for kinetic energy is not known or is different, the function can be replaced by an additional neural network with an approach similar to that for V .

heading: 3.3 TRANSLATIONAL AND ROTATIONAL SYMMETRY OF MCLNN
text: Theorem 1 MCLNN exhibits translational and rotational symmetry. PROOF. Consider, the Lagrangian of the system of N -particles as L(q, q̇) = T (q̇) − V (q). First, we focus on the translational symmetry and next on rotational symmetry.
Lemma 1 MCLNN exhibits translational symmetry. PROOF. When the particles are subjected to a translation ~ , the updated positions are given by q′ == {~qi + ~ | ni ∈ P}. The updated Lagrangian of the system is given by L(q′, q̇). To prove L(q′, q̇) = L(q, q̇), we need to prove V (q′) = V (q) as the other term in the Lagrangian remains unaffected. Applying the transformation on the positions and computing the l2 norm,
q′ij = √ [(~qi + ~ )− (~qj + ~ )] · [(~qi + ~ )− (~qj + ~ )])
= √ (~qi − ~qj) · (~qi − ~qj)
=qij (8)
which implies V (q′) = V (q). Therefore, L(q′, q̇) = L(q, q̇).
Lemma 2 MCLNN exhibits rotational symmetry. PROOF. When the system is subjected to pure rotation, the positions are transformed as Qq, where Q is an orthogonal tensor representing a rotation. Correspondingly, the Lagrangian is modified as L(Qq, q̇). As in the case of translation, we need only to prove that V (Qq) = V (q). To prove this, it is worth recalling that, for an orthogonal tensor Q, QQ t = Q tQ = I . Now, the proof can be obtained by computing q′ij as:
q′ij = √ (Q~qi −Q~qj) · (Q~qi −Q~qj)
= √ Q(~qi − ~qj) ·Q(~qi − ~qj)
= √ QQT (~qi − ~qj) · (~qi − ~qj)
= √ (~qi − ~qj) · (~qi − ~qj)
= qij (9)
This implies that V (Qq) = V (q) and hence, L(q, q̇) = L(Qq, q̇). Thus, we demonstrate that the Lagrangian in MCLNN exhibits both translational and rotational symmetry. Invoking Noether’s theorem, it can be proven that the Lagrangian preserving these symmetries will consequently conserve of linear and angular momenta, respectively. (For proof, see App. A.4 and A.5).
Combining Lemma 1 and Lemma 2, we get that the MCLNN will respect the laws of conservation energy, linear momentum, and angular momentum.

heading: 3.4 GENERALIZABILITY AND INTERPRETABILITY
text: The MCLNN exhibits a granular structure, which applies the computations on individual or pairwise entities that are aggregated to obtain the total Lagrangian of the system. This structure, by design, allows model generalizability to unseen system sizes. Once trained on an n-body system, the MCLNN can be applied on an m-body system (m 6= n) directly by replacing n with m in Eq. 7. Finally, we discuss interpretability of MCLNN. The main function learned by MCLNN is the Vij , which depends on qij . Since, Vij adds to form the total potential energy V (q) of the system, it is reasonable to assume that Vij represents the pair-wise potential energy. As we demonstrate later in Sec. 4.4, we show that Vij indeed corresponds to the pair-wise potential energy of the ground truth.

heading: 4 EXPERIMENTS
text: In this section, we benchmark MCLNN and establish: • Accuracy: MCLNN is more accurate in modeling trajectories of N -body systems than
LNN (Cranmer et al., 2020a). • Conservation of Physics: Consistent with the theoretical analysis in Sec. 3.3, MCLNN preserves
the laws of physics more comprehensively than LNN • Generalizability and Interpretability: MCLNN generalizes to unseen data better than LNN and
enables higher interpretability due to predicting the potential energy instead of the Lagrangian directly.
The codebase along with all baselines used and the datasets are available at the anonymous repository https://anonymous.4open.science/r/nbodyMCLNN-2618/.

heading: 4.1 EXPERIMENTAL SETUP
text: We use JAX (Bradbury et al., 2020), a high-performance numerical computing python package and JAX MD (Schoenholz & Cubuk, 2020), a JAX based molecular dynamics package for the n-body simulations.

heading: 4.1.1 TASKS
text: To benchmark the performance of MCLNN, we conduct experiments on three multi-body systems, namely, linear and non-linear springs, and gravitational systems.
• Linear spring. We simulate the dynamics of three particles (e.g.: balls with volume tending to zero) connected with each other by linear springs. The initial conditions of the simulations include the positions and velocities of the three particles and the stiffness of the spring. The potential energy of two balls connected by a spring is given by Vij = 12k(qij − q0)
2, where qij is the instantaneous distance between the springs and q0 is the equilibrium distance, and k is the stiffness of the spring. The force between connected pair of particles is computed as Fij = −dVijdqij = −k(qij − q0). Note that the value of k is kept constant for all the springs.
• Non-linear spring. We simulate the dynamics of three particles connected by non-linear springs. The initial conditions of the simulations include the positions and velocities of the three particles. The potential energy of two particles connected by a non-linear spring is given by Vij = 12k(qij − q0)
4, where qij is the instantaneous distance between the springs and q0 is the equilibrium distance, and k is the stiffness of the spring. The force between connected pair of particles is computed as Fij = −2k(qij − q0)3. The value of k is kept constant for all the springs. • Gravitational system. We simulate the dynamics of four particles interacting with each other under the gravitational force. The initial conditions of the simulations include the positions and velocities of the four masses. The potential energy of a pair of particles with mass m1 and m2 if given by Vij = −Gm1m2qij , where G is universal gravitational constant and qij is the instantaneous distance between the pair of particles. The force between the pair of particles is computed as Fij = Gm1m2
q2ij . Note that the mass m of all the particles are kept constant.
Data Generation: We use forward simulation ofN particles to generate training data corresponding to each task. First, an N -body configuration with given initial positions and velocities are considered. The potential energy of the system is computed from the positions using analytical expressions associated with each task as described above. The force on each particle is then calculated using as the gradient of potential energy, which is used to compute the acceleration. Finally, the velocityVerlet algorithm is used for time integration to obtain the updated positions of the particles. Please see the appendix (A.6) for configurations (including the number of particles, initial configuration, and time step) of each of the tasks.

heading: 4.1.2 BASELINE
text: The original LNN takes both the positions and velocities of all the particles as input and gives the total Lagrangian as the output. In the present case, we make the assumption on the functional form of kinetic energy (see Eq.5). As such, to ensure a fair comparison, a slightly modified version of the original LNN is considered here as the baseline. For this baseline, the positions are given as input to the neural network which gives the potential energy V (~q) as output. The kinetic energy is computed using the functional form, from which the Lagrangian is computed. All the remaining training procedures are maintained same as in the original LNN. It should be noted that decoupling the T and V terms and providing the expression for T should make the learning easier for the LNN and hence should give improved performance in comparison to the original LNN. This model will be referred to as baseline or baseline LNN, henceforth.

heading: 4.1.3 METRICS
text: To analyze the performance of the MCLNN in comparison to the baseline, the choice of the appropriate metrics is crucial. Since the systems considered here are chaotic, slight differences in the initial conditions or predictions will lead to differences in the trajectories. More importantly, given the current state of a system, multiple “correct” trajectories may exist as long as they represent the degenerate or equivalent configurational states of the system. As such, purely computing the error in trajectory with respect to the ground truth may not be representative of the performance of the model. Learning and predicting the Lagrangian ensures that the state of the system is represented ac-
curately, resulting in the prediction of the same or equivalent trajectories of the system. Further, the trajectories should respect the spatial and temporal symmetries resulting in conservation of Hamiltonian (total energy), linear, and angular momenta. Thus, we focus on the Mean Absolute Error (MAE) between predicted (1) Lagrangian, (2) Hamiltonian (total energy), (3) linear and (4) angular momenta, with respect to the ground truth.

heading: 4.1.4 TRAINING
text: For training MCLNN, 100 different trajectories, each having 20 points with fixed time intervals, are used for all three tasks. Each trajectory has different initial conditions and hence represent different configurations that are accessible to the system. The model performance is evaluated by comparing the predicted trajectory with respect to the ground truth for the same initial state.
For training LNN, we sample 10,000 data points (set of (q, q̇) as input) and q̈ as output) from forward simulations. The training is performed using the mean squared error between the acceleration predicted by the EL equation and the ground truth. The detailed configurations and parameters associated with simulation and training of each task is provided in App. A.6.

heading: 4.1.5 INFERENCE
text: To asses the long term stability of models, we perform forward simulations using trained LNN and MCLNN models for a prolonged time interval for each task. The initial condition and other parameters are kept same for both LNN and MCLNN. Furthermore, the initial conditions of the trajectories during inference are different from those encountered during the training. Consequently, all trajectories during evaluation are unseen. We record Lagrangian, Hamiltonian, linear momentum and angular momentum for the whole trajectory. The predicted forward simulation is compared with the ground truth data as generated in Sec. 4.1.1.

heading: 4.2 ACCURACY
text: First, we focus on the linear spring system with three particles. Figure 2(a.1)-(a.4) shows the Lagrangian, Hamiltonian (total energy), linear, and angular momenta of the system. The error in these quantities predicted by the baseline LNN and MCLNN are shown in Figure 2(b.1)-(b.4). We observe that the Lagrangian predicted by the baseline LNN starts diverging after 30 time steps of the forward simulation. We notice that the divergence in the Lagrangian is accompanied by a divergence in the linear momentum as well. Up on further simulation, the angular momentum also starts diverging. In contrast, the Lagrangian predicted by the MCLNN follows the ground truth with very little error. Further, the Hamiltonian, and linear and angular momenta remain conserved in the MCLNN. This suggests that the trajectory predicted by the MCLNN is accurate and stable with no long term drift or divergence. It is worth noting that, for the baseline LNN, despite the poor predictions of the Lagrangian at higher values of time steps, the total energy remain conserved. This points to the fact that the energy conservation is a constraint enforced by the EL equation and is not a reflection on the quality Lagrangian. Any value of Lagrangian can still yield a constant Hamiltonian (energy), provided the constraints as imposed by the EL equations on the Lagrangian are satisfied. As such, LNNs should be evaluated considering multiple metrics as demonstrated here.
Now, we focus on the non-linear spring and the gravitational system. Figure 3(a) and (c) shows the Lagrangian, Hamiltonian (total energy), linear, and angular momenta of the non-linear spring and gravitational systems. The error in these quantities predicted by the baseline LNN and MCLNN are shown in Figure 2(b) and (d). As in the case of the linear spring system, we observe that there is a long term drift in the Lagrangian predicted by the baseline LNN. Similarly, the linear and angular momenta of the baseline also diverge in a few time steps. In contrast, the Lagrangian predicted by the MCLNN exhibits a good match with the ground truth. It should be noted that the multiparticle interacting systems are chaotic in nature and hence the exact trajectory of the ground truth and simulated system may vary. However, the low values of error in the Lagrangian along with the long-term stability of both non-linear spring and gravitational system suggests the realistic nature of the trajectory predicted by MCLNN. In addition, we observe that the linear and angular momenta of the system remain conserved in the MCLNN. Altogether, the results confirm that the MCLNN can successfully learn the generative model p(T) that can sample the trajectory effectively from the physics-constrained configurational and temporal space.

heading: 4.3 GENERALIZABILITY TO UNSEEN SYSTEM SIZES
text: To demonstrate the generalizability of MCLNN, we use the generative model p(T) for each of the tasks learned during the training. The models are then used to predict trajectories of systems
with different number of particles. Figure 4 shows the trajectories for linear and non-linear springs with six particles (MCLNN trained on three particles), and gravitational system with eight particles (MCLNN trained on four particles), predicted by the MCLNN in comparison to the ground truth. We observe that the Lagrangian predicted by the MCLNN exhibits excellent match with the ground truth. In addition, all other quantities such as energy and momenta remain conserved (see Appendix Fig. 7). This suggests that once the generative model p(T) has been learned, MCLNN can generalize it to predict the trajectory of systems with any number of particles. Note that the baseline LNN cannot simulate systems of different sizes as the positions of all the particles are given simultaneously as the input to the neural network.

heading: 4.4 INTERPRETABILITY
text: Figure 5 shows the Vij predicted by the MCLNN in comparison to the pair-wise potential energy obtained analytically from ground truth. The shaded region represents the qij values that were present in the training set. We observe that the MCLNN is able to learn the pairwise potential energy function accurately for the qij values in the training set. Further, the MCLNN is able to interpolate values excellently and extrapolate reasonably. It is worth noting that the MCLNN is trained only on the trajectories of the system. Hence, no information regarding the potential energy, force, or even acceleration is given to the neural network during training. This suggests that MCLNN is indeed able to learn the physics of the problem by training purely on the trajectory.

heading: 5 CONCLUSION
text: We introduced a new framework, namely, momentum conserving Lagrangian neural networks (MCLNN), for incorporating physics-based priors in neural networks for accurate simulations of multi-particle systems. We demonstrated that the MCLNN respects the symmetries in space and time, leading to conservation of energy, and linear and angular momentum. We showed that the incorporation of these additional conservation laws of momenta makes the Lagrangian of the system stable by avoiding any long-term drift in it. This, in turn, results in a realistic simulation of multibody systems. Further, we showed that the MCLNN once trained, can generalize to systems of any size. Finally, we demonstrated that the MCLNN is highly interpretable and provides direct insights into the interaction laws governing the dynamics of multi-particle systems. This, in turn, allows one to verify the realistic nature of the function learned by the MCLNN.
At this juncture, it is worth discussing some of the open questions and shortcomings of MCLNN that can be addressed as part of future works. (i) MCLNN assumes interaction between all the particles in the system. A graph-based MCLNN with contrastive loss can potentially address this challenge by incorporating the topology of the particle system along with differential importance for nearby particles. (ii) MCLNN considers only pair-wise interactions for computing the Vij . However, there could be additional interactions involving three (Vijk), four (Vijkl) or even higher particles. This could be addressed by feature engineering or graph-based MCLNN. (iii) Similarly, simulations involving different types of interacting particles is challenging in MCLNN as it may lead to non-unique solutions. (iv) In addition, present work can be extended to address more challenging problems, for instance, to learn the Lagrangian of system with non-conservative forces, and to learn generalized kinetic energy functions, while maintaining the granularity and generalizability.

heading: 6 REPRODUCIBILITY STATEMENT
text: Please find all codes and experiments at the following anonymous link: https: //anonymous.4open.science/r/nbodyMCLNN-2618/. In Appendix A.6, we have given additional details to make the work completely reproducible. In particular, we have: (1) additional derivations and proofs in App. A.3,A.4,A.5, (2) initial conditions, task configurations, and parameters for the training of neural networks used for each task in App. A.6, (3) detailed notation in App. A.7. The dataset generation (using simulations) step is added to the code before training step and hence there are no additional dataset files are provided.

heading: A APPENDIX

heading: A.1 BASELINE LNN
text: Figure 6 shows the predicted force with respect to the actual force for the training data of baseline for the three tasks.

heading: A.2 GENERALIZABILITY TO UNSEEN SYSTEM SIZES
text: Figure 7 shows the generalizability of the MCLNN to unseen system sizes for the three tasks considered. For all the three tasks, the number of particles considered are twice that of the training system size. We observe that the Lagrangian for the unseen system is predicted in excellent agreement to the ground truth for all the three tasks. In addition, the Hamiltonian (total energy), and momenta are conserved in all the three tasks, confirming the realistic nature of the simulations by MCLNN.

heading: A.3 COMPUTING ACCELERATION FROM EULER-LAGRANGE EQUATIONS
text: The acceleration of particles can be obtained from the Euler-Lagrange equations as follows.
∂
∂q̇
dL dt = ∂L ∂q
∂
∂q̇
( ∂L
∂q
dq dt + ∂L ∂q̇ dq̇ dt
) = ∂L
∂q
∂
∂q̇
( ∂L
∂q q̇+
∂L ∂q̇ q̈
) = ∂L
∂q
∂
∂q̇
∂L ∂q q̇+ ∂ ∂q̇ ∂L ∂q̇ q̈ = ∂L ∂q
∂
∂q̇
∂L ∂q̇ q̈ = ∂L ∂q − ∂ ∂q̇ ∂L ∂q q̇
q̈ =
( ∂
∂q̇
∂L
∂q̇
)−1 [ ∂L
∂q − ∂ ∂q̇ ∂L ∂q q̇ ] q̈ = (∇q̇q̇L)−1 [∇qL− (∇q̇qL) q̇]

heading: A.4 TRANSLATIONAL SYMMETRY
text: The conservation of linear momentum can be derived from translational symmetry of Lagrangian as follows.
L(q, q̇) = L(q+ , q̇)
q ′ = q+
=⇒ ∂q ′
∂q = 1∑ ∂L ∂q = ∑ ∂L ∂q ′ = ∑ ∂L ∂q ∂q ∂q ′ = ∑ ∂L ∂q .1
δL = ∑ δq ∂L
∂q δL = ∑ ∂L
∂q δL = ∑ ∂L
∂q
since,
δL = 0 =⇒ ∑ ∂L
∂q = 0∑ ∂L
∂q = 0∑ d
dt
∂L ∂q̇ = 0
d
dt ∑ ∂L ∂q̇
= 0∑ ∂L ∂q̇
= constant∑ pi = constant (linear momentum)
where pi = ∂L∂q̇i

heading: A.5 ROTATIONAL SYMMETRY
text: The conservation of angular momentum can be derived from rotational symmetry of Lagrangian as follows.
Under an infinitesimal rotation δθ
δ~qi = δθ × ~qi
and
δ~̇qi = δθ × ~̇qi
since, Lagrangian does not change with infinitesimal rotation δθ
δL = ∑ i ∂L ∂~qi δ~qi + ∑ i ∂L ∂~̇qi δ~̇qi = 0
using generalized momentum
∂L ∂~̇qi = ~pi
then EL equation gives
d dt ~pi − ∂L ∂~qi = 0
~̇pi = ∂L
∂~qi =⇒ δL = ∑ i ~̇pi · δ~qi + ∑ i
~pi · δ~̇qi = 0∑ i
~̇pi· (δθ × ~qi) + ~pi · (δθ × ~̇qi) = 0∑ i
δθ· (~qi × ~̇pi) + δθ· (~̇qi × ~pi) = 0∑ i δθ· [(~qi × ~̇pi) + (~̇qi × ~pi)] = 0
∑ i δθ· d dt (~qi × ~pi) = 0
because δθ is arbitrary, then ∑ i d dt (~ri × pi) = 0
d dt ( ∑ i ~ri × pi) = 0
angular momentum = ∑ i (~ri × ~pi) = constant
A.6 INITIAL CONDITIONS AND TASK CONFIGURATION
The number of hidden units are chosen from hyperparamter search given in Table 1. The activation function for all hidden units are square plus f(x) = (x+ √ x2+4) 2 which is similar to soft plus (Cranmer et al., 2020a). We use ADAM (Kingma & Ba, 2015) optimiser for model training with learning rate given in respective configuration tables. We use velocity-Verlet for time integration during trajectory evolution. Mass of all the particles are maintained as 1.0 units.

heading: A.6.1 LINEAR SPRING SYSTEM
text: Initial position = [ 0.486657678894505 0.755041888583519 0.0 −0.681737994414464 0.293660233197210 0.0 −0.022596327468640 −0.612645601255358 0.0 ]
Initial velocity = [−0.182709864466916 0.363013287999004 0.0 −0.579074922540872 −0.748157481446087 0.0 0.761784787007641 0.385144193447218 0.0 ]
dt = 0.01 mass = 1.0 stride = 10 runs = 20 = points per trajectory lr = 1.0e− 3 layers = [10, 10] epochs = 100000 samples = 100 = number of trajectories seed = 100 time step = dt× stride = 0.1

heading: A.6.2 NON-LINEAR SPRING SYSTEM
text: Initial position = [ 0.486657678894505 0.755041888583519 0.0 −0.681737994414464 0.293660233197210 0.0 −0.022596327468640 −0.612645601255358 0.0 ]
Initial velocity = [−0.182709864466916 0.363013287999004 0.0 −0.579074922540872 −0.748157481446087 0.0 0.761784787007641 0.385144193447218 0.0 ] dt = 0.01 mass = 1.0 stride = 10 runs = 20 = points per trajectory lr = 1.0e− 3 layers = [10, 10] epochs = 100000 samples = 100 = number of trajectories seed = 100 time step = dt× stride = 0.1

heading: A.6.3 GRAVITATIONAL SYSTEM
text: Initial position =  1.0 0.0 0.09.0 0.0 0.011.0 0.0 0.0 −1.0 0.0 0.0  Initial velocity =
0.0 0.05 0.00.0 −0.05 0.00.0 0.65 0.0 0.0 −0.65 0.0  dt = 0.01 mass = 1.0 stride = 10 runs = 20 = points per trajectory lr = 1.0e− 3 layers = [10, 10] epochs = 100000 samples = 100 = number of trajectories seed = 100 time step = dt× stride = 0.1

heading: A.7 DEFAULT NOTATION
text: a A scalar (integer or real)
~a A vector
A A matrix
In Identity matrix with n rows and n columns
I Identity matrix with dimensionality implied by context
e(i) Standard basis vector [0, . . . , 0, 1, 0, . . . , 0] with a 1 at position i
~ai Vector corresponding to particle i
Ai,j Element i, j of matrix A
Ai,: Row i of matrix A
A:,i Column i of matrix A
Ai,j,k Element (i, j, k) of a 3-D tensor A
A:,:,i 2-D slice of a 3-D tensor dy dx Derivative of y with respect to x
∂y ∂x Partial derivative of y with respect to x ∇~xL Gradient of L with respect to ~x ∇Xy Matrix derivatives of y with respect to X ∂f ∂x Jacobian matrix J ∈ Rm×n of f : Rn → Rm∫
f(x)dx Definite integral over the entire domain of x∫ S f(x)dx Definite integral with respect to x over the set S
P (a) A probability distribution over a discrete variable
p(a) A probability distribution over a continuous variable, or over a variable whose type has not been specified

Review summaries:
This paper enhances Lagrangian neural networks by adding conservation of the angular and linear momenta. It does so by enforcing symmetry with respect to translation and rotation of the system in the Lagrangian. Experiments performed on a set of physical systems shows angular and linear momenta conservation reduce the divergence between the (implicitly) learned and exact dynamic of those systems. 

This paper is motivated by LNN and aim to enforce momentum conservation in learning based on the Lagrangian framework. The goal is straightforward and simple, but I have not seen similar work in the literature. The kinetic energy is assumed to be known and pairwise potential energy is learned through a neural network. Since the distance between each pair of particles is constructed as the input to the potential energy network, the learned model enforces energy as well as momentum conservation. The proposed model is tested on three Interacting particle systems with homogeneous pairwise interaction. 

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.