Stiffness-aware neural network for learning Hamiltonian systems
Authors: SENWEI Liang
Zhongzhan Huang
Hong Zhang
Venue: ICLR 2022 Poster (2022)
Keywords: Hamiltonain systems
Neural network
Stiff dynamical systems
Data-driven method

Abstract: We propose stiffness-aware neural network (SANN), a new method for learning Hamiltonian dynamical systems from data. SANN identifies 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 classification 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 fields. 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 significant improvement in accuracy.

Published as a conference paper at ICLR 2022
STIFFNESS- NEURAL NETWORK FOR LEARNING HAMILTONIAN SYSTEMS
Senwei Liang Purdue University
Zhongzhan Huang Sun Yat-sen University
Hong Zhang Argonne National Laboratory hongzhang@an goo
1iang 3eguu33ueuuuu edu
huangzhzh230 ma Saall..yysu.e. www.ssueedu.co cn
ABSTRACT
We propose stiffness-aware neural network (SANN). new method for learning Hamiltonian dynamical systems from data. SANN identifies and splits the train- ing data into stiff and nonstiff portions based on a stiffness-awars index, a simple, yet yet effective metric we introduce to quantify the stiffness of the dynamical system. This classification along with resampling technique allows us to apply different time integration strategies such as step size adaptation to better capture the dy- namical characteristies of the Hamiltonian vector fields. We evaluate SANN on complex physical systems including a three-body problem and billiard model. We show that SANN more stable and can better preserve energy when compared with the ofte-of-the-art methods, leading to significant improvement in accuracy.
INTRODUCTION
Data-driven modeling of dynamical systems provides a computationally inexpensive approach for exploiting the scientific law of the physical processes and predicting future phenomena (Giannakis)
& Majda, 2014) Harlim et al., 2021, Gu et al. 2021 Mou et al. 202T). With the superior approx- imation and generalization capacity (Luu et al. 2020 Shen et al. 2021) of neural networks (NNs),
important advances have been made in learning Hamil summitoniann 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 learn- ing this conserved quantity from data (Greydanus et al., 2019 DiPietro 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 let al. 2020 Marchal 1990). especially for complex systems that are intrinsically stiff. When learn- ing 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. 202T Wang et al. 2020).
To illustrate the difficulty in learning stiff Hamiltonian dynamics, let consider a three-body prob- lem that describes the interactions of three particles under gravitational force. According to phys- ical law, the repulsive force of two particles increases dramatically when they get close to each other. Hence, the close three-body interaction intensifies stiffness phenomena (Huang & Leimkuh- ler 1997). Fig. [ shows the reference orbits of three partieles and the orbits learned by the Hamilto- nian neural network (HNN) (Greydanus et al. 2019). famous approach that directly approximates the Hamiltonian with a neural network. In Fig. [., 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) ee al. 2019). However, the particles deviate from the reference orbits rapidly after close encounter, and the energy becomes nonconserved (Fig. 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 fitting a smooth function with the fast decay in the frequency domain (Xu al., 2220 Cao et al. 2019). This implicit bias impedes components, such as singularities in an N-body problem.
the NN from capturing
Correeponddnnce should be addressed hongzhang@an .gov

---

Published as a conference paper at ICLR 2022
Example 1
(b)
(C)
[d]
system energy
1996
Time step 32
Time step 58
Time step 64
Time step = 76
time step
45
100
Example 2
(o)
(h)
system energy
-5
-22
-10
-83
step
step
step
step
time step
Reference
HNN (Dotted Line)
Starting point
Direction
Figure I: First four columns show comparison of the reference orbits and the Hamiltonian neural network (HNN) orbits of three particles. The orbits of the different particles are displayed with dif- ferent colors. The rightmost figures show the energy comparison. The reference orbits are simulated by the RKF45 solver with the ground truth Hamiltonian, and the energy is conserved. The energy of the HNN changes dramatically at the close interaction of two particles (see c, f), and eventually the HNN orbits diverge from the reference.
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 small proportion of trajectories corresponds to stiff dynamics. As statistical example, the 1, 000 independent simulations 2019). 91.4% of the trajectories contain close encounter, but average only 4.2% of the time intervals within the trajectories contain close encounter.
three-body trajectories following (Chen et al.
To mitigate the implicit bias caused by NN-based optimization and the imbalanced stiffness propor- tion problem, we propose a new method called the stiffiess-aware neural network (SANN) based on the stiffness classification of the training data. In approach, introduce stiffness- aware index (SAI) as a simple, yet effective metric to classify the time intervals into stiff and nonstiff portions. For training efficiency, we integrate the Hamiltonian dynamics over different intervals with different step sizes based on their classification. 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.
2. We validate the SANN method with complex Hamiltonian dynamics including three-body prob- lem and billiard model. Extensive numerical results show that SANN can accurately predict the the stiff dynamics and significantly outperform the existing methods.
PRELIMINARIES
2.1 HAMILTONIAN SYSTEM
The Hamiltonian system describes the continuous evolution of states in the phase space (p, q), where E R represents the generalized momentum and q Ra denotes position coordinates. The Hamiltonian H(p. q) 22 - TR. is denoted as the total energy of the system at (p, q). The dynamics can be described with H(p q) by
dq 0H dt
dp
0H
(1)
3p
dt
dq
With Eqn. (1. H(p q) is conservative during the evolution du ox classical mechanics, the Hamiltonian is usually expressed as the sum of the kinetic and potential
dp 8 dg 0. In
ap dq
2

---

Published as a conference paper at ICLR 2022
Data Processing
Hamiltonian Help.g
Numerical Integration
Stiff?
Burc Buicueeee
p'
Mp
Kinetic energy
P p
Use large Step size
Step size
Potential energy
ODE Solver
Trajectory
Figure 2: Workflow of SANN.
energies. We interested in learning the separable Hamiltonian, namely, H(p. T(p) + V(q), where T(p) kinetic energy while V(q) is potential energy. In this paper our goal is to lear the Hamiltonian of dynamical system. What is available the time series of observations {p",, q²) )) The Hamiltonian is parameterized by neural network as follows,
k(p. q) = Ho(p. q) p Mp + p(q: W),
where Φ R + I is fully connected NN, M E R is a trainable matrix, and 0 {M, W set of all trainable parameters. Once we obtain He. the trajectories of the learned Hamiltonian system can be simulated by Eqn. ① with an ordinary differential equation (ODE) solver.
2.2 ODE SOLVER FOR THE HAMILTONIAN SYSTEM
We use Forward qº), At) denote a one-step integration of Eqn. (1 from (p qD) over the step size At. The Euler method first-order Runge- Kutta method adopted in (Greydanus et al. 2019). When the Fuler method is used, the integration becomes
(p 1+1 q 1+11 Forward((p'' ,''), H. 4.41 (.) 2 (p' (p",,""')+)" +- OH 8H
(1_- - 4.).
0q ōp
The Fuler 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- error-based step size control has great potential for improving computational efficiency. 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 combined ODE system and estimating the error on all batch elements, as noted in Chen et al.) (2018) 8). Moreover, the step size is limited by the stiffest element in a batch, making it difficult to use large step size especially when the batch size is large.
A symplectic integrator, numerical method that conserves the energy quantity, is widely used learning Hamiltonian system (Chen et al. 2019: Zhong et al. 2020). Leapfrog is a second-order symplectic integrator designed for separable Hamiltonian. The Leapfrog scheme is in Appendix C
3 STIFENESS-AWARE NEURAL NETWORK
In this section we introduce our method, called SANN (stiffness- neural network), with its workflow shown Fig. 2 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 classification. Next, we resample the stiff intervals balance their ratio and avoid biased training.
3.1 IDENTIFYING THE STIFF INTERVAL
In this part we first discuss the stiffness index (SI) that is used to characterize stiffness for an ODE and propose the diffness-awaar index (SAD) 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.
du
dt f(u).
(4)
3

---

Published as a conference paper at ICLR 2022
SI at the state utt) is defined by max (R((),)) where A 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 Equ. (4) is linear system with constant coefficients, namely, di Au, and A E R is diagonalizable matrix with eigenvalues (Ar) and corresponding eigenvectors (o; )i= Then the solution of the linear system is (t) = Σx= C, Let us suppose that Re(A,) < 0, i = I,... n. We have exst + as - co. Hence, max { Re():))} reveals the fastest speed of decaying to O. max Re (A,)l) is large, the integrator needs small step size to reduce the local truncation error.
Definition of SAL Given the observations { from the system (4), we define SAI at u'x by
(5)
Intuitively, SAI characterizes the relative changing speed of the state u from time t; to f,+1- The norm of finite difference || - --+ ||2 approximates f(xz) 1l2, which is used in time parameter- ization of the adaptive time step method (Huang & Leimkuhler 1997). In our data-driven setting, we cannot compute SI as the analytic expression of f unavailable. But SAI can serve as a proxy of ST in time series data, which will be demonstrated in Section
Next, we split the training data into stiff and nonstiff portions by classifying the time interval as either stiff or nonstiff based on SAL. For time series observation (u" J.N. Ji=1. first compute SAI for each time interval. Let SAT; = |u u
We then rank { The interval ="--"=""
f411-, ll2.
(t,, te+1) with larger SAI, is supposed be stiffer part. We y (0. 1) denote the stiff ratio, a hyperparamettr to determine the ratio of the stiff portion. The intervals belonging to the top-y of [SAI, 1 are identified as stiff intervals, and the others are classified as nonstiff. In other words, the intervals in the (1 7) 100-th percentile of SAI are identified nonstiff intervals.
3.2 TRAINING THE ESTIMATED HAMILTONIAN
Loss function. Let {(p", hop"s,00 )
Jizt be the given trajectory from Hamiltonian system, Given pair of consecutive states, (p"i, ,4') and (p" qh*+), we integrate Eqn. ① from i; to obtain the estimated solution (f q at 1 To achieve more accurate solution, we conduct the integration over
[t,, t441 with small time step 511-4, where S is an integer hyperparameter called partition. Specifically, we proceed as follows:
l+t1 - 12 1,... ,, (6)
(pitais ())))/s) (p "" q/=44-118)) Ho.
S
where (p" )(t) (p't,q*) q²)). 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 Ho can be optimized by minimizing the mean squared error between the prediction and the ground truth:
N-1
Ip
ph 112 + f q
N
with stochastic optimization algorithms such as Adam KKingma Ba 2014).
Resampling, As demonstrated in Section D the stiff phenomena may comprise only small pro- portion in the trajectory of the Hamiltonian system. We usually pick the stiff ratio y 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 in- tervals K times. A typical choice for replication K is 1=x such that the number of the stiff intervals is comparable to that of the nonstiff intervals.
4 EXPERIMENTS
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
4

---

Published as a conference paper at ICLR 2022
(a)
(b) Billlard (energy)
Billiard (MSE)
0.5
SANN (Ours)
SANN [Ours)
Potentiall
HNN
HNN
SRNN
0.4
SRNN
MAE
Reletoe
0.1
step
Figure 3: (a) Profile the potential function for the billiard model; q q (gx: qy) E R2 denotes the position. Comparison of (b) relative energy enror and (c) MSE for billiard model.
(Ours)
HNN
Reference
SANN (Ours)
SRNN
HNN
Time
Comparison of billiard orbits simulated using the Hamiltonian learned by different meth- ods. We can see that SANN produces orbits that are almost identical to those of the reference. We provide results in Appendix and dynamic graphs website
Figure 4:
HNN (Greydanus et al. 2019) and SRNN (Chen et al. 2019). All approaches aim to learn He q) Eqn. (1) a representation of the dynamics. HNN trains He (p. q) by minimizing the error between the partial derivatives of (p. q) and time derivatives approximated by data. SRNN utilizes symplectic solvers to conduct multistep integration and trains Ho q) by minimizing the error between the prediction and the data. SANN uses 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 He(p, q) for different methods, simulate the dynamics using the Leapfrog solver with the same settings during testing. We provide guidance for hyperparameter selection in Appendix G and additional experiments on the Pendulum- problem which involves inseparable Hamiltonian in Appendix
4.1 BILLIARD MODEL
The billiard model has wide application real-world physical systems, spanning quantum-classical correspondence (Stöckmann & Stein 1990). lasers (Stone) 2010). quantum dots (Ponomarenko ct al. 2008), and nanodevices (Chen et al.) 2016). The billiard model that we consider describes
billiard bouncing between ring with soft boundaries (Choudhary et al. 2020 The Hamiltonian defined as
H(p, q) - 11p1 + (1 exp (" -a- - =aale))))
| (go. 021l))
exp
(8)
5

---

Published as a conference paper at ICLR 2022
where p R² and q IR² are the momentum and the position of billiard on a 2D plane, respec- tively; Tas Tg are the radius of the outer circle and the inner circle, respectively; s is the softness of the boundary; and (4o, 0) is shift of the inner circle from the center. As shown in Fig. (3a),, 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 Ta, ra. 8, and O0 are set to be 1.0, 0.5, 0.05, and 0.1, respec- tively. We use the physics-informed model (p, q) defined in Egn. 2): è a 3-hidden-layer fully connected NN with width 256 per layer and a rational activation function (Boullé 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 Š be 10 for the stiff interval and 2 for the nonstiff interval. The loss function (7) is optimized by Adam with batch size of 1,024, and we use an initial learning rate of 0.001 for 500 epochs. The learning rate follows cosine decay with the increasing training epoch. During testing, we integrate ODE using the Leapfrog integrator with fixed step size of 0.005.
Results. SANN can predict the dynamics more accurately than HNN and SRNN can. Fig. 3 b.
<) shows the mean square error (MSE) with time and relative energy error, which is defined as the ratio of difference between the current energy and the initial energy to the initial energy. We can see that SANN maintains much smaller energy error compared with HNN 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 implicit bias is given in Appendix A
A discussion on
Three-body (energy)
Three-body (MSE)
Three-body (energy)
Three-body (MSE)
Relative energy error
SRAN
GAAa: [fum))
GoARi (Sum))
Reletive energy errar
SJMUUN tor thout resampingg
SAMUN for Un resarpging!
S/00 En "ssempingg
Subal to thou: resamoing:
MSE
step
step
Time step
step
Figure 5: The first two figures show the comparison of relative energy error and MSE with varied time for the three-body problem. The last two figures show the comparison of performance with and without resampling.
4.2 THREE-BODY PROBLEM
We consider three-body problem that describes the motion of three particles under Newtonian gravitational force. The Hamiltonian of the three-body system is given by
3
1
Gm.my
H(p. q) =
Σ
(9)
q 9
15KjSS
where P. E R2, qi E R and my are the momentum, position, and mass of the ith particle, i = 3. e (pt, P2, Pa) and q a (qı, (qı,q2, 92, qr), 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, = =m = my = and G = 1 in our experiments. To approximate the Hamiltonian of three-body dynamics, we use physics-ifformed Ha as follows,
Ho q) - p Mp + Φ[[ - 92||2, || - qa /2. || 4all2; W).
(10)
where & is defined as in Section[4.1 and we incorporate the physics information where the potential function of pairwise distances. The training data comprises 1,000 trajectories, and the length

---

Published as a conference paper at ICLR 2022
of each trajectory N is 60 with the step size 0.1. 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 y to be 10% and the partition S to be 50 for the stiff interval and 10 for the nonstiff interval. When training, we minimize the loss function 7 by Adam with the hatch size of 1,024. Ho is trained with an initial learning rate of 0.001; the learning rate decays for 3,000 epochs. During evaluation, the Leapfrog integrator with fixed step size of 0,01 is used.
Results. SANN learns more accurate Hamiltonian than HNN and SRNN do on the three-body problem. Fig. 6 shows the orbíts of three particles simulated by different methods and system energy comparison. The relative energy error and MSE of orbits simulated by different methods are presented in Fig. 5 In Fig. [ 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, the last time step SANN's MSE is 63.84% lower than SRNN's and 92.65% lower than HNN's. s.
SANA (Ount)
2:3
HAV
1.0
HNN
SRNN
SANN (Ours)
00
200
Time step
Referencc
SA (Ous)
5214
0.0
HNN
SRNN
SANN (Ours)
Time step
System
SANN (Ours)
Time step
Figure 6: The first three columns show the comparison of the reference orbits (solid curves) and the orbits learned by different method (dotted curves). is the initial position, and is the direction. The orbit of the different particles presented by different colors. The rightmost figures show the energy comparison. We provide more results Appendix E and dynamic graphs in website
STIFENESS METRIC
In Section 3.1 SAI introduced in order to characterize the stiffness of a trajectory. In this section we show thcoretically 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 representative example, Fig. 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 particles get close. In three-body dynam-
7

---

Published as a conference paper at ICLR 2022
Example
Example
Trrowbodyy
Timme
Figure 7: Comparison of SAI and SI with time. The color markers refer to the occurrence of the stiffness phenomena; the same color corresponds to the same time period. The trends of SAI and SI consistent. We provide more results on Appendix E and dynamic graphs in website
ics, the stiff phenomena are known to occur at the close interaction of particles, which is in good correspondence to large values. More important, SAI keeps the same trend with SI as time goes by. and SAI reaches the peak the same time interval SI does. Through extensive numerical validation like this, we find that SAI can characterize the stiffness of a trajectory very well.
SAI is 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 definition, SI depends on the eigenvalue Am of the linear ODE with the maximum absolute real part. We show in Theorem (with proof in Appendix D) that SAI depends all the eigenvalues. In ideal setting where fm = 1, & = 0, vi # m, SAI is roughly equivalent to SI.
Theorem 1. Let A be a n n symmetric real matrix and {A,);-1 be its n distinct eigenvalues. Then for the linear system
dæ
de Ax with nonzero initial condition x00) u", we have
u
(x) 1/2 + O(A))
(11)
1
At
1=1
where u <<tt) and l 2 0, 1, n. with Σ 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 finite difference term in Eqn. (5). but may change the norm of the state. However, the stiffness classification in our method is not sensitive to coordinate translation because the classification relies only on the ranking (not their values) of the SAIs for the time steps in a trajectory. See Appendix B for detailed analysis.
6 ABLATION STUDY
In this section we explore the effect of varied resampling replication K introduced in Section 3.2 and activation functions used in Hp 9). The experiments are conducted the three-body problem. T shows the training MSE, testing MSE, relative energy error, and training time (in seconds) for each epoch. Fig. S shows the performance with and without resampling. Fig. g displays the performance of different activation functions used in the NN.
Resampling, From Exp. 2-Exp. 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%.
Efficiency. In Exp. 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 significantly. 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é et al. 2020) used in the NN to approximate the Hamiltonian. Even though theoretically ReLU NN has
8

---

Published as a conference paper at ICLR 2022
good approximation property, it has worse performance in our problem where the dynamics involve a first-order derivative. Rational NN is slightly better than Tanh NN. According to Occam's 's Ra- (Lattimore & Hutter 201T), intuitively prefer to formulate the physical system with simple expression and avoid a complicated one such as Tanh NN, a highly nonlinear function. Rational NN, which has the form of rational fraction, is more suitable for modeling the Hamiltonian.
Exp.
Stilt Ratio ?
Partition (nomstill)
Purtition (still)
Replication
Train MSE
Test MSE
Energy Eruo
Training Time
Exp.
0.1
10
2.00-6
0.0513
0.0611
13:52
Exp.
0.1
10
2.0e-7
0.2173
0.7884
10.08
Exp.
0.1
10
50
4
1.5e-7
0.0460
0.1026
19.82
Exp.
0.1
10
6
7.4e-8
0.0294
0.0581
26.36
Exp.
0.1
8
5.5e-8
0.0154
0.0228
32.35
Exp.
0.1
9
4.36-8
0.0234
0.0457
35.35
Exp. 7
0.5
50
50
6,20-7
0.0554
0.0900
34.75
Table 1: Performance comparison under different parameter settings. "Exp." is short for experiment.
Three-body (MSE)
Three-boce body (energy)
Billiard (MSE)
Billiard (energy)
Ratiorw
Ratlon
Relative we er Everrg error
RaUA Tarh
Tarh
Reletive emergy
logiMSEi
MSE
sc
Time step
Time step
Time Step
Time step
Figure 8: Performance comparison of different activation functions the three-body problem and billiard model.
RELATED WORK
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 Cherifi, 2020 Zhong et al., 202T). Greydanus et al. (019);; 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 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. (2021) )); Jin al. (2020)
design NNs with a symplectic structure to characterize the system, while Chen & Tao (2021) learn symplectic map of the Hamiltonian dynamics. [Zhong al. (2020) incorporate the physical prior into the dynamics parameterization.
Learning stiff dynamics. It has been shown that stiffness may lead failures in data-driven mod- elling (Wang et al., (2200) Kim et al. [222]] Parmar et al. [2021]. Kim et al. (202T) propose a scaling strategy that mitigates the stiffness of the dynamics stabilize the gradient calculation. This ap-
proach, however, is applicable only to problems where stiffness is caused by widely separated time scales and endures during the whole trajectory.
CONCLUSION
We propose SANN (stiffness-aware neural network) to learn the stiff Hamiltonian system. We also propose new metric SAI (stiffness-aware index) to classify the training data into stiff and nonstiff portions. This classification along with a resampling technique allows us to apply step size adapta- tion strategies to better capture the dynamical characteristics of the Hamiltonian vector fields. On complex physical systems including the three-body problem and the billiard model, our method out- performs the state-of-art approaches with significant margin. Our method has potential to extend to other types of stiff dynamical systems, not limited to learning the Hamiltonian; such an extension left as future work.

---

Published as a conference paper at ICLR 2022
ACKNOWLEDGEMENT
This material is based upon work supported by the U.S. Department of Energy, Office of Science, Office of Advanced Scientific Computing Research, Scientific Discovery through Advanced Com- puting (SciDAC) program through the FASTMath Institute under contract DE-AC02-06CHI1 1357 at Argonne National Laboratory. acknowledges the support of the Ross-Lynn fellowship from Purdue University.
REFERENCES
Richard Aiken. Stiff computation, volume 169. Oxford University Press Oxford, 1985.
David Arrowsmith and Colin M Place. Dynamical systems: differential equations, maps, and chaotic behaviour, volume 5. CRC Press, 1992.
Nicolas Boullé, Yuji Nakatsukasa, and Alex Townsend. Rational neural networks. arXiv preprint arXiv:2004.01900 2020.
Yuan Cao, Zhiying Fang. Yue Wu, Ding-Xuan Zhou, and Quanquan Gu. Towards understanding the spectral bias of deep learning. arXiv preprint 1912.01 2019.
Renyi Chen and Molei Tao. Data-driven prediction of general hamiltonian dynamics via learning Exactly-symplecte maps. International Conference on Machine Learning, pp. 1717-1727. PMLR, 2021.
Ricky T. Q. Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural Ordinary Dif- ferential Equations. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa- Desa-Bianchi, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 31. Curran As- sociates, Inc., 2018. URL URL/https::////oceediinn neur
neur ips. htt.ppaaprr201//1110 pdf
Shaowen Chen, Zheng Han, Mirza M. Elahi, K. M. Masum Habib, Lei Wang. Bo Wen, Yuanda Gao, Takashi Taniguchi, Kenji Watanabe, James Hone, Avik W. Ghosh, and Cory Dean. Electron optics with p-n junctions in ballistic graphene. Science, 353(630) 2016. doi: 10. 1126/science..af5488...
Zhengdao Chen, Jianyu Zhang, Martin Arjovsky, and Léon Bottou. Symplectic recurrent neural networks. arXiv preprint urXiv:1909) 13334, 2019.
Karim Cherifi. An overview recent machine learning techniques for port hamiltonian systems. Physica D: Nonlinear Phenomena, 411:132620, 2020.
Anshul Choudhary, John F Lindner, Elliott G Holliday, Scott T Miller, Sudeshna Sinha, and William Ditto. Physics-enhanced neural networks learn order and chaos. Physical Review E, 101(6):062007, 2020.
Daniel DiPietro, Shiying Xiong, and Bo Zhu. Sparse Symplectically Integrated Neural Networks. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp.
Curran Associates, Inc., 2020.
Erwin Fehlberg. Low-order classical Runge-Kutta formulas with stepsize control and their applica- tion to some heat transfer problems, volume 315. National aeronautics and space administration, 1969.
Mare Finzi, Ke Alexander Wang, and Andrew G Wilson. Simplifying Hamiltonian and Lagrangian Neural Networks via Explicit Constraints. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 13880- 13889. Curran Associates, Inc., 2020.
Dimitrios Giannakis and Andrew J Majda. Data-driven methods for dynamical systems: Quantifying predictability and extracting spatiotemporal patterns. Wiley & Sons, 2014.
10

---

Published as a conference paper at ICLR 2022
Samuel Greydanus, Misko Dzamba, and Jason Yosinski. Hamiltonian neural networks. In Advances in Neural Information Processing Systems, volume 32, 2019.
Yiqi Gu, John Harlim, Senwei Liang, and Haizhao Yang. Stationary Density Estimation of It)
Diffusions Using Deep Learning. arXiv preprint arXiv:2109 N..099200.03392,, 2021.
John Harlim, Shixiao Jiang, Senwei Liang, and Haizhao Yang. Machine learning for prediction with missing dynamics. Journal of Computational Physics, 129:22, 2021.
Weizhang Huang and Benedict Leimkuhler. The adaptive Verlet method. SIAM Journal on Scientific Computing, 18(1):2392556 1997.
Pengzhan Jin, Zhen Zhang, Aiqing Zhu, Yifa Tang, and George Em Karniadakis. SympNets: In- trinsic suructure-preservve symplectic networks for identifying Hamiltonian systems. Neural Networks, 132:166-1793 179, 2020.
Suyong Kim, Weiqi Ji, Sili Deng, and Christopher Rackauckas. Stiff neural ordinary differential equations. arXiv preprint arXiv:2103.15341, 2021.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.660 42860, 2014.
Tor Lattimore and Marcus Hutter. No Free Lunch versus Occam's Razor in Supervised Learning, In Algorithmic Probubility and Friends, 2011.
Lu, Z. Shen, Yang, and S. Zhang. Deep Network Approximation for Smooth Functions. arXiv e-prints, arXiv:2001.055 003000000,03040,, 2020.
C. Marchal. The Three-body Problem. Advances in Industrial Engineering. Elsevier, 1990. ISBN 9780444874405.
Changhong Mou, Birgul Koc, Omer San, Leo G Rebholz, and Traian Iliescu. Data-driven variational multiscale reduced order models. Computer Methods in Applied Mechanics and Engineering, 373: 13340, 2021.
Mihir Parmar, Mathew Halm, and Michael Posa. Fundamental challenges in deep learning for stiff contact dynamics. arXiv preprint arXiv:2103. 15406, 2021.
L. A. Ponomarenko, F. Schedin, M. I. Katsnelson, R. Yang, E. W. Hill, K. S. Novoselov, and A. K. Geim. Chaotic Dirac Billiard Graphene Quantum Dots. Science, 32065874) 2008 doi: 10.1 10.1see/sce 1154663.
Basri Ronen, David Jacobs, Yoni Kasten, and Shira Kritchman. The convergence rate of neural net- works for learned functions of different frequencies. Advances in Neural Information Processing Systems, 32:4761-4771, 2019.
Zuowei Shen, Haizhao Yang, and Shijun Zhang. Deep Network with Approximation Error Being Reciprocal of Width to Power of Square Root of Depth. Neural Computation, 2021.
H-J Stöckmann and J Stein. Quantum chaos in billiards studied by microwave absorption. Physical review letters, 64(19):2215, 1990.
A Douglas Stone. Chaotic billiard lasers. Nature, 665(7299):666666977 697, 2010.
Yunjin Tong, Shiying Xiong, Xingzhe He, Guanghan Pan, and Bo Zhu. Symplectic neural networks in Taylor series form for Hamiltonian systems. Journal of Computational Physics, 437:110325 2021. ISSN 0021-9991.
Peter Toth. Danilo Rezende, Andrew Jaegle, Sébastien Racanière, Aleksandar Botev, and Irina Higgins. Hamiltonian generative networks. In International Conference on Learning Represen- tations, 2020. URL
Sifan Wang, Yujun Teng, and Paris Perdikaris. Understanding and mitigating gradient pathologies in physics-ifformed neural networks. arXiv preprint 1rXXi:2001.045366 2020.
11

---

Published as a conference paper at ICLR 2022
Jared Willard, Xiaowei Jia, Shaoming Xu, Michael Steinbach, and Vipin Kumar. Integrating physics-based modeling with machine learning: A survey. arXiv preprint arXiv:2003.04919, 1(1):1-34, 2020.
Zhi-Qin John Xu, Yaoyu Zhang, Tao Luo, Yanyang Xiao, and Zheng Ma. Frequency Principle: Fourier Analysis Sheds Light on Deep Neural Networks. Communications in Computational Physics, 286))1176-11777 2020. ISSN 1991-7120.
Yaofeng Desmond Zhong, Biswadip Dey, and Amit Chakraborty. Symplectic ODE-Net: Learning Hamiltonian Dynamics with Control. In International Conference on Learning Representations, 2020.
Yaofeng Desmond Zhong, Biswadip Dey, and Amit Chakraborty. Benchmarking monsy-conservvng neural networks for learning dynamics from data. In Learning for Dynamics and Control, pp- 1218-1229, PMLR, 2021.
12

---

Published as a conference paper at ICLR 2022
A MITIGATING IMPLICIT BIAS
In this section, we empirically demonstrate that our SANN can mitigate the implicit bias of NN- based 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 large number of training epochs is used.
A.1 LEARNED POTENTIAL FUNCTION
The stiffness of the billiard model is mainly caused by the sharp boundary of the potential function. Hence, the learned potential function reflect 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 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 learn a smooth inner circle boundary, which indicates these methods suffer from the implicit bias mentioned in Section contrast, the edge and corner of the SANN potential function appear much sharper.
the boundaries of the ring. From
A.2 INFLUENCE OF NUMBER EPOCHS
Ronen et al 2019 points out that it requires ( time learn function of frequency k for NN-based optimization. However, we observe that the extending the training time cannot mitigate the implicit bias for SRNN 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. I0 shows the training error versus different numbers of epochs. well MSE and the energy error of the simulated trajectories. We can see that the training loss of SRNN converges the same value energy error do not improve much as the number of epochs grows. Similar behaviours are observed for HNN.
when the number of epochs becomes extremely large, and the MSE and the
Raterence
SÁNN
SRNN
HNN
Side View
8
3r
Reference
SANN
SRNN
HNN
Top View
4x
Figure 9: The learned potential function of the billiard model.
13

---

Published as a conference paper at ICLR 2022
Three-body (Train loss)
Three-body (MSE)
Three-body (energy)
spode
Relative energy errOT
spache STIN LOC essore
STAN 30c opechn
Train loss (log)
opects
MSE
0.5
1913 200 spocts
1954
8.5
P-NS 224 epoctss
speche
epocta
(a)
2002 40c0 acco D000 20000 34000 Epochs
(b)
2
Time step
100
(C)
Time step
200
Figure 10: Comparison different methods in training loss, MSE, relative energy error for different numbers of training epochs.
14

---

Published as a conference paper at ICLR 2022
B GENERALITY OF SAI-BASED STIFENESS CLASSIFICATION
In this section, we empirically demonstrate the SAI-based stiffness classification is not sensitive to the translation of the coordinate. The purpose is to show that this classification strategy can identify the stiff part of the trajectory under different kinds of coordinate translation.
Specifically, translate the position coordinates of the three-body trajectories with kinds of translations, fixed direction and random direction. For fixed direction, the position coordinates are added with fixed vector v (0. 0, v, ....... v), e.g., (p" p"i,q' q")) + V. For random direction, the position coordinates are added to with random vector = (0, ... v1, Un) and {#t are i..d random variables uniformly distributed on [0. 1].
Fig. [1 a) and [Tbb highlight the stiff intervals (marked in yellow) that our method identifies using the SAIs calculated in different coordinate systems. Three examples (one in each column) are chosen for illustration. Fig. 1 a)) shows the classification results (x = 0.1) for the fixed translation using 7 widely different values of v (v = 0.1, 10, ..................... 10%). Fig- 11bb b) shows the classification results for
occurrences of the random vector uniformly sampled from [0, u] (v 0.1, 10, 10%). Fig. 12
illustrates the same analysis for three additional examples.
We can see that (1) the SAI-based classification captures the stiff parts of the trajectory (those with large SI values) successfully for all scenarios, and (2) coordinate translation has almost no im- pact the classification results. This is expected by construction. The second term of Eqn. (5) u l| is translation-invariant quantity. Coordinate translation changes only the norm of 611-6 the states, so it causes the SAIs to be scaled for all the time intervals. Therefore, when sorting the intervals based on SAI, we maintain similar order for these intervals and identify the stiff parts accurately.
Example
Example
Example 3
(8)
Tie sirp
Tive
(b)
Time Mtt
Fime
Tiee
(c)
5M food tradeiiee Je33
(Tack buriitaane 1ul1
moation 342
Tran (A01
ioad tardasion Ie3)
Food hunslatee Jodi
S4 ffoee transition
54 (Tai trmntini (an)
step
Time step (t)
Time step (t)
Figure 11: (a) SAI-based classification for trajectory translated by fixed vector for different
values of v (v = 0.1, 10.... ... 109). Stiff intervals are marked in yellow. (b) SAI-based classifica-
tion for trajectory translated by uniform random vector v [0, v] (u 0.1, 10..... 110)). (c)
Comparison of SI, SAI of the original trajectory, and SAI of the translated trajectory.
15

---

Published as a conference paper at ICLR 2022
Example
Example
Example
Thee
Time waa
(b)
(c)
rhord Moed Iopolation AM Food trnoiatie let
FFod tram mranitiion 3e11 (faed
Hird translater Te01
Hhaellia muilher Soll
flord crandattor 1o4)
Time step
step (t)
Time step (t)
Figure 12: (a) SAI-based classification for trajectory translated by fixed vector for different
values of v (v = 0.1, ...,, 10%). Stiff intervals are marked in yellow. (b) SAI-based classifica-
tion for trajectory translated by uniform random vector v [0, v] (u 0.1, 10,... .0..). (c)
Comparison of SI, SAI of the original trajectory, and SAI of the translated trajectory.
16

---

Published as a conference paper at ICLR 2022
C LEAPFROG INTEGRATION METHOD
The forward step of Leapfrog on separable Hamiltonian H(p. q) = T(p) + V(q) is given as fol- lows (Chen et 2019).
p 1 = =p"
1 - - -t) V
2
qltt = qr + -t++1-t))
- -1, V ++)).
na =p ti (lit1 1,) V. ((
D PROOF OF THEOREM
Theorem 1. Let A be an n x n symmetric real matrix and { be its n distinet eigenvalues.
Then for the linear system dæ = Ar with nonzero initial condition x00) ud, we have dt
u
(AR) 1/2 +0(A0),
(12)
|
At
1=1
where u' x(At) and & 2 0, 1,.... n with Ei= 2 =1
Proof. Let {u,}, be the eigenvectors corresponding to the eigenvalues { Since A is symmetric real matrix, { form set of orthogonal vectors. Without loss of generalization, we assume {v:}² is standard orthogonal basis. Let (0) = w Di=1 CrU;; then ([(),, Cx)lla ||22||2. The solution of the linear system is ax(1) Σ i=1 G.aee Therefore,
lim 1 u² Ar-0
u
1
Ei= crate Ei= G:ts
ll2
lim At
| At
|
At
cd,0%
(13)
1/2
===
We let l, A df
and have Σ l, 1.
17

---

Published as a conference paper at ICLR 2022
E SUPPLEMENTARY RESULTS
In this section we provide some supplementary results to the experiments of the main text.
Fig. 13 gives comparison of the predicted trajectories of the billiard model, and this is 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. 13 show the orbits of three particles learned by different methods, and these results are a supplement to Fig. 6 We can see that HNN and SRNN diverge from the orbits after 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 with time, and these results are a supplement Fig. These four examples in Fig. [6 support that the SAI keeps the same trend as SI and achieves the peak the same time SI does.
Reference
SANN (Ours)
SRNN
HNN
Time
Reference
SANN (Ours)
SRNN
HNN
Time
Reference
SANN (Ours)
SRNN
HNN
C
Time
Reference
SANN (Ours)
SRNN
HNN
Time
(Supplement to Fig. Comparison of billiard orbits simulated using the Hamiltonian learned by different methods.
Figure 13:
18

---

Published as a conference paper at ICLR 2022
0.
System
6.8
0.7
Reforence
SANN (Durs)
SRMN
HWN
HNN
SRNN
SANN
(Ours)
60 Time step
00
system energy
0.8
0.8
Reterence
(0und)
HNN
SRNN
SANN (Ours)
20
60
80
100
Time step
1.0
system eneroy
0.8
0.6
6.4
0.2
HNN
SRNN
SANN (Ours)
Time step
0.2
System energy
0.0
0.7
0.6
Relemance
SAMN (Dura)
HNN
SRNN
SANN (Ours)
60 Time step
0.7
system energy
0.8
0.6
Referencc SANN (0urs)
SRNN
HNN
SRNN
SANN (Ours)
100 20
60 60
Time step
Figure 14: (Supplement Fig. 6) The first three columns show comparison of reference orbits (solid curves) and the orbits learned by different method (dotted curves). is the initial position, and & is the direction.
19

---

Published as a conference paper at ICLR 2022
1.0
energy
0.0
6.8
0.8
Refemenc
SAMN (Suna)
0.7
3RM
HWN
HNN
SRNN
SANN (Ours)
60
step
Refemence
system energy
SANN (Dura)
4.0
2.0
0.0
HNN
SRNN
SANN (Ours)
20
40
60
00
100
Time step
system
0.9
6.
0.8
6.
HNN
SRNN
SANN (Ours)
Time step
Figure 15: (Supplement Fig. 6 The first three columns show the comparison of the reference orbits (solid curves) and the orbits learned by different methods (dotted curves). position, and is the direction.
is the initial
L-46
SN
SAI
12.5
1:37
1-37
15
t=48
1:46
-13
23-
-60
75
96
Three-body
stepiti
Step(()
EB9
c=63
t=40
1.0
-23
-23
di viz
20
Three-body
Time stepit)
Three body
Time stepit)
Figure 16: (Supplement to Fig. Comparison of SAI and SI with time. The color markers refers to the occurrence the stiffness phenomena, and the same color corresponds to the same time period. The trends of SAI and SI consistent.
20

---

Published as a conference paper at ICLR 2022
F NOISY DATA
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) the clean training data. Figures 17 and [8 respeectvvely 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 difficulty in learning the chaotic dynamics when the data is disturbed by noise.
Billiard noSse 0.0
Billiard (MSE)-noise 0.05
Billiard (MSE)-noise 0.1
Billiard (MSE)-noise 0.2
MSE
SAMM 19004
SAMN (Our))
c1-
SRAV
SPMM
SRMV
400
40
step
Time step
Time step
step
Figure 17: Comparison of MSE of the learned dynamical systems under noisy data with different noise levels.
Billiard (energyl-noise 0.0
Billiard (energy)-noise 0.05
Billiard (energyl-noise 0.1
Billiard (energy)-noise 0.2
Relative energy error 5
SANN (Durs)
SANN (Ours)
[Ours]
1OL1D Aummn Recnne PO
HNN
HNN
energy
SRNN
SANIN
2.5
Reletive
SANN [Ours!
HNN SRNN
200
.00
630
1000
2002
4000
M0X
3M0
1000
stop
Time step
stop
step
Figure 18: Comparison of relative energy error of the learned dynamical systems under noisy data with different noise levels.
21

---

Published as a conference paper at ICLR 2022
G HYPERPARAMETER SELECTION
We determine the stiff ratio based on the validation dataset. We investigate the distribution about the SAI over 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. I9 One can see that large SAIs mainly concentrate around the order index from 90 to 100. Hence, we choose 90- 90-percentile (y = 0.1) as the threshold.
10
5AI normallizel
0.0
0.6
90% nonstiff
04-
10%
0.
stiff
20
100
Index
Figure 19: The increasing order of the normalized SAIs over the three-body dataset of 1,000 trajec- tories. The shaded area represents the 95% confidence interval.
For the partition S for nonstiff and stiff portions, we perform the grid search for the pairs of pa- rameters as did in Choudhary et al. (202)). We use the validation loss as a proxy of the final performance, In particular, first 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.
22

---

Published as a conference paper at ICLR 2022
H INSEPARABLE HAMILTONIAN
SANN can also be applied to inseparable Hamiltonian. Like HNN, we use an NN that takes p and q an input to approximate the Hamiltonian directly on two problems: chain of N pendulums Finzi et al. 2020) and the three-body problem.
In the Pendulum-N problem, let my be the mass of the i-th pendulum, & be the rigid rod between (d-1)- (-th and i-th pendulum, be the generalized momentum and qi be the angle between the i-th pendulum and the y-axis. The Hamiltonian can be defined
H(p. q) = 2D ((a))) - ΣΣ
mily cos qk:
(14)
k=
where the mass matrix has complicated form M()), cos(q: 9 )glz EN fk= my- We consider N and N = 6, and set = m, = f = 1- 1,2, 1,2,... N. We simulate training set with 200 trajectories and the testing set with 100 trajectories by RKF45. The initial states are uniformly and randomly sampled from [-0.25, 0.25]. Each trajectory contains 100 time steps with step size of 0.03.
The bode problem is described in Section 4.2
It has a separable Hamiltonian, but we treat
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, set the partition parameter S to be 50 and 10 for the N-chain pendulum and the three-body problem, respectively. Table the train MSE and test MSE for SANN and HNN. We can see that SANN outperforms HNN in all the tests.
Dataset
Method
Train MSE
Test MSE
HNN
3. 3.80F-08
8.40F-05
Pendulum-4
SANN
1.80E-08
4.60E-05
HNN
1.60E-077
4.80E-04
Pendulum-6
SANN
1.00E-07
3.90E-04
HNN 4,20E-06
1.60E+00
Three-body
SANN
8.60E-07
7.00E-01
Table 2: Train MSE and test MSE for HNN and SANN.
23


Reviews:
[{"review_id": "t58voUXg6ys", "reviewer": "Reviewer_iTTU", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The paper is well-written, and it conveys the key ideas in a very precise manner. The authors have done an excellent job in highlighting the need for a stiffness-aware approach. Figure 1 and the related discussions in the Introduction section provide a compelling motivation for the problem.\n\nThe idea is straightforward but quite effective, as one can see from the experimental results. However, unless some aspects of the paper are further improved or better explained, it remains challenging to evaluate the real contribution and impact of this work.\n\nThe paper mentions that the mass matrix ($M$) is diagonal and subsequently it uses a set of trainable parameters to learn the individual elements of this diagonal matrix. Would you please confirm if this is indeed the case or the writing is conveying a wrong picture? For a large class of Hamiltonian systems, the mass/inertia matrix is not diagonal, and its entries depend on the generalized coordinates; for example, please consider a $k$-link pendulum with joint angles as the generalized coordinates. Therefore, if $M$ is indeed assumed to be a diagonal matrix, the scope of this work is very restricted unless the authors can show with additional experiments that the proposed approach holds true even when the mass matrix is non-diagonal and position-dependent.\n\nBasri et al. (*The Convergence Rate of Neural Networks for Learned Functions of Different Frequencies*, NeurIPS 2019) have demonstrated that the number of epochs to learn a function has a squared relationship with its frequency. As the stiffness of an ODE is related to the presence of very high-frequency components in its solution, it may be possible that SRNN or HNN can achieve comparable accuracy if they are trained over a sufficiently long period. But the current set of experiments do not provide a precise answer to this equation. Therefore, I would strongly encourage the authors to run an additional experiment that trains SRNN and HNN with a very high number of epochs.\n\nFinzi et al. (*Simplifying Hamiltonian and Lagrangian Neural Networks via Explicit Constraints*, NeurIPS 2020) have shown that the $L_1$ loss functions exhibit better performance while inferring dynamics from data. It would be helpful to see if the same holds true for this problem as well. Therefore, the authors should consider carrying out an additional ablation study that shows how the performance varies when the loss function uses the $L_1$ norm.\n\nRecent work by Kim et al. (*Stiff Neural Ordinary Differential Equations*, arXiv:2103.15341) has extended the scope of Neural ODEs to stiff differential equations. I would encourage the authors to consider this approach (if possible) as an additional baseline.\n\nThe Related Work section has missed some relevant prior work that enforces Hamiltonian dynamics while using a neural network to infer dynamics from data. Please refer to the following survey papers and the references therein for further details about the relevant prior work: *Integrating Physics-Based Modeling with Machine Learning: A Survey* (Willard et al., arXiv:2003.04919), *An overview on recent machine learning techniques for Port Hamiltonian systems* (Cherifi, Physica D, 2020), and *Benchmarking Energy-Conserving Neural Networks for Learning Dynamics from Data* (Zhong et al., L4DC 2021).\n\nSome additional minor comments:\n\n--- The second paragraph of the Introduction mentions: \"the particles deviate from the reference orbits rapidly after collision\". As two particles cannot collide in this setting, the authors should rephrase the sentence with \"after a close encounter\".\n\n--- The authors should consider introducing $gamma$, i.e., the hyperparameter that denotes the stiffness ratio threshold, as a percentile number. It would make the point clearer.\n\n--- Also, \"the larger(stiffer)\" should be either \"a larger(stiffer)\" or \"the largest(stiffest)\".", "summary_of_the_paper": "Incorporating Hamiltonian dynamics based inductive bias into deep neural networks has gained significant attention in the machine learning community over the last few years. However, in many real physical systems, the underlying ODEs can exhibit stiffness, i.e., numerical instability, over some time intervals or for certain values of initial conditions and/or parameter choice. As a result, the existing realizations of Hamiltonian-preserving neural networks underperform in these scenarios. This paper proposes a solution to overcome this issue. In particular, it introduces an easy to compute stiffness-aware index to split the training data points into stiff and nonstiff groups, which are then treated with different integration schemes (with different values for integration time interval and time steps). When demonstrated on relevant physical systems, the proposed approach shows improved accuracy in prediction and energy conservation.", "summary_of_the_review": "This paper has proposed a straightforward but effective solution that can infer Hamiltonian dynamics even when the governing ODEs exhibit numerical stiffness. The problem is very well-motivated, and the paper explains the core ideas in a very precise way. However, I have mentioned in the *main review*, some aspects (especially the assumption that $M$ is a diagonal matrix with entries that are position-independent) should be properly addressed or explained to increase the concreteness and overall clarity of the paper.\n\n********** Post Rebuttal Response **********\nI would like to thank the authors for addressing the prior concerns. I have updated my score.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "S_WN9ytlrdF", "reviewer": "Reviewer_MCLP", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strong points.\n- A stiffness-aware approach for learning Hamiltonian systems is novel.\n- The experiments show that SANN can accurately simulate complex Hamiltonian systems than baseline methods, HNN and SRNN.\n- This paper is well-written.\n\nWeak points.\n- There are some hyper-parameters to be determined manually.\n- The compared methods are the bare minimum.\n- There are several concerns about the experimental setting.\n\nComments.\n1. The task of learning physical dynamics from data has been of great interest recently. A key idea of incorporating stiffness into the learning scheme is novel and exciting. The proposed method is simple but effective; however, there are the hyper-parameters $\\gamma, S$ to be manually determined. Especially, the ratio of the stiff portion, $\\gamma$, could be critical for performance. Could the authors explain how to estimate this hyper-parameter for various physical systems, including the $M$-body problem when $M$ is large?\n\n2. The proposed method classifies the intervals into binary labels, that is, stiff or nonstiff. It might be helpful to model the continuous stiff-ratio for each interval that is learnable.\n\n3. The authors assume the separable Hamiltonian $H(p,q)=T(p)+V(q)$. Can the proposed learning framework apply to the inseparable generalized model?\n\n4. The HNN (Greydanus et al., 2019) does not consider the separable assumption. In the experiments, did the authors use this assumption in the HNN learning? Also, The input of HNN is the partial derivatives, unlike SANN and HRNN. How did the authors give the input to HNN?\n\n5. Why did the authors use the Leapfrog solver? Isn't the simple Euler method appropriate?", "summary_of_the_paper": "This paper presents a new method, stiffness-aware neural network (SANN), for learning Hamiltonian systems from data. The authors define a stiffness-aware index (SAI) for classifying the training data into stiff and nonstiff samples. Based on the classification result, the step size for the numerical solver is adjusted, and the number of samples for training is balanced. The effectiveness of SANN is demonstrated using chaotic Hamiltonian systems, i.e., a three-body problem and billiard model.", "summary_of_the_review": "This paper addresses the interesting problem and is well-motivated. Also, the proposed approach is novel, and the experimental results are insightful. Although I tend to accept this paper, I have some concerns, which I detail in the main review.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "4wfL3lM1-5", "reviewer": "Reviewer_MYXe", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "### Positives\n\nIt is an insightful suggestion that the stiffness is a bottleneck to learning of a physical system.\n\nIt is surprising that a simple oversampling is enough.\n\n### Negatives\n\nSection 5 demonstrates that SAI is a good approximation to stiffness index (SI), but this might not hold for different coordinate systems. SI assumes that the origin of the coordinate system is an equilibrium point. In other words, the bias from the equilibrium point is already subtracted from the position. However, SAI does not (or cannot in practice). The norm of the state, and thereby, SAI depend on the coordinate system. It is preferable to valid the generality of SAI.\n\nThe sampling strategy is heuristic. As shown in Table 1, the results are sensitive to the hyperparameter tuning. A guidance justified theoretically is preferable.\n\n### Minor Comments\n\nIt might be an insightful suggestion that the neural network tends to learn a smooth function and this implicit bias is a bottleneck to learning of a stiff system. However, this suggestion is not validated by experiments, and it is unclear whether the proposal resolves this problem. An additional experiment or analysis is not mandatory, but may improve the contribution of the present study.\n\nFigure 1 is a bit confusing. I prefer that the scales of the axes are fixed within each example.\n\n### After discussion\n\nAll my concerns were addressed by the additional experiments and explanations. I update the score from 5 to 6.", "summary_of_the_paper": "This study proposes the stiff-aware index (SAI) for an ordinary differential equation and the training strategy that uses samples with large SAIs more frequently. A neural network has the implicit bias to tend to learn a smooth function. Only a limited portion of data obtained from a stiff system exhibits a rapid change, in other words, the training data is imbalanced between a gradual change and a rapid change. Hence, without SAI, it is difficult for a neural network to learn a stiff dynamics. The contribution of SAI is confirmed using a three-body problem.\n", "summary_of_the_review": "This study is based on an insightful suggestion, and the proposed method is simple but effective. However, the strategy is heuristic, and the generality is unclear.\n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "0Z7ld3rAEiA", "reviewer": "Reviewer_RqeB", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Separating the stiff and non-stiff parts of the data when training a Hamiltonian network is a novel idea. And the results show certain advantages of using this method. And experiments are included to discuss the influence of resampling, integration partitions, and activation functions.\nHowever, only two examples are shown in the experiments, it would be more convincing if experiments of more Hamiltonian systems can be conducted.\n\nSome additional questions:\nWhen predicting future states, does the model use s fixed time step? (or still using a smaller timestep for stiff parts? If so, how is the stiffness calculated?)\nAre there significant advantages of using an additional trainable term p^TMp, instead of training a single network that represents p^TMp + φ(q;W)? \n\n", "summary_of_the_paper": "This paper proposes to improve the learning of a Hamiltonian system, by characterizing the stiffness of the time series data.\n\nA stiffness-aware index is first used to classify the time interval into stiff and nonstiff. Then, during the training of the Hamiltonian network, the stiff part is integrated using a smaller timestep, and also sampled more frequently in the training data.", "summary_of_the_review": "There are certain contributions of this paper on proposing characterizing the stiffness of the time series data, which shows its advantage of improving the performance of the Hamiltonian network.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper introduces the Stiffness-aware neural network (SANN) for improving numerical stability in Hamiltonian neural networks. To this end, the authors introduce the stiffness-aware index (SAI) to classify time intervals into stiff and non-stiff portions, and propose to adapt the integration scheme accordingly.

The paper initially received three weak accept and one weak reject recommendations. The main limitations pointed out by reviewers relate to missing references from the literature, assumptions behind the proposed approach (e.g. structure of the mass matrix, separable Hamiltonian), and clarifications on experiments including additional baselines and hyper-parameter settings. 
The rebuttal did a good job in answering reviewers' concerns: RiTTU increased his rating to a clear accept, and RMYXe increased his rating to weak accept.
Eventually, there is a consensus among reviewers to accept the paper. 

The AC's own readings confirmed the reviewers' recommendation. The method is straightforward yet effective, and the paper is well written. The effectiveness of the proposed approach is shown in different contexts. Since several complex systems exhibit chaotic characteristics, the paper brings a meaningful contribution to the community.
Decision: Accept (Poster)