Venue: ICLR 2022 Conference

Reviews:

final_score: rating: 8: accept, good paper
review_content: 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.
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.

The 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.

The 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.

Basri 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.

Finzi 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.

Recent 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.

The 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).

Some additional minor comments:

--- 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".

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

--- Also, "the larger(stiffer)" should be either "a larger(stiffer)" or "the largest(stiffest)".
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.

********** Post Rebuttal Response **********
I would like to thank the authors for addressing the prior concerns. I have updated my score.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: 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.
main_review: Strong points.
- A stiffness-aware approach for learning Hamiltonian systems is novel.
- The experiments show that SANN can accurately simulate complex Hamiltonian systems than baseline methods, HNN and SRNN.
- This paper is well-written.

Weak points.
- There are some hyper-parameters to be determined manually.
- The compared methods are the bare minimum.
- There are several concerns about the experimental setting.

Comments.
1. 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?

2. 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.

3. The authors assume the separable Hamiltonian $H(p,q)=T(p)+V(q)$. Can the proposed learning framework apply to the inseparable generalized model?

4. 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?

5. Why did the authors use the Leapfrog solver? Isn't the simple Euler method appropriate?
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.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: This paper proposes to improve the learning of a Hamiltonian system, by characterizing the stiffness of the time series data.

A 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.
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.
However, only two examples are shown in the experiments, it would be more convincing if experiments of more Hamiltonian systems can be conducted.

Some additional questions:
When 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?)
Are there significant advantages of using an additional trainable term p^TMp, instead of training a single network that represents p^TMp + φ(q;W)? 


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.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: 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.

main_review: ### Positives

It is an insightful suggestion that the stiffness is a bottleneck to learning of a physical system.

It is surprising that a simple oversampling is enough.

### Negatives

Section 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.

The sampling strategy is heuristic. As shown in Table 1, the results are sensitive to the hyperparameter tuning. A guidance justified theoretically is preferable.

### Minor Comments

It 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.

Figure 1 is a bit confusing. I prefer that the scales of the axes are fixed within each example.

### After discussion

All my concerns were addressed by the additional experiments and explanations. I update the score from 5 to 6.
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.
review_title: null

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)