S TIFFNESS - AWARE NEURAL NETWORK FOR LEARNING H AMILTONIAN SYSTEMS
A BSTRACT
We propose stiffness-aware neural network (SANN), a new method for learning
Hamiltonian dynamical systems from data.
SANN identiﬁes and splits the training data into stiff and nonstiff portions based on a stiffness-aware index, a simple, yet effective metric we introduce to quantify the stiffness of the dynamical system.
This classiﬁcation along with a resampling technique allows us to apply different time integration strategies such as step size adaptation to better capture the dynamical characteristics of the Hamiltonian vector ﬁelds.
We evaluate SANN on complex physical systems including a three-body problem and billiard model.
We show that SANN is more stable and can better preserve energy when compared with the state-of-the-art methods, leading to signiﬁcant improvement in accuracy.
I NTRODUCTION
Data-driven modeling of dynamical systems provides a computationally inexpensive approach for exploiting the scientiﬁc law of the physical processes and predicting future phenomena (Giannakis & Majda, 2014; Harlim et al., 2021;
Gu et al., 2021; Mou et al., 2021).
With the superior approximation and generalization capacity (Lu et al., 2020; Shen et al., 2021) of neural networks (NNs), important advances have been made in learning Hamiltonian systems (Greydanus et al., 2019; Finzi et al., 2020;
Tong et al., 2021).
Based on the physical property that the total energy of the system (also called Hamiltonian) must be conserved, many studies model the Hamiltonian system by learning this conserved quantity from data (Greydanus et al., 2019; DiPietro et al., 2020).
Even though some successes have been achieved it remains challenging to learn the estimated system to capture the exact physical law, because of the elusive and chaotic characteristics of the systems (Choudhary et al., 2020; Marchal, 1990), especially for complex systems that are intrinsically stiff.
When learning stiff dynamics, the NN optimization easily leads to an unstable solution or biased estimation due to the lack of constraint on parameters (Kim et al., 2021;
Wang et al., 2020).
To illustrate the difﬁculty in learning stiff Hamiltonian dynamics, let us consider a three-body problem that describes the interactions of three particles under gravitational force.
According to physical law, the repulsive force of two particles increases dramatically when they get close to each other.
Hence, the close three-body interaction intensiﬁes stiffness phenomena (Huang & Leimkuhler, 1997).
Fig.
1 shows the reference orbits of three particles and the orbits learned by the Hamiltonian neural network (HNN) (Greydanus et al., 2019), a famous approach that directly approximates the Hamiltonian with a neural network.
In Fig.
1(a, e), the HNN orbits coincide with the reference when the three particles are far from each other, which is consistent with the results in (Greydanus et al., 2019).
However, the particles deviate from the reference orbits rapidly after close encounter, and the energy becomes nonconserved (Fig.
1(c, f)).
This reveals although the NN is trained with the data containing close encounter, the NN does not fully capture the dynamics and easily diverges when stiffness grows.
Two primary factors may affect the accuracy of the learned Hamiltonian.
Implicit bias.
NN-based optimization has an implicit bias toward ﬁtting a smooth function with the fast decay in the frequency domain (Xu et al., 2020; Cao et al., 2019).
This implicit bias impedes the NN from capturing high-frequency components, such as singularities in an N-body problem.
Imbalanced stiffness proportion.
The stiffness of the dynamical system changes with time and varies across different trajectories (e.g., the same system with different initial conditions).
It is not uncommon that only a small proportion of trajectories corresponds to stiff dynamics.
As a statistical example, in the 1 , 000 independent simulations of three-body trajectories following (Chen et al., 2019), 91.4% of the trajectories contain close encounter, but on average only 4 . 2% of the time intervals within the trajectories contain close encounter.
To mitigate the implicit bias caused by NN-based optimization and the imbalanced stiffness proportion problem, we propose a new method called the stiffness-aware neural network (SANN) based on the stiffness classiﬁcation of the training data.
In our approach, we introduce a stiffness-aware index (SAI) as a simple, yet effective metric to classify the time intervals into stiff and nonstiff portions.
For training efﬁciency, we integrate the Hamiltonian dynamics over different intervals with different step sizes based on their classiﬁcation.
To balance the ratio between the stiff group and the nonstiff group and avoid biased training, we resample the stiff intervals.
Our contributions are as follows.
We identify the importance of the stiffness concept in learning Hamiltonian systems from the time series data.
We show that SAI is easy to calculate and can be used to effectively determine stiffness intervals of the data.
We validate the SANN method with complex Hamiltonian dynamics including a three-body problem and billiard model.
Extensive numerical results show that SANN can accurately predict the the stiff dynamics and signiﬁcantly outperform the existing methods.
P RELIMINARIES 2. H AMILTONIAN SYSTEM
The Hamiltonian system describes the continuous-time evolution of states in the phase space ( p , q ) , where p ∈ R n represents the generalized momentum and q ∈ R n denotes position coordinates.
The Hamiltonian H ( p , q ) : R 2 n → R 1 is denoted as the total energy of the system at ( p , q ) .
The dynamics can be described with H ( p , q ) by
[Equation]
With Eqn.
(1), H ( p , q ) is conservative during the evolution as d H d t = ∂
H ∂ p · d p d t + ∂ H ∂ q · d q d t = 0 .
In classical mechanics, the Hamiltonian is usually expressed as the sum of the kinetic and potential energies.
We are interested in learning the separable Hamiltonian, namely, H ( p , q )
=
T ( p ) +
V ( q ) , where T ( p ) is kinetic energy while V ( q ) is potential energy.
In this paper our goal is to learn the Hamiltonian of a dynamical system.
What is available is the time series of observations { ( p t i , q t i ) } Ni =1 .
The Hamiltonian is parameterized by a neural network as follows,
[Equation]
where φ : R n → R 1 is a fully connected NN, M ∈ R n × n is a trainable matrix, and θ (cid:44) { M , W } is a set of all trainable parameters.
Once we obtain H θ , the trajectories of the learned Hamiltonian system can be simulated by Eqn.
(1) with an ordinary differential equation (ODE) solver.
2. ODE S OLVER FOR THE H AMILTONIAN S YSTEM
We use Forward(( p 0 , q 0 ) , H , ∆ t ) to denote a one-step integration of Eqn.
(1) from ( p 0 , q 0 ) over the step size ∆ t .
The Euler method is a ﬁrst-order Runge–Kutta method adopted in (Greydanus et al., 2019).
When the Euler method is used, the integration becomes
[Equation]
The Euler method is usually not suitable for stiff dynamics because of its poor stability property.
Another popular numerical method (Fehlberg, 1969) for the Hamiltonian system is Runge—Kutta- –Fehlberg (RKF45).
RKF45 allows the step size to be changed adaptively based on an estimate of the local truncation error.
This error-based step size control has great potential for improving computational efﬁciency.
However, its effectiveness may be hampered by the use of minibatches, which are crucial for training.
Minibatching adds an additional dimension to Eqn.
Controlling the step size requires considering a combined ODE system and estimating the error on all batch elements, as noted in Chen et al.
Moreover, the step size is limited by the stiffest element in a batch, making it difﬁcult to use a large step size especially when the batch size is large.
A symplectic integrator, a numerical method that conserves the energy quantity, is widely used in learning Hamiltonian system (Chen et al., 2019; Zhong et al., 2020).
Leapfrog is a second-order symplectic integrator designed for a separable Hamiltonian.
The Leapfrog scheme is in Appendix C.
S TIFFNESS -A WARE N EURAL N ETWORK
In this section we introduce our method, called SANN (stiffness-aware neural network), with its workﬂow shown in Fig.
First, we propose SAI (stiffness-aware index) to classify the time interval as either stiff or nonstiff.
During training, we integrate the NN-parameterized Hamiltonian dynamics over the interval with different step sizes based on their classiﬁcation.
Next, we resample the stiff intervals to balance their ratio and avoid biased training.
3. I DENTIFYING THE STIFF INTERVAL
In this part we ﬁrst discuss the stiffness index (SI) that is used to characterize stiffness for an ODE and propose the stiffness-aware index (SAI) to characterize the stiffness for the time series data.
SI reveals the fastest rate of change of state.
We consider the dynamics of the state u ,
[Equation]
SI at the state u ( t ) is deﬁned by max {| Re ( λ i ) |} , where λ i is the eigenvalue of the Jacobian matrix of Eqn.
(4) (Aiken, 1985).
To illustrate that SI reveals the fastest rate of change of state, we consider that Eqn.
(4) is a linear system with constant coefﬁcients, namely, d u dt
=
Au , and A ∈ R n × n is a diagonalizable matrix with eigenvalues { λ i } ni =1 and corresponding eigenvectors { v i } ni =1 .
Then the solution of the linear system is u ( t ) =
(cid:80) ni
=1 c i v i e λ i t .
Let us suppose that Re ( λ i ) < 0 , i = 1 , ..., n .
We have e λ i t → 0 as t → ∞ .
Hence, max {| Re ( λ i ) |} reveals the fastest speed of decaying to 0 .
If max {| Re ( λ i ) |} is large, the integrator needs a small step size to reduce the local truncation error.
Deﬁnition of SAI.
Given the observations { u t i } Nt =1 from the system (4), we deﬁne SAI at u t i by
[Equation]
Intuitively, SAI characterizes the relative changing speed of the state u from time t i to t i +1 .
The norm of ﬁnite difference (cid:13)(cid:13) u ti +1 − u ti t
i +1 − t i (cid:13)(cid:13) 2 approximates (cid:13)(cid:13)
f ( u t i ) (cid:107) 2 , which is used in time parameterization of the adaptive time step method (Huang & Leimkuhler, 1997).
In our data-driven setting, we cannot compute SI as the analytic expression of f is unavailable.
But SAI can serve as a proxy of SI in time series data, which will be demonstrated in Section 5.
Next, we split the training data into stiff and nonstiff portions by classifying the time interval as either stiff or nonstiff based on SAI.
For a time series observation { u t i } Ni =1 , we ﬁrst compute SAI for each time interval.
Let SAI i = 1 (cid:107) u ti (cid:107) 2 (cid:13)(cid:13)
u ti +
− u ti t
i +1 − t i (cid:13)(cid:13) 2 .
We then rank { SAI i } N − 1 i =1 .
The interval ( t i , t i +1 ) with a larger SAI i is supposed to be a stiffer part.
We use γ ∈ (0 , 1) to denote the stiff ratio, a hyperparameter to determine the ratio of the stiff portion.
The intervals belonging to the top- γ of { SAI i } N − 1 i =1 are identiﬁed as stiff intervals, and the others are classiﬁed as nonstiff.
In other words, the intervals in the (1 − γ ) × 100 -th percentile of SAI are identiﬁed as nonstiff intervals.
3. T RAINING THE ESTIMATED H AMILTONIAN
Loss function.
Let { ( p t i , q t i ) }
Ni =1 be the given trajectory from Hamiltonian system.
Given a pair of consecutive states, ( p t i , q t i ) and ( p t i +1 , q t i +1 ) , we integrate Eqn.
(1) from t i to obtain the estimated solution ( ˆp t i +1 , ˆq t i +1 ) at t i +1 .
To achieve a more accurate solution, we conduct the integration over [ t i , t i +1 ] with a small time step t i +1 − t i S , where S is an integer hyperparameter called partition.
Speciﬁcally, we proceed as follows:
[Equation]
where ( ˆp t i , ˆq t i ) (cid:44) ( p t i , q t i ) .
For the stiff equation, the step size of the numerical integration must be small enough for the stable solution.
Therefore, the large partition S should be chosen for the stiff portion while the small partition S is used for the nonstiff portion.
The parameter θ in H θ can be optimized by minimizing the mean squared error between the prediction and the ground truth:
[Equation]
with stochastic optimization algorithms such as Adam (Kingma & Ba, 2014).
Resampling.
As demonstrated in Section 1, the stiff phenomena may comprise only a small proportion in the trajectory of the Hamiltonian system.
We usually pick the stiff ratio γ with a value less than 0.5.
To avoid biased regression and unstable training caused by imbalanced categories of intervals, we adopt a random resampling technique that balances the data by replicating the stiff intervals K times.
A typical choice for replication K is 1 − γγ such that the number of the stiff intervals is comparable to that of the nonstiff intervals.
E XPERIMENTS
To evaluate the performance of SANN, we use two complex Hamiltonian systems: the billiard model and the three-body problem.
We compare SANN with two famous approaches, including
[Figure]
HNN (Greydanus et al., 2019) and SRNN (Chen et al., 2019).
All approaches aim to learn H θ ( p , q ) in Eqn.
(1) as a representation of the dynamics.
HNN trains H θ ( p , q ) by minimizing the error between the partial derivatives of H θ ( p , q ) and time derivatives approximated by data.
SRNN utilizes symplectic solvers to conduct multistep integration and trains H θ ( p , q ) by minimizing the error between the prediction and the data.
SANN uses a similar loss function but with stiffness awareness as shown in Section 3.2.
For a fair comparison, we use the Leapfrog solver for both
SRNN and SANN.
To evaluate the accuracy of the learned H θ ( p , q ) for different methods, we simulate the dynamics using the Leapfrog solver with the same settings during testing.
We provide a guidance for hyperparameter selection in Appendix G and additional experiments on the PendulumN problem which involves inseparable Hamiltonian in Appendix H.
4. B ILLIARD MODEL
The billiard model has wide application in real-world physical systems, spanning quantum-classical correspondence (St ¨ ockmann & Stein, 1990), lasers (Stone, 2010), quantum dots (Ponomarenko et al., 2008), and nanodevices (Chen et al., 2016).
The billiard model that we consider describes a billiard bouncing between a ring with soft boundaries (Choudhary et al., 2020).
The Hamiltonian is deﬁned as
[Equation]
where p ∈ R 2 and q ∈ R 2 are the momentum and the position of a billiard on a 2D plane, respectively; r α , r β are the radius of the outer circle and the inner circle, respectively; s is the softness of the boundary; and ( q 0 , 0) is a shift of the inner circle from the center.
As shown in Fig. 3(a), the potential function becomes sharp near the boundaries of the ring, causing a rapid change of the momenta of the billiard ball when it gets close to the boundaries.
Experimental setup.
The parameters r α , r β , s , and q 0 are set to be 1 . 0 , 0 . 5 , 0 . 05 , and 0 . 1 , respectively.
We use the physics-informed model H θ ( p , q ) deﬁned in Eqn.
(2); φ is a 3-hidden-layer fully connected NN with width 256 per layer and a rational activation function (Boull´e et al., 2020).
We simulate the training set with 100 trajectories and the testing set with 30 trajectories by RKF45.
Each trajectory contains 1,000 time steps with step size of 0.1.
We set the stiff ratio γ to be 10% and the partition S to be 10 for the stiff interval and 2 for the nonstiff interval.
The loss function (7) is optimized by Adam with a batch size of 1,024, and we use an initial learning rate of 0.001 forepochs.
The learning rate follows cosine decay with the increasing training epoch.
During testing, we integrate ODE (1) using the Leapfrog integrator with a ﬁxed step size of 0.005.
Results.
SANN can predict the dynamics more accurately than HNN and SRNN can.
Fig.
3(b, c) shows the mean square error (MSE) with time and relative energy error, which is deﬁned as the ratio of difference between the current energy and the initial energy to the initial energy.
We can see that SANN maintains a much smaller energy error compared with HNN or SRNN.
On average,
SANN reduces the relative energy error by 56.07% compared with SRNN and 92.13% compared with HNN.
The gap in terms of MSE enlarges as the number of time steps increases.
At the last time step, the MSE of SANN is 56.91% lower than that of SRNN and 83.44% lower than that of
HNN.
Fig.
4 gives a comparison of the predicted trajectories.
SANN recovers visually the same trajectories as the reference, while the trajectories obtained with SRNN and HNN clearly diverge from the reference.
Additional results for noisy data are given in Appendix F. A discussion on implicit bias is given in Appendix A.
4. T HREE - BODY PROBLEM
We consider a three-body problem that describes the motion of three particles under Newtonian gravitational force.
The Hamiltonian of the three-body system is given by
[Equation]
where p i ∈
R 2 , q i ∈ R 2 and m i are the momentum, position, and mass of the i th particle, i = 1 , 2 , 3 , p (cid:44) ( p 1 , p 2 , p 3 ) and q (cid:44) ( q 1 , q 2 , q 3 ) , and G is the gravitational constant.
One can see that the potential function incurs singularities when the distance between any two particles is small.
Experiment setup.
We set m 1 = m 2 = m 3 = 1 and G = 1 in our experiments.
To approximate the Hamiltonian of three-body dynamics, we use a physics-informed H θ as follows,
[Equation]
where φ is deﬁned as in Section 4.1, and we incorporate the physics information where the potential is a function of pairwise distances.
The training data comprises 1,000 trajectories, and the length of each trajectory N is 60 with the step size 0 .
The random initialization of initial states follows (Chen et al., 2019).
The testing data set consists of 300 trajectories with length N of 100 and step size of 0 . 1 .
We set the stiff ratio γ to be 10% and the partition S to be 50 for the stiff interval andfor the nonstiff interval.
When training, we minimize the loss function (7) by Adam with the batch size of 1,024.
H θ is trained with an initial learning rate of 0.001; the learning rate decays for 3, epochs.
During evaluation, the Leapfrog integrator with ﬁxed step size of 0.01 is used.
Results.
SANN learns a more accurate Hamiltonian than HNN and SRNN do on the three-body problem.
Fig.
6 shows the orbits of three particles simulated by different methods and a system energy comparison.
The relative energy error and MSE of orbits simulated by different methods are presented in Fig.
In Fig.
6, in the beginning the orbits of all methods coincide with the reference.
However, we see that the orbits of HNN and SRNN rapidly drift away from the reference, and the energy changes dramatically whenever two orbits meet.
On the contrary, SANN produces orbits that are nearly identical to the orbits of the reference, and the energy remains roughly constant.
From
Fig.
5, we see that SANN achieves lower energy error and MSE compared with HNN and SRNN.
On average, SANN reduces the relative energy error by 67.80% compared with SRNN and 96.30% compared with HNN.
Also, at the last time step SANN’s MSE is 63.84% lower than SRNN’s and 92.65% lower than HNN’s. is
S TIFFNESS M ETRIC In
Section 3.1 SAI was introduced in order to characterize the stiffness of a trajectory.
In this section we show theoretically and numerically that SAI can serve as a proxy of SI in a data-driven scenario.
SAI keeps the same trend as SI.
With the analytic expression of dynamics in three-body problem and a given trajectory, we compute the SI and SAI for each time step of the trajectory, respectively.
As a representative example, Fig.
7 shows the orbits of three particles and the corresponding curves of SAI and SI with time.
We can see that SI remains roughly a small constant when three particles are far away from each other but rises sharply when two particles get close.
In three-body dynam- ics, the stiff phenomena are known to occur at the close interaction of particles, which is in good correspondence to large SI values.
More important, SAI keeps the same trend with SI as time goes by, and SAI reaches the peak at the same time interval as SI does.
Through extensive numerical validation like this, we ﬁnd that SAI can characterize the stiffness of a trajectory very well.
SAI is a generalized SI.
For a theoretical analysis, we follow the standard technique that linearizes the ODE locally (Arrowsmith & Place, 1992).
In our method, SAI is computed by adjacent states with small step size, and the adjacent states satisfy the linearized ODE.
By deﬁnition, SI depends on the eigenvalue λ m of the linear ODE with the maximum absolute real part.
We show in Theorem(with proof in Appendix D) that SAI depends on all the eigenvalues.
In an ideal setting where (cid:96) m = 1 , (cid:96) i = 0 , ∀ i (cid:54) = m , SAI is roughly equivalent to SI.
Theorem 1.
Let A be a n × n symmetric real matrix and { λ i } ni =1 be its n distinct eigenvalues.
Then for the linear system d x dt =
Ax with nonzero initial condition x (0)
= u 0 , we have
[Equation]
(11) where u 1 = x (∆ t ) and (cid:96) i ≥ 0 , i = 1 , ..., n with (cid:80) ni =1 (cid:96) i = 1 .
One limitation of our stiffness metric is that the value of SAI may change with the coordinate system.
For example, a coordinate translation does not change the norm of the ﬁnite difference term in
Eqn.
(5), but may change the norm of the state.
However, the stiffness classiﬁcation in our method is not sensitive to coordinate translation because the classiﬁcation relies only on the ranking (not their values) of the SAIs for the time steps in a trajectory.
See Appendix B for a detailed analysis.
A BLATION STUDY
In this section we explore the effect of varied resampling replication K introduced in Section 3.2 and activation functions used in H θ ( p , q ) .
The experiments are conducted on the three-body problem.
Table 1 shows the training MSE, testing MSE, relative energy error, and training time (in seconds) for each epoch.
Fig.
5 shows the performance with and without resampling.
Fig.
8 displays the performance of different activation functions used in the NN.
Resampling.
From Exp. 2–Exp. 6, we see that with the increasing resampling replication for stiff portion, the learned Hamiltonian dynamics can achieve smaller errors, but the training time grows accordingly.
For example, when the replication K changes from 6 to 8, the testing MSE decreases by 47.61% but the training time increases by 22.72%.
Efﬁciency.
In Exp.
7 where all training intervals are integrated with partition 50, we achieve a more accurate solution compared with Exp. 2 where some portions of intervals are integrated with partition 10.
However, the training time increases signiﬁcantly.
On the other hand, despite the time increment caused by resampling, Exp. 4 achieves comparable performance to Exp. 7 but reduces over 40% of training time.
To trade off between the accuracy and time cost, we can choose an appropriate resampling replication.
Activation functions.
We compare the performance of ReLU, Tanh, and Rational (Boull´e et al., 2020) used in the NN to approximate the Hamiltonian.
Even though theoretically ReLU NN has a good approximation property, it has worse performance in our problem where the dynamics involve a ﬁrst-order derivative.
Rational NN is slightly better than Tanh NN.
According to Occam’s Razor (Lattimore & Hutter, 2011), we intuitively prefer to formulate the physical system with a simple expression and avoid a complicated one such as Tanh NN, a highly nonlinear function.
Rational NN, which has the form of a rational fraction, is more suitable for modeling the Hamiltonian.
[Table]
R ELATED W ORK
Learning Hamiltonian system from data.
Enforcing the conservative law into the system structure becomes a powerful and popular tool to learn the Hamiltonian system from data (Willard et al., 2020; Cheriﬁ, 2020; Zhong et al., 2021).
Greydanus et al.
; Choudhary et al.
(2020) use an
NN to approximate the H ( p , q ) instead of learning the dynamics directly.
This idea also applies to learning the conserved quantities from images (Toth et al., 2020).
To improve the accuracy of integration, Chen et al.
(2019) conduct multi-step integration with a symplectic solver.
Finzi et al.
(2020) simplify the learning process by coordinate transformation and enforcing the constraints of new coordinates.
To learn the Hamiltonian dynamics directly, Tong et al.
; Jin et al.
(2020) design NNs with a symplectic structure to characterize the system, while Chen & Tao (2021) learn a symplectic map of the Hamiltonian dynamics.
Zhong et al.
(2020) incorporate the physical prior into the dynamics parameterization.
Learning stiff dynamics.
It has been shown that stiffness may lead to failures in data-driven modelling (Wang et al., 2020;
Kim et al., 2021; Parmar et al., 2021).
Kim et al.
(2021) propose a scaling strategy that mitigates the stiffness of the dynamics to stabilize the gradient calculation.
This approach, however, is applicable only to problems where stiffness is caused by widely separated time scales and endures during the whole trajectory.
C ONCLUSION
We propose SANN (stiffness-aware neural network) to learn the stiff Hamiltonian system.
We also propose a new metric SAI (stiffness-aware index) to classify the training data into stiff and nonstiff portions.
This classiﬁcation along with a resampling technique allows us to apply step size adaptation strategies to better capture the dynamical characteristics of the Hamiltonian vector ﬁelds.
On complex physical systems including the three-body problem and the billiard model, our method outperforms the state-of-art approaches with a signiﬁcant margin.
Our method has potential to extend to other types of stiff dynamical systems, not limited to learning the Hamiltonian; such an extension is left as future work. A CKNOWLEDGEMENT
This material is based upon work supported by the U.S. Department of Energy, Ofﬁce of Science,
Ofﬁce of Advanced Scientiﬁc Computing Research, Scientiﬁc Discovery through Advanced Computing (SciDAC) program through the FASTMath Institute under contract DE-AC02-06CH11357 at
Argonne National Laboratory.
S L. acknowledges the support of the Ross-Lynn fellowship from Purdue University. A M ITIGATING IMPLICIT BIAS
In this section, we empirically demonstrate that our SANN can mitigate the implicit bias of NNbased optimization.
First, the potential functions learned by different methods are compared, and we can see that SANN can learn the potential function more accurately than other methods.
Then, we show the performance for excessive number of epochs to demonstrate that the baseline methods cannot capture the dynamics well even when a large number of training epochs is used.
A. L EARNED POTENTIAL FUNCTION
The stiffness of the billiard model is mainly caused by the sharp boundary of the potential function.
Hence, the learned potential function can reﬂect the accuracy of the simulated dynamics.
Fig.shows the top view and side view of the potential function of the billiard model from different methods.
The reference potential function becomes sharp near the boundaries of the ring.
From the top view, compared with SRNN and HNN, SANN learns a potential function that is close to the reference.
From the side view, the potential functions of SRNN and HNN tend to learn a smooth inner circle boundary, which indicates these methods suffer from the implicit bias mentioned in
Section 1.
In contrast, the edge and corner of the SANN potential function appear much sharper.
A. I NFLUENCE OF NUMBER OF EPOCHS
Ronen et al.
(2019) points out that it requires O ( k 2 ) time to learn a function of frequency k for
NN-based optimization.
However, we observe that the extending the training time cannot mitigate the implicit bias for SRNN or HNN when learning the stiff Hamiltonian systems.
We increase the number of epochs for training SRNN and HNN from 3 , 000 to 6 , 000 , 10 , 000 , and 15 , 000 .
Fig.
10 shows the training error versus different numbers of epochs, as well as MSE and the energy error of the simulated trajectories.
We can see that the training loss of SRNN converges to the same value even when the number of epochs becomes extremely large, and the MSE and the energy error do not improve much as the number of epochs grows.
Similar behaviours are observed for HNN.
B G ENERALITY OF SAI- BASED STIFFNESS CLASSIFICATION
In this section, we empirically demonstrate the SAI-based stiffness classiﬁcation is not sensitive to the translation of the coordinate.
The purpose is to show that this classiﬁcation strategy can identify the stiff part of the trajectory under different kinds of coordinate translation.
Speciﬁcally, we translate the position coordinates of the three-body trajectories with two kinds of translations, ﬁxed direction and random direction.
For ﬁxed direction, the position coordinates are added with a ﬁxed vector v = (0 , · · · , 0 , v, · · · , v ) , e.g., ( p t i , q t i ) + v .
For a random direction, the position coordinates are added to with a random vector v = (0 , · · · , 0 , v 1 , · · · , v n ) and { v i } are i.i.d random variables uniformly distributed on [0 , 1] .
Fig.
11(a) and 11(b) highlight the stiff intervals (marked in yellow) that our method identiﬁes using the SAIs calculated in different coordinate systems.
Three examples (one in each column) are chosen for illustration.
Fig.
11(a) shows the classiﬁcation results ( γ = 0 .
for the ﬁxed translation usingwidely different values of v ( v
=
Fig.
11(b) shows the classiﬁcation results foroccurrences of the random vector
v uniformly sampled from [0 , v ] ( v = 0 .
Fig.illustrates the same analysis for three additional examples.
We can see that (1) the SAI-based classiﬁcation captures the stiff parts of the trajectory (those with large SI values) successfully for all scenarios, and (2) coordinate translation has almost no impact on the classiﬁcation results.
This is expected by construction.
The second term of Eqn.
(5) (cid:13)(cid:13)
u ti +
− u ti t
i +1 − t i (cid:13)(cid:13) 2 is a translation-invariant quantity.
Coordinate translation changes only the norm of the states, so it causes the SAIs to be scaled for all the time intervals.
Therefore, when sorting the intervals based on SAI, we can maintain a similar order for these intervals and identify the stiff parts accurately.
C L EAPFROG INTEGRATION METHOD
The forward step of Leapfrog on separable Hamiltonian H ( p , q ) =
T ( p ) + V ( q ) is given as follows (Chen et al., 2019),
[Equation]
D P ROOF OF T HEOREM
Theorem 1 .
Let A be an n × n symmetric real matrix and { λ i } ni =1 be its n distinct eigenvalues.
Then for the linear system d x dt =
Ax with nonzero initial condition x (0)
= u 0 , we have
[Equation]
where u 1 = x (∆ t ) and (cid:96) i ≥ 0 , i = 1 , ..., n with (cid:80) ni =1 (cid:96) i = 1 .
Proof.
Let { v i } ni =1 be the eigenvectors corresponding to the eigenvalues { λ i } ni =1 .
Since A is a symmetric real matrix, { v i } ni =1 form a set of orthogonal vectors.
Without loss of generalization, we assume { v i } ni =1 is standard orthogonal basis.
Let x (0) =
u 0 = (cid:80) ni
=1 c i v i ; then (cid:107) ( c 1 , ..., c n ) (cid:107)
= (cid:107) u 0 (cid:107) 2 .
The solution of the linear system is x ( t ) =
(cid:80) ni
=1 c i v i e λ i t .
Therefore,
[Equation]
We let (cid:96) i (cid:44)
c 2 i (cid:80) ns =
1 c 2 s ≥ 0 and have (cid:80) ni =1 (cid:96) i = 1 .
E S UPPLEMENTARY R ESULTS
In this section we provide some supplementary results to the experiments of the main text.
Fig.
13 gives a comparison of the predicted trajectories of the billiard model, and this is a supplement to Fig. 4.
Again, we see that SANN recovers visually the same trajectories as the reference.
For the three-body problem, Fig. 14 and Fig.
15 show the orbits of three particles learned by different methods, and these results are a supplement to Fig.
We can see that HNN and SRNN diverge from the orbits after a period of time while SANN produces orbits that are nearly identical to the orbits of the reference and the energy is roughly conserved.
Fig.
16 displays the comparison of SAI and SI with time, and these results are a supplement to Fig.
These four examples in Fig.
16 support that the SAI keeps the same trend as SI and achieves the peak at the same time as SI does.
Reference
SANN (Ours)
SRNN
HNN
Reference
SANN (Ours)
SRNN
HNN
Reference
SANN (Ours)
SRNN
HNN
Reference
SANN (Ours)
SRNN
HNN
[Equation]
In this section we evaluate our method on the noisy training data on Billiard model.
We add the
Gaussian noise with different noise levels (standard deviation) to the clean training data.
Figuresand 18 respectively show the MSE and relative energy error with time under noisy data with different noise level.
We can see that all the learned dynamical systems become unstable and the energy gets nonconserved.
We see that HNN, SRNN, and SANN have difﬁculty in learning the chaotic dynamics when the data is disturbed by noise.
G H YPERPARAMETER SELECTION
We determine the stiff ratio γ based on the validation dataset.
We investigate the distribution about the SAI over a trajectory and estimate the ratio of the stiff portion for the dataset.
Take the dataset of the three-body problem as an example.
We normalize the SAIs of a trajectory to [0 , 1] and show the increasing order of SAIs in Fig.
One can see that large SAIs mainly concentrate around the order index from 90 to 100 .
Hence, we choose 90 -percentile ( γ = 0 .
1 ) as the threshold.
For the partition S for nonstiff and stiff portions, we perform the grid search for the pairs of parameters as did in Choudhary et al.
We use the validation loss as a proxy of the ﬁnal performance.
In particular, we ﬁrst conduct 10-epoch training and sort the pairs of parameters based on the validation loss.
Then, we choose the parameters with the smallest validation loss.
H I NSEPARABLE H AMILTONIAN
SANN can also be applied to inseparable Hamiltonian.
Like HNN, we use an NN that takes p and q as an input to approximate the Hamiltonian directly on two problems: a chain of N pendulums(Finzi et al., 2020) and the three-body problem.
In the Pendulum-N problem, let m i be the mass of the i -th pendulum, (cid:96) i be the rigid rod between ( i − 1) -th
and i -th pendulum, p i be the generalized momentum and q i be the angle between the i -th pendulum and the y -axis.
The Hamiltonian can be deﬁned as
[Equation]
where the mass matrix has a complicated form M ( q ) i,j = cos( q i − q j ) (cid:96) i (cid:96) j (cid:80)
Nk
[Equation]
max { i,j } m k .
We consider N = 4 and N = 6 , and set g
= m i = (cid:96) i = 1 , i = 1 , 2 , · · · , N .
We simulate the training set with 200 trajectories and the testing set with 100 trajectories by RKF45.
The initial states are uniformly and randomly sampled from [ − 0 .
Each trajectory contains 100 time steps with a step size of 0.03.
The three-body problem is described in Section 4.2.
It has a separable Hamiltonian, but we treat it as an inseparable Hamiltonian in this experiment.
We compare SANN with HNN using these two problems.
For both methods, the Euler method is used for training while RKF45 is used for predicting.
For SANN, we set the partition parameter S to be 50 and 10 for the N -chain pendulum and the three-body problem, respectively.
Table 2 shows the train MSE and test MSE for SANN and HNN.
We can see that SANN outperforms HNN in all the tests.
[Table]