Title: Physics-Informed Regularization for Domain-Agnostic Dynamical System Modeling
Venue: NeurIPS 2024
Accepted: yes
Keywords: Physics-informed Neural Networks;  Dynamical Systems ; Graph Neural Networks
Abstract: Learning complex physical dynamics purely from data is challenging due to the intrinsic properties of systems to be satisfied. Incorporating physics-informed priors, such as in Hamiltonian Neural Networks (HNNs), achieves high-precision modeling for energy-conservative systems. However, real-world systems often deviate from strict energy conservation and follow different physical priors. To address this, we present a framework that achieves high-precision modeling for a wide range of dynamical systems from the numerical aspect, by enforcing Time-Reversal Symmetry (TRS) via a novel regularization term. It helps preserve energies for conservative systems while serving as a strong inductive bias for non-conservative, reversible systems. While TRS is a domain-specific physical prior, we present the first theoretical proof that TRS loss can universally improve modeling accuracy by minimizing higher-order Taylor terms in ODE integration, which is numerically beneficial to various systems regardless of their properties, even for irreversible systems. By integrating the TRS loss within neural ordinary differential equation models, the proposed model TREAT demonstrates superior performance on diverse physical systems. It achieves a significant 11.5% MSE improvement in a challenging chaotic triple-pendulum scenario, underscoring TREAT’s broad applicability and effectiveness.

Physics-Informed Regularization for Domain-Agnostic Dynamical System Modeling

Abstract

Learning complex physical dynamics purely from data is challenging due to the intrinsic properties of systems to be satisfied. Incorporating physics-informed priors, such as in Hamiltonian Neural Networks (HNNs), achieves high-precision modeling for energy-conservative systems. However, real-world systems often deviate from strict energy conservation and follow different physical priors. To address this, we present a framework that achieves high-precision modeling for a wide range of dynamical systems from the numerical aspect, by enforcing Time-Reversal Symmetry (TRS) via a novel regularization term. It helps preserve energies for conservative systems while serving as a strong inductive bias for non-conservative, reversible systems. While TRS is a domain-specific physical prior, we present the first theoretical proof that TRS loss can universally improve modeling accuracy by minimizing higher-order Taylor terms in ODE integration, which is numerically beneficial to various systems regardless of their properties, even for irreversible systems. By integrating the TRS loss within neural ordinary differential equation models, the proposed model TREAT demonstrates superior performance on diverse physical systems. It achieves a significant 11.5% MSE improvement in a challenging chaotic triple-pendulum scenario, underscoring TREAT's broad applicability and effectiveness. Code and further details are available at here.

1 Introduction

Dynamical systems, spanning applications from physical simulations (Kipf et al., 2018; Wang et al., 2020; Lu et al., 2022; Huang et al., 2023; Luo et al., 2023a; Xu et al., 2024; Luo et al., 2024) to robotic control (Li et al., 2022; Ni and Qureshi, 2022), are challenging to model due to intricate dynamic patterns and potential interactions under multi-agent settings. Traditional numerical simulators require extensive domain knowledge for design, which is sometimes unknown (Sanchez-Gonzalez et al., 2020), and can consume significant computational resources (Wang et al., 2024). Therefore, directly learning dynamics from the observational data becomes an attractive alternative.

Existing deep learning approaches (Sanchez-Gonzalez et al., 2020; Pfaff et al., 2021; Han et al., 2022a) usually learn a fixed-step transition function to predict system dynamics from timestamp t to timestamp t +1 and rollout trajectories recursively. The transition function can have different inductive biases, such as Graph Neural Networks (GNNs) (Pfaff et al., 2020; Martinkus et al., 2021; Lam et al., 2023; Cao et al., 2023) for capturing pair-wise interactions among agents through message passing. Most recently, neural ordinary differential equations (Neural ODEs) (Chen et al.,

1

Time-Reversal

Dynamical Systems under

Energy-Conservative

(a) High-Precision Modeling of Dynamical Systems

𝑑𝑡

= -𝐹(𝑅 𝑧 )

𝑡

backward

forward

𝑧 =

𝑞, 𝑝

𝑅 𝑧

= (𝑞, -𝑝)

Identical

positions

(

𝑞

)

(b.2) Time-Reversal Symmetry

TRS Loss

Injecting

Physical Prior

Reducing Error

Accumulation over

Integration Steps

(b.1) Physical Priors

Integration steps

Euler's method

Grorund truth

𝑧

!"#

= 𝑧

!

+

𝑑𝐹

𝑑𝑧

!

∆𝑡 + 𝑂 ∆𝑡

$

+..

Numerical Errors

(b.3) Error Accumulation

2018; Rubanova et al., 2019) have emerged as a potent solution for modeling system dynamics in a continuous manner, which offer superior prediction accuracy over discrete models in the long-range, and can handle systems with partial observations. In particular, GraphODEs (Huang et al., 2020; Luo et al., 2023b; Zang and Wang, 2020; Jiang et al., 2023; Luo et al., 2023c) extend NeuralODEs to model interacting (multi-agent) dynamical systems, where agents co-evolve and form trajectories jointly.

However, the complexity of dynamical systems necessitates large amounts of data. Models trained on limited data risk violating fundamental physical principles such as energy conservation. A promising strategy to improve modeling accuracy involves incorporating physical inductive biases (Raissi et al., 2019; Cranmer et al., 2020). Existing models like Hamiltonian Neural Networks (HNNs) (Greydanus et al., 2019; Sanchez-Gonzalez et al., 2019) strictly enforce energy conservation, yielding more accurate predictions for energy-conservative systems. However, not all real-world systems strictly adhere to energy conservation, and they may adhere to various physical priors. Other methods that model both energy-conserving and dissipative systems, as well as reversible systems, offer more flexibility (Zhong et al., 2020; Gruber et al., 2024). Nevertheless, they often rely on prior knowledge of the system and are also limited to systems with corresponding physical priors. Such system diversity largely limits the usage of existing models which are designed for specific physical prior.

To address this, we present a framework that achieves high-precision modeling for a wide range of dynamical systems from the numerical aspect, by enforcing Time-Reversal Symmetry (TRS) via a novel regularization term. Specifically, TRS posits that a system's dynamics should remain invariant when time is reversed (Lamb and Roberts, 1998). To incorporate TRS, we propose a simple-yet-effective self-supervised regularization term that acts as a soft constraint. This term aligns forward and backward trajectories predicted by a neural network and we use GraphODE as the backbone. We theoretically prove that the TRS loss effectively minimizes higher-order Taylor expansion terms during ODE integration, offering a general numerical advantage for improving modeling accuracy across a wide array of systems, regardless of their physical properties. It forces the model to capture fine-grained physical properties such as jerk (the derivatives of accelerations) and provides more regularization for long-term prediction. We also justify our TRS design choice, showing case its superior performance both analytically and empirically. We name the model as TREAT ( T imeRe vers a l Symme t ry ODE).

Note that TRS itself is a physical prior, that is broader than energy conservation as depicted in Figure 1(b.1). It covers classical energy-conservative systems such as Newtonian mechanics, and also non-conservative, reversible systems like Stokes flow (Pozrikidis, 2001), commonly encountered in microfluidics (Kim and Karrila, 2013; Cao and Li, 2018; Cao et al., 2019). Therefore, TRS loss achieves high-precision modeling from both the physical aspect, and the numerical aspect as shown in Figure 1(a), making it domain-agnostic and widely applicable to various dynamical systems. We systematically conduct experiments across 9 diverse datasets spanning across 1.) single-agent, multi-agent systems; 2.) simulated and real-world systems; and 3.) systems with different physical

priors. TREAT consistently outperforms state-of-the-art baselines, affirming its effectiveness and versatility across various dynamic scenarios.

Our primary contributions can be summarized as follows:

We introduce TREAT, a powerful framework that achieves high-precision modeling for a wide range of systems from the numerical aspect, by enforcing Time-Reversal Symmetry (TRS) via a regularization term.

We establish the first theoretical proof that the time-reversal symmetry loss could in general help learn more fine-grained and long-context system dynamics from the numerical aspect, regardless of systems' physical properties (even irreversible systems). This bridges the specific physical implication and the general numerical benefits of the physical prior -TRS.

We present empirical evidence of TREAT's state-of-the-art performance in a variety of systems over 9 datasets, including real-world & simulated systems, etc. It yields a significant MSE improvement of 11.5% on the challenging chaotic triple-pendulum system.

2 Preliminaries and Related Work

We represent a dynamical system as a graph G = ( V , E ) , where V denotes the node set of N agents 3 and E denotes the set of edges representing their physical interactions. For simplicity, we assumed G to be static over time. Single-agent dynamical system is a special case where the graph only has one node. In the following, we use the multi-agent setting by default to illustrate our model. We denote X ( t ) ∈ R N × d as the feature matrix at timestamp t for all agents, with d as the feature dimension. Model input consists of trajectories of feature matrices over M historical timestamps X ( t -M : -1 ) = { X ( t -M ) , . . . , X ( t -1 ) } and G . The timestamps t -1 , · · · , t -M < 0 can have non-uniform intervals and take any continuous values. Our goal is to learn a neural simulator f θ ( · ) : [ X ( t -M : -1 ) , G ] → Y ( t 0: K ) , which predicts node dynamics Y ( t ) in the future on timestamps 0 = t 0 < · · · < t K = T sampled within [0 , T ] . We use y i ( t ) to denote the targeted dynamic vector of agent i at time t . In some cases when we are only predicting system feature trajectories, Y ( · ) ≡ X ( · ) .

2.1 NeuralODE for Dynamical Systems

NeuralODEs (Chen et al., 2018; Rubanova et al., 2019) are a family of continuous models that define the evolution of dynamical systems by ordinary differential equations (ODEs). The state evolution can be described as: ˙ z i ( t ) := d z i ( t ) dt = g ( z 1 ( t ) , z 2 ( t ) · · · z N ( t )) , where z i ( t ) ∈ R d denotes the latent state variable for agent i at timestamp t . The ODE function g is parameterized by a neural network such as Multi-Layer Perception (MLP), which is automatically learned from data. GraphODEs (Poli et al., 2019; Huang et al., 2020; Luo et al., 2023b; Wen et al., 2022; Huang et al., 2024) are special cases of NeuralODEs, where g is a Graph Neural Network (GNN) to capture the continuous interaction among agents.

GraphODEs have been shown to achieve superior performance, especially in long-range predictions and can handle data irregularity issues. They usually follow the encoder-processor-decoder architecture, where an encoder first computes the latent initial states z 1 ( t 0 ) , · · · z N ( t 0 ) for all agents simultaneously based on their historical observations as in Eqn 1.

Then the GNN-based ODE predicts the latent trajectories starting from the learned initial states. The latent state z i ( t ) can be computed at any desired time using a numerical solver such as RungeKuttais (Schober et al., 2019) as:

Finally, a decoder extracts the predicted dynamics ˆ y i ( t ) based on the latent states z i ( t ) for any timestamp t :

However, vanilla GraphODEs can violate physical properties of a system, resulting in unrealistic predictions. We therefore propose to inject physics-informed regularization term to make more accurate predictions.

2.2 Time-Reversal Symmetry (TRS)

Consider a dynamical system described in the form of d x ( t ) dt = F ( x ( t )) , where x ( t ) ∈ Ω is the observed states such as positions. The system is said to follow the TimeReversal Symmetry if there exists a reversing operator R : Ω ↦→ Ω such that (Lamb and Roberts, 1998):

where ◦ denote the action of functional R on the function x .

Intuitively, we can assume x ( t ) is the position of a flying ball and the con- ventional reversing operator is defined as R : x ↦→ R ◦ x , R ◦ x ( t ) = x ( -t ) . This implies when x ( t ) is a forward trajectory position with initial position x (0) , x ( -t ) is then a position in the timereversal trajectory, where x ( -t ) is calculated using the same function F , but with the integration time reversed, i.e. dt ↦→ d ( -t ) . Eqn 4 shows how to create the reverse trajectory of a flying ball: at each position, the velocity (i.e., the derivative of position with respect to time) should be the opposite. In neural networks, we usually model trajectories in the latent space via z (Sanchez-Gonzalez et al., 2020), which can be decoded back to real observation state i.e. positions. Therefore, we apply the reversal operator for z .

Now we introduce a time evolution operator ϕ τ such that ϕ τ ◦ z ( t ) = z ( t + τ ) for arbitrary t, τ ∈ R . It satisfies ϕ τ 1 ◦ ϕ τ 2 = ϕ τ 1 + τ 2 , where ◦ denotes composition. The time evolution operator helps us to move forward (when τ > 0 ) or backward (when τ < 0 ) through time, thus forming a trajectory. Based on (Lamb and Roberts, 1998), in terms of the evolution operator, Eqn 4 implies:

which means that moving forward t steps and then turning to the opposite direction is equivalent to firstly turning to the opposite direction and then moving backwards t steps 4 . Eqn 5 has been widely used to describe time-reversal symmetry in existing literature (Huh et al., 2020; Valperga et al., 2022). Nevertheless, we propose the following lemma, which is more intuitive to understand and straightforward to guide the design of our time-reversal regularizer.

Lemma 2.1. Eqn 5 is equivalent to R ◦ ϕ t ◦ R ◦ ϕ t = I , where I denotes identity mapping.

Lemma 2.1 means if we move t steps forward, then turn to the opposite direction, and then move forward for t more steps, it shall restore back to the same state. This is illustrated in Figure 2 where the reverse trajectory should be the same as the forward trajectory. 5 It can be understood as rewinding a video to the very beginning. The proof of Lemma 2.1 is in Appendix A.2.

3 Method: TREAT

We present a novel framework TREAT that achieves high-precision modeling for a wide range of systems from the numerical aspect, by enforcing Time-Reversal Symmetry (TRS) via a regularization

𝐳

𝐢

𝐟𝐰𝐝

(𝐭

𝐊

)

𝐳

𝐢

𝐟𝐰𝐝

(𝐭

𝟎

)

𝐳

𝐢

𝐟𝐰𝐝

(𝐭

𝟏

)

𝐳

𝐢

𝐟𝐰𝐝

(𝐭

𝐊(𝟏

)

𝒛

𝒊

𝐫𝐞𝐯

(𝒕

𝑲

.

)

……

Latent forward trajectory

Latent reverse trajectory

𝝓

𝑻

𝑅

𝑅

𝝓

𝑻

𝑹∘ 𝝓

𝑻

∘ 𝑹 ∘ 𝝓

𝑻

= 𝑰

𝒛

𝒊

𝐫𝐞𝐯

(𝒕

𝑲(𝟏

.

)

𝒛

𝒊

𝐫𝐞𝐯

(𝒕

𝟏

.

)

𝒛

𝒊

𝐫𝐞𝐯

(𝒕

𝟎

.

)

……

𝑂𝐷𝐸 𝑆𝑜𝑙𝑣𝑒

𝑔, (𝑧

!

(𝑡

"

) … 𝑧

#

(𝑡

"

) , (𝑡

"

…𝑡

$

))

Time-Reversal Symmetry Constraint

𝑡

%&

𝑡

%'

𝑂

!

𝑂

(

𝑂

'

𝑡

%(

𝑡

%!

Processor: Learnable ODE function

Forward trajectory

) 𝒚

𝒊

𝐟𝐰𝐝

Reverse trajectory

) 𝒚

𝒊

𝐫𝐞𝐯

Encoder

Input

𝑿

,

𝓖

Output

𝒀

𝑂

!

𝑂

"

𝑂

#

𝑡

"

𝑡

$%!

𝑡

$

𝑂

!

𝑂

'

𝑡

!

𝑂

(

𝑧

"

(0)

𝑧

#

(0)

𝑧

!

(0)

+ 𝑦

(

)*+

𝑡

= 𝑓

,-.

(𝑧

(

)*+

(𝑡))

Decoder

! 𝒚

𝒊

𝐟𝐰𝐝

(𝒕

𝑲

)

! 𝒚

𝒊

𝐫𝐞𝐯

(𝒕

𝟎

*

)

𝒛

𝒊

𝐟𝐰𝐝

(𝒕

𝟎

)

𝒛

𝒊

𝐟𝐰𝐝

(𝒕

𝑲+𝟏

) 𝒛

𝒊

𝐟𝐰𝐝

(𝒕

𝑲

)

𝒛

𝒊

𝐟𝒘𝒅

(𝒕

𝟏

)

𝑅

𝝓

𝑻

𝝓

𝑻

! 𝒚

𝒊

𝐟𝐰𝐝

(𝒕

𝟎

)

! 𝒚

𝒊

𝐫𝐞𝐯

(𝒕

𝑲

*

)

= 4

𝐢1𝟏

𝐍

4

𝒌1𝟎

𝐊

5

5) 𝒚

𝒊

𝐟𝐰𝐝

𝑡

7

-) 𝒚

𝒊

𝐫𝐞𝐯

(𝑡

897

:

)

𝟐

𝟐

𝓛

𝐫𝐞𝐯𝐞𝐫𝐬𝐞

𝑹

! 𝒚

𝒊

𝐟𝐰𝐝

(𝒕

𝟏

)

! 𝒚

𝒊

𝐫𝐞𝐯

(𝐭

𝐊+𝟏

*

)

Latent dynamics

! 𝒚

𝒊

𝐟𝐰𝐝

(𝒕

𝑲+𝟏

)

! 𝒚

𝒊

𝐫𝐞𝐯

(𝐭

𝟏

*

)

1 𝑦

!

"#$

𝑡

= 𝑓

%&'

(𝑧

!

"#$

(𝑡))

𝑧

(

)*+

(𝑡

=

) = 𝑓

->.

(𝑋 𝑡

9?:9A

, 𝒢)

term. It improves modeling accuracy regardless of systems' physical properties. We first introduce our architecture design, followed by theoretical analysis to explain its numerical benefits.

TREAT uses GraphODE (Huang et al., 2020) as the backbone and flexibly incorporates TRS as a regularization term based on Lemma 2.1. This term aligns model forward and reverse trajectories. In practice, our model predicts the forward trajectories at a series of timestamps { t k } K k =0 as ground truth observations are discrete, where 0 = t 0 < t 1 < · · · < t K = T . The reverse trajectories are also at the same series of K timestamps so as to be aligned with the forward one, which we denote as { t ′ k } K k =0 satisfying 0 = t ′ 0 < t ′ 1 < · · · < t ′ K = T . It's important to note that the values of the time variable t ′ k in the reverse trajectories do not represent real time, but serve as indexes of reverse trajectories. This leads to the relation t ′ K -k = T -t k , which means the reverse trajectories at timestamp t ′ K -k correspond to the forward trajectories at time t k . For example, t ′ 0 = T -t K = 0 . It indicates t ′ 0 and t K are both pointing to the same real time T , which is the ending point of the forward trajectory as shown in Figure 3. Based on Lemma 2.1, the difference of the two trajectories at any observed time should be small, i.e. z fwd ( t k ) ≈ z rev ( t ′ K -k ) . This serves as the guideline for our regularizer design. The weight of the regularizer is also adjustable to adapt different systems. The overall framework is depicted in Figure 3.

3.1 Time-Reversal Symmetry Loss and Training

Forward Trajectory Prediction and Reconstruction Loss. For multi-agent systems, we utilize the GNN operator described in (Kipf et al., 2018) as our ODE function g ( · ) , which drives the system to move forward and output the forward trajectories for latent states z fwd i ( t ) at each continuous time t ∈ [0 , T ] and each agent i .We then employ a Multilayer Perceptron (MLP) as a decoder to predict output trajectories ˆ y fwd i ( t ) based on the latent states. We summarize the whole procedure as:

To train the model, we use the reconstruction loss that minimizes the L2 distance between predicted forward trajectories { ˆ y fwd i ( t k ) } K k =0 and the ground truth trajectories { y i ( t k ) } K k =0 as :

Reverse Trajectory Prediction and Regularization Loss. We design a novel time-reversal symmetry loss as a soft constraint to flexibly regulate systems' behavior based on Lemma 2.1. Specifically, we first compute the latent reverse trajectories z rev ( t ) by starting from the ending state of the forward one, traversed back over time. We then employ the decoder to output dynamic trajectories y rev ( t ) .

Next, based on Lemma 2.1, if the system follows Time-Reversal Symmetry , the forward and backward trajectories shall be exactly overlap. We thus design the reversal loss by minimizing the L2 distances between model forward and backward trajectories decoded from the latent trajectories:

Finally, we jointly train TREAT as a weighted combination of the two losses:

where α is a positive coefficient to balance the two losses based on different targeted systems.

Remark. The computational time of L reverse is of the same scale as the reconstruction loss L pred . As the computation process of the reversal loss is to first use the ODE solver to generate the reverse trajectories, which has the same computational overhead as computing the forward trajectories, and then compute the L2 distances.

3.2 Theoretical Analysis of Time-Reversal Symmetry Loss

We next theoretically show that the time-reversal symmetry loss numerically helps to improve prediction accuracy in general, regardless of systems' physical properties. Specifically, we show that it minimizes higher-order Taylor expansion terms during the ODE integration steps.

Theorem 3.1. Let ∆ t denote the integration step size in an ODE solver and T be the prediction length. The reconstruction loss L pred defined in Eqn 7 is O ( T 3 ∆ t 2 ) . The time-reversal loss L reverse defined in Eqn 9 is O ( T 5 ∆ t 4 ) .

We prove Theorem 3.1 in Appendix A.3. From Theorem 3.1, we can see two nice properties of our proposed time-reversal loss: 1) Regarding the relationship to ∆ t , L reverse is optimizing a highorder term ∆ t 4 , which forces the model to predict fine-grained physical properties such as jerk (the derivatives of accelerations). In comparison, the reconstruction loss optimizes ∆ t 2 , which mainly guides the model to predict the locations/velocities accurately. Therefore, the combined loss enables our model to be more noise-tolerable; 2) Regarding the relationship to T , L reverse is more sensitive to total sequence length ( T 5 ), thus it provides more regularization for long-context prediction, a key challenge for dynamic modeling.

TRS Loss Design Choice. We define L reverse as the distance between model forward trajectories and backward trajectories. Based on the definition of TRS in Sec. 2.2, there are other implementation choices. One prior work TRS-ODE (Huh et al., 2020) designed a TRS loss based on Eqn 5, where a reverse trajectory shares the same starting point as the forward one. However, we show that our implementation based on Lemma 2.1 to approximate time-reversal symmetry has a lower maximum error compared to their implementation below, supported by empirical experiments in Sec. 4.2.

Lemma 3.2. Let L reverse be the TRS implementation of TREAT based on Lemma 2.1, L reverse 2 be the one in (Huh et al., 2020) based on Eqn 5. When the reconstruction loss defined in Eqn 7 of both methods are equal, and the two TRS losses are equal, i.e. L reverse = L reverse 2 , the maximum error between the reversal and ground truth trajectory for each agent, i.e. MaxError gt _ rev = max k ∈ [ K ] ∥ y i ( t k ) -ˆ y rev i ( t ′ K -k ) ∥ 2 for i = 1 , 2 · · · N , made by TREAT is smaller.

We prove Lemma 3.2 in Appendix A.4. Another implementation is to minimize the distances between model backward trajectories and ground truth trajectories. When both forward and backward trajectories are close to ground-truth, they are implicitly symmetric. The major drawback is that at the early stage of learning when the forward is far away from ground truth ( L pred ), such implicit regularization does not force time-reversal symmetry, but introduces more noise.

4 Experiments

Datasets. We conduct systematic evaluations over five multi-agent systems including three 5-body spring systems (Kipf et al., 2018), a complex chaotic pendulum system and a real-world motion capture dataset (Carnegie Mellon University, 2003); and four single-agent systems including three spring systems (with only one node) and a chaotic strange attractors system (Huh et al., 2020).

The settings of spring systems include: 1) conservative, i.e. no interactions with the environments, we call it Simple Spring ; 2) non-conservative with frictions, we call it Damped Spring ; 3) nonconservative with periodic external forces, we call it Forced Spring . The Pendulum system contains three connected sticks in a 2D plane. It is highly sensitive to initial states, with minor disturbances leading to significantly different trajectories (Shinbrot et al., 1992; Awrejcewicz et al., 2008). The realworld motion capture dataset (Carnegie Mellon University, 2003) describes the walking trajectories of a person, each tracking a single joint. We call it Human Motion . The strange attractor consists of symmetric attractor/repellor force pairs and is chaotic (Sprott, 2015). It is also highly sensitive to the initial states (Koppe et al., 2019). We call it Attractor .

Towards physical properties, Simple Spring and Pendulum are conservative and reversible; Force Spring and Attractor are reversible but non-conservative; Damped Spring are irreversible and nonconservative. For Human Motion , it does not adhere to specific physical laws since it is a real-world dataset. Details of the datasets and generation pipelines can be found inAppendix C.

Task Setup. We conduct evaluation by splitting trajectories into two halves: [ t 1 , t M ] , [ t M +1 , t K ] where timestamps can be irregular. We condition the first half of observations to make predictions for the second half as in (Rubanova et al., 2019). For spring datasets and Pendulum , we generate irregular-sampled trajectories and set the training samples to be 20,000 and testing samples to be 5,000 respectively. For Attractor , We generate 1,000 and 50 trajectories for training and testing respectively following Huh et al. (2020). 10% of training samples are used as validation sets and the maximum trajectory prediction length is 60. Details can be found in Appendix C.

Baselines. Wecompare TREAT against three baseline types: 1) pure data-driven approaches including LG-ODE (Huang et al., 2020) and LatentODE (Rubanova et al., 2019), where the first one is a multiagent approach considering pair-wise interactions, and the second one is a single-agent approach that predicts each trajectory independently; 2) energy-preserving HODEN (Greydanus et al., 2019); and 3) time-reversal TRS-ODEN (Huh et al., 2020).

The latter two are single-agent approaches and require initial states as given input. To handle missing initial states in our dataset, we approximate the initial states for the two methods via linear spline interpolation (Endre Süli, 2003). In addition, we substitute the ODE network in TRS-ODEN with a GNN (Kipf et al., 2018) as TRS-ODENGNN, which serves as a new multi-agent approach for fair comparison. HODEN cannot be easily extended to the multi-agent setting as replacing the ODE function with a GNN can violate energy conservation of the original HODEN. For running LGODE and TREAT on single-agent datasets, we only include self-loop edges in the graph G = ( V , E ) , which makes the ODE function g a simple MLP. Implementation details can be found in Appendix D.2.

4.1 Main Results

Table 1 shows the prediction performance on both multi-agent systems and single-agent systems measured by mean squared error (MSE). We can see that TREAT consistently surpasses other models, highlighting its generalizability and the efficacy of the proposed TRS loss.

For multi-agent systems, approaches that consider interactions among agents (LG-ODE, TRSODENGNN, TREAT) consistently outperform single-agent baselines (LatentODE, HODEN, TRSODEN), and TREAT achieves the best performance across datasets.

The chaotic nature of the Pendulum system and the Attractor system, with their sensitivity to initial states 6 , poses extreme challenges for dynamic modeling. This leads to highly unstable predictions for models like HODEN and TRS-ODEN, as they estimate initial states via inaccurate linear spline interpolation (Endre Süli, 2003). In contrast, LatentODE, LG-ODE, and TREAT employ advanced encoders that infer latent states from observed data and demonstrate superior accuracy. Among them, TREAT achieves the most accurate predictions, further showing its robust generalization capabilities.

We observe that misapplied inductive biases can degrade results, which limits the usage of physicsinformed methods that are designed for individual physical prior such as HODEN. HODEN only excels on energy-conservative systems, such as Simple Spring compared with LatentODE and TRSODENin the multi-agent setting. Its performance drop dramatically on Force Spring , Damped Spring , and Attractor . Note that HODEN naively forces each agent to be energy-conservative, instead of the whole system. Therefore, it performs poorly than LG-ODE, TREAT in the multi-agent settings.

For the Human Motion dataset, characterized by its dynamic ambiguity as it does not adhere to specific physical laws, we cannot directly determine whether it is conservative or time-reversal. For such a system with an unknown nature, TREAT outperforms other purely data-driven methods significantly, showcasing its strong numerical benefits in improving prediction accuracy across diverse system types. This is also shown by its superior performance on Damped Spring , which is irreversible.

(a) Simple Spring

(b) Damped Spring

(c) Forced Spring

(d) Pendulum

LG-ODE

TREAT

TRS-ODEN

HODEN

LG-ODE

TREAT

TRS-ODEN

HODEN

LG-ODE

TREAT

TRS-ODEN

HODEN

LG-ODE

TREAT

TRS-ODEN

HODEN

(a) Simple Spring

(b) Damped Spring

(c) Forced Spring

(d) Pendulum

TREAT

LG-ODE

TREAT

LG-ODE

TREAT

LG-ODE

TREAT

LG-ODE

4.2 Ablation and Sensitivity Analysis

Ablation on implementation of L reverse . We conduct two ablation by changing the implementation of L reverse discussed in Sec. 3.2: 1) TREAT L rev = gt-rev , which computes the reversal loss as the L2 distance between ground truth trajectories to model backward trajectories; 2) TREAT L rev = rev2, which implements the TRS loss based on Eqn 5 as in TRS-ODEN (Huh et al., 2020). From the last block of Table 1, we can clearly see that our implementation achieves the best performance against the two.

Evaluation across prediction lengths. We vary the maximum prediction lengths from 20 to 60 and report model performance as shown in Figure 4. As the prediction step increases, TREAT consistently maintains optimal prediction performance, while other baselines exhibit significant error accumulations. The performance gap between TREAT and baselines widens when making long-range predictions, highlighting the superior predictive capability of TREAT.

Evaluation across different α . We vary the values of the coefficient α defined in Eqn 10, which balances the reconstruction loss and the TRS loss. Figure 5 demonstrates that the optimal α values being neither too high nor too low. This is because when α is too small, the model tends to neglect the TRS physical bias, resulting in error accumulations. Conversely, when α becomes too large, the model can emphasize TRS at the cost of accuracy. Nonetheless, across different α values, TREAT consistently surpasses the purely data-driven LG-ODE, showcasing its superiority and flexibility in modeling diverse dynamical systems.

We study TREAT's sensitivity towards solver choice and observation ratios in Appendix E.1 and Appendix E.2 respectively.

LG-ODE

TREAT

HODEN

LG-ODE

TREAT

HODEN

LG-ODE

TREAT

HODEN

Ground Truth

TANGO

LG-ODE

Energy

HODEN

(a) Simple

Spring

(b) Damped

Spring

(c) Forced

Spring

TREAT

LG-ODE

TREAT

HODEN

LG-ODE

TREAT

HODEN

LG-ODE

TREAT

HODEN

(a) Simple Spring

(b) Damped Spring

(c) Forced Spring

(d) Pendulum

×1#

!"#

×1#

!"#

×1#

!"#

×1#

!"$

TREAT

LG-ODE

TREAT

LG-ODE

TREAT

LG-ODE

TREAT

LG-ODE

LG-ODE Reversal Loss

LG-ODE Reversal Loss

LG-ODE Reversal Loss

LG-ODE Reversal Loss

4.3 Visualizations

Trajectory Visualizations. Model predictions and ground truth are visualized in Figure 6. As HODEN is a single-agent baseline that individually forces every agent's energy to be constant over time which is not valid, the predicted trajectories is having the largest errors and systems' total energy is not conserved for all datasets. The purely data-driven LG-ODE exhibits unrealistic energy patterns, as seen in the energy spikes in Simple Spring and Force Spring . In contrast, TREAT, incorporating reversal loss, generates realistic energy trends, and consistently produces trajectories closest to the ground truth, showing its superior performance.

Reversal Loss Visualizations To illustrate the issue of energy explosion from the purely data-driven LG-ODE, we visualize the TRS loss over training epochs from LG-ODE 7 and TREAT in Figure 7. As results suggest, LG-ODE has increased TRS loss over training epochs, meaning it is violating the time-reversal symmetry sharply, in contrast to TREAT which has decreased reversal loss over epochs.

5 Conclusions

We propose TREAT, a deep learningframework that achieves high-precision modeling for a wide range of dynamical systems by injecting time-reversal symmetry as an inductive bias. TREAT features a novel regularization term to softly enforce time-reversal symmetry by aligning predicted forward and reverse trajectories from a GraphODE model. Notably, we theoretically prove that the regularization term effectively minimizes higher-order Taylor expansion terms during the ODE integration, which serves as a general numerical benefit widely applicable to various systems (even irreversible systems) regardless of their physical properties. Empirical evaluations on different kinds of datasets illustrate TREAT's superior efficacy in accurately capturing real-world system dynamics.

6 Limitations

Currently, TREAT only incorporates inductive bias from the temporal aspect, while there are many important properties in the spatial aspect such as translation and rotation equivariance (Satorras et al., 2021; Han et al., 2022b; Xu et al., 2022). Future endeavors that combine biases from both temporal and spatial dimensions could unveil a new frontier in dynamical systems modeling.

References

Awrejcewicz, G. Kudra, and G. Wasilewski. 2008. Chaotic zones in triple pendulum dynamics observed experimentally and numerically. Applied Mechanics and Materials (2008), 1-17.

Cao, M. Chai, M. Li, and C. Jiang. 2023. Efficient learning of mesh-based physical simulation with bi-stride multi-scale graph neural network. In International Conference on Machine Learning . PMLR, 3541-3558.

Cao, X. Gao, and R. Li. 2019. A liquid plug moving in an annular pipe-Heat transfer analysis. International Journal of Heat and Mass Transfer 139 (2019), 1065-1076.

Cao and R. Li. 2018. A liquid plug moving in an annular pipe-Flow analysis. Physics of Fluids 30, 9 (2018).

Carnegie Mellon University. 2003. Carnegie-Mellon Motion Capture Database. http://mocap. cs.cmu.edu Online database.

Chang, L. Meng, E. Haber, L. Ruthotto, D. Begert, and E. Holtham. 2018. Reversible architectures for arbitrarily deep residual neural networks. In Proceedings of the AAAI conference on artificial intelligence , Vol. 32.

T. Q. Chen, B. Amos, and M. Nickel. 2021. Learning Neural Event Functions for Ordinary Differential Equations. International Conference on Learning Representations (2021).

Q. Chen, Y. Rubanova, J. Bettencourt, and D. Duvenaud. 2018. Neural Ordinary Differential Equations. In Advances in Neural Information Processing Systems .

Cranmer, S. Greydanus, S. Hoyer, P. Battaglia, D. Spergel, and S. Ho. 2020. Lagrangian neural networks. arXiv preprint arXiv:2003.04630 (2020).

F. M. Endre Süli. 2003. An Introduction to Numerical Analysis . Cambridge University Press. 293 pages.

Greydanus, M. Dzamba, and J. Yosinski. 2019. Hamiltonian neural networks. Advances in Neural Information Processing Systems (2019).

Anthony Gruber, Kookjin Lee, and Nathaniel Trask. 2024. Reversible and irreversible bracket-based dynamics for deep graph neural networks. Advances in Neural Information Processing Systems 36 (2024).

Jiaqi Han, Wenbing Huang, Hengbo Ma, Jiachen Li, Joshua B. Tenenbaum, and Chuang Gan. 2022a. Learning Physical Dynamics with Subequivariant Graph Neural Networks. In Advances in Neural Information Processing Systems , Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (Eds.). https://openreview.net/forum?id=siG_S8mUWxf

Jiaqi Han, Wenbing Huang, Tingyang Xu, and Yu Rong. 2022b. Equivariant graph hierarchy-based neural networks. Advances in Neural Information Processing Systems 35 (2022), 9176-9187.

Hu, Y. Dong, K. Wang, and Y. Sun. 2020. Heterogeneous Graph Transformer. In Proceedings of the 2020 World Wide Web Conference .

Zijie Huang, Jeehyun Hwang, Junkai Zhang, Jinwoo Baik, Weitong Zhang, Dominik Wodarz, Yizhou Sun, Quanquan Gu, and Wei Wang. 2024. Causal Graph ODE: Continuous Treatment Effect Modeling in Multi-agent Dynamical Systems. In Proceedings of the ACM Web Conference 2024 (Singapore, Singapore) (WWW '24) . 4607-4617.

Huang, Y. Sun, and W. Wang. 2020. Learning Continuous System Dynamics from IrregularlySampled Partial Observations. In Advances in Neural Information Processing Systems .

Huang, Y. Sun, and W. Wang. 2021. Coupled Graph ODE for Learning Interacting System Dynamics. In Proceedings of the 27th ACM SIGKDD Conference on Knowledge Discovery and Data Mining .

Huang, Y. Sun, and W. Wang. 2023. Generalizing Graph ODE for Learning Complex System Dynamics across Environments. In Proceedings of the 29th ACM SIGKDD Conference on Knowledge Discovery and Data Mining (KDD '23) . 798-809.

Huh, E. Yang, S. J. Hwang, and J. Shin. 2020. Time-Reversal Symmetric ODE Network. In Advances in Neural Information Processing Systems .

Jiang, Z. Huang, X. Luo, and Y. Sun. 2023. CF-GODE: Continuous-Time Causal Inference for Multi-Agent Dynamical Systems. In Proceedings of the 29th ACM SIGKDD Conference on Knowledge Discovery and Data Mining .

Kim and S. J. Karrila. 2013. Microhydrodynamics: principles and selected applications . Courier Corporation.

Kipf, E. Fetaya, K. Wang, M. Welling, and R. Zemel. 2018. Neural Relational Inference for Interacting Systems. arXiv preprint arXiv:1802.04687 (2018).

Koppe, H. Toutounji, P. Kirsch, S. Lis, and D. Durstewitz. 2019. Identifying nonlinear dynamical systems via generative recurrent neural networks with applications to fMRI. PLoS computational biology 15, 8 (2019), e1007263.

Lam, A. Sanchez-Gonzalez, M. Willson, P. Wirnsberger, M. Fortunato, F. Alet, S. Ravuri, T. Ewalds, Z. Eaton-Rosen, W. Hu, A. Merose, S. Hoyer, G. Holland, O. Vinyals, J. Stott, A. Pritzel, S. Mohamed, and P. Battaglia. 2023. Learning skillful medium-range global weather forecasting. Science 382, 6677 (2023), 1416-1421.

S. Lamb and J. A. Roberts. 1998. Time-reversal symmetry in dynamical systems: a survey. Physica D: Nonlinear Phenomena (1998), 1-39.

Li, F. Xia, R. Martín-Martín, M. Lingelbach, S. Srivastava, B. Shen, K. E. Vainio, C. Gokmen, G. Dharan, T. Jain, A. Kurenkov, K. Liu, H. Gweon, J. Wu, L. Fei-Fei, and S. Savarese. 2022. iGibson 2.0: Object-Centric Simulation for Robot Learning of Everyday Household Tasks. In Proceedings of the 5th Conference on Robot Learning .

Liu, A. Kumar, J. Ba, J. Kiros, and K. Swersky. 2019. Graph normalizing flows. Advances in Neural Information Processing Systems 32 (2019).

Loshchilov and F. Hutter. 2019. Decoupled weight decay regularization. In The International Conference on Learning Representations .

Lu, S. Lin, G. Chen, and J. Pan. 2022. ModLaNets: Learning Generalisable Dynamics via Modularity and Physical Inductive Bias. In Proceedings of the 39th International Conference on Machine Learning (Proceedings of Machine Learning Research, Vol. 162) , Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato (Eds.). PMLR, 1438414397.

Xiao Luo, Yiyang Gu, Huiyu Jiang, Hang Zhou, Jinsheng Huang, Wei Ju, Zhiping Xiao, Ming Zhang, and Yizhou Sun. 2024. PGODE: Towards High-quality System Dynamics Modeling. In Forty-first International Conference on Machine Learning .

Xiao Luo, Haixin Wang, Zijie Huang, Huiyu Jiang, Abhijeet Sadashiv Gangan, Song Jiang, and Yizhou Sun. 2023a. CARE: Modeling Interacting Dynamics Under Temporal Environmental Variation. In Thirty-seventh Conference on Neural Information Processing Systems . https: //openreview.net/forum?id=lwg3ohkFRv

Luo, J. Yuan, Z. Huang, H. Jiang, Y. Qin, W. Ju, M. Zhang, and Y. Sun. 2023b. HOPE: Highorder Graph ODE For Modeling Interacting Dynamics. In Proceedings of the 40th International Conference on Machine Learning .

Xiao Luo, Jingyang Yuan, Zijie Huang, Huiyu Jiang, Yifang Qin, Wei Ju, Ming Zhang, and Yizhou Sun. 2023c. Hope: High-order graph ode for modeling interacting dynamics. In International Conference on Machine Learning . PMLR, 23124-23139.

Karolis Martinkus, Aurelien Lucchi, and Nathanaël Perraudin. 2021. Scalable graph networks for particle simulations. In Proceedings of the AAAI Conference on Artificial Intelligence , Vol. 35. 8912-8920.

Ni and A. H. Qureshi. 2022. Ntfields: Neural time fields for physics-informed robot motion planning. arXiv preprint arXiv:2210.00120 (2022).

North. 2021. Formulations of classical mechanics. Forthcoming in A companion to the philosophy of physics. Routledge (2021).

Pfaff, M. Fortunato, A. Sanchez-Gonzalez, and P. Battaglia. 2021. Learning Mesh-Based Simulation with Graph Networks. In International Conference on Learning Representations .

Tobias Pfaff, Meire Fortunato, Alvaro Sanchez-Gonzalez, and Peter W Battaglia. 2020. Learning mesh-based simulation with graph networks. arXiv preprint arXiv:2010.03409 (2020).

Poli, S. Massaroli, J. Park, A. Yamashita, H. Asama, and J. Park. 2019. Graph neural ordinary differential equations. arXiv preprint arXiv:1911.07532 (2019).

Pozrikidis. 2001. Interfacial dynamics for Stokes flow. J. Comput. Phys. 169, 2 (2001), 250-301.

Raissi, P. Perdikaris, and G. E. Karniadakis. 2019. Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations. Journal of Computational physics 378 (2019), 686-707.

Rubanova, R. T. Chen, and D. K. Duvenaud. 2019. Latent ordinary differential equations for irregularly-sampled time series. In Advances in Neural Information Processing Systems .

Sanchez-Gonzalez, V. Bapst, K. Cranmer, and P. Battaglia. 2019. Hamiltonian Graph Networks with ODE Integrators. In Advances in Neural Information Processing Systems .

Sanchez-Gonzalez, J. Godwin, T. Pfaff, R. Ying, J. Leskovec, and P. W. Battaglia. 2020. Learning to Simulate Complex Physics with Graph Networks. In Proceedings of the 37th International Conference on Machine Learning .

G. Satorras, E. Hoogeboom, and M. Welling. 2021. E (n) equivariant graph neural networks. In International conference on machine learning . PMLR, 9323-9332.

Schober, S. Särkkä, and P. Hennig. 2019. A probabilistic model for the numerical solution of initial value problems. In Statistics and Computing . 99-122.

Sepp and S. Jürgen. 1997. Long Short-term Memory. Neural computation (1997).

Shinbrot, C. Grebogi, J. Wisdom, and J. A. Yorke. 1992. Chaos in a double pendulum. American Journal of Physics 6 (1992), 491-499.

C. Sprott. 2015. Symmetric time-reversible flows with a strange attractor. International Journal of Bifurcation and Chaos 25, 05 (2015), 1550078.

Stachowiak and T. Okada. 2006. A numerical analysis of chaos in the double pendulum. Chaos, Solitons & Fractals 2 (2006), 417-422.

C. Tolman. 1938. The Determiners of Behavior at a Choice Point. Psychological Review 45, 1 (1938), 1-41.

Valperga, K. Webster, D. Turaev, V. Klein, and J. Lamb. 2022. Learning Reversible Symplectic Dynamics. In Proceedings of The 4th Annual Learning for Dynamics and Control Conference .

Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, ˙ U. Kaiser, and I. Polosukhin. 2017. Attention is All you Need. In Advances in Neural Information Processing Systems .

Haixin Wang, Yadi Cao, Zijie Huang, Yuxuan Liu, Peiyan Hu, Xiao Luo, Zezheng Song, Wanjia Zhao, Jilin Liu, Jinan Sun, et al. 2024. Recent Advances on Machine Learning for Computational Fluid Dynamics: A Survey. arXiv preprint arXiv:2408.12171 (2024).

Wang, K. Kashinath, M. Mustafa, A. Albert, and R. Yu. 2020. Towards physics-informed deep learning for turbulent flow prediction. In Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery and Data Mining .

Wen, H. Wang, and D. Metaxas. 2022. Social ODE: Multi-agent Trajectory Forecasting with Neural Ordinary Differential Equations. In European Conference on Computer Vision .

Minkai Xu, Jiaqi Han, Aaron Lou, Jean Kossaifi, Arvind Ramanathan, Kamyar Azizzadenesheli, Jure Leskovec, Stefano Ermon, and Anima Anandkumar. 2024. Equivariant Graph Neural Operator for Modeling 3D Dynamics. In Forty-first International Conference on Machine Learning . https: //openreview.net/forum?id=dccRCYmL5x

Minkai Xu, Lantao Yu, Yang Song, Chence Shi, Stefano Ermon, and Jian Tang. 2022. Geodiff: A geometric diffusion model for molecular conformation generation. arXiv preprint arXiv:2203.02923 (2022).

Zang and F. Wang. 2020. Neural dynamics on complex networks. In Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery and Data Mining .

Yaofeng Desmond Zhong, Biswadip Dey, and Amit Chakraborty. 2020. Dissipative symoden: Encoding hamiltonian dynamics with dissipation and control into deep learning. arXiv preprint arXiv:2002.08860 (2020).

A Theoretical Analysis

A.1 Implementation of the Time-Reversal Symmetry Loss

Algorithm 1 The implementation of L reverse

Require: latent initial states z fwd i ( t 0 ) ; the ODE function g ( · ) ; number of agents N :

1: for each i ∈ N do

2: Compute the latent forward trajectory at timestamps { t k } k =0 : z fwd i ( t k ) = ODE-Solver ( g, [ z fwd 1 ( t 0 ) , z fwd 2 ( t 0 ) ... z fwd N ( t 0 )] , t k ) . Reach the final state z fwd i ( t K ) .

K

3: The initial state of the reverse trajectory is defined as z rev i ( t ′ 0 ) = z fwd i ( t K ) , and the dynamics of the system which is the ODE function g ( · ) is also reversed as -g ( · ) .

4: Compute the latent reverse trajectory at timestamps { t ′ k } K k =0 , z rev i ( t ′ k ) = ODE-Solver ( g, [ z rev 1 ( t ′ 0 ) , z rev 2 ( t ′ 0 ) ... z rev N ( t ′ 0 )] , t ′ k ) .

5: ˆ y fwd i ( t k ) = f DEC ( z fwd i ( t k )) , ˆ y rev i ( t ′ k ) = f DEC ( z rev i ( t ′ k ))

7: L reverse = ∑ N i =1 ∑ K k =0 ∥ ∥ ∥ ˆ y fwd i ( t k ) -ˆ y rev i ( t ′ K -k ) ∥ ∥ ∥ 2 2

6: end for

A.2 Proof of Lemma 1

Proof. The definition of time-reversal symmetry is given by:

Here, R is an involution operator, which means R ◦ R = I.

First, we apply the time evolution operator ϕ t to both sides of Eqn 11:

Next, we apply the involution operator R to both sides of the equation:

Since R ◦ R = I, we finally arrive at:

Simplifying, we obtain:

which means the trajectories can overlap when evolving backward from the final state.

A.3 Proof of Theorem 3.1

Let ∆ t denote the integration step size in an ODE solver and T be the prediction length. The time stamps of the ODE solver are { t j } T j =0 , where t j +1 -t j = ∆ t for j = 0 , · · · , T ( T > 1) . Next suppose during the forward evolution, the updates go through states z fwd ( t j ) = ( q fwd ( t j ) , p fwd ( t j )) for j = 0 , · · · , T , where q fwd ( t j ) is position, p fwd ( t j ) is momentum, while during the reverse evolution they go through states z rev ( t j ) = ( q rev ( t j ) , p rev ( t j )) for j = 0 , · · · , T , in reverse order. The ground truth trajectory is z gt ( t j ) = ( q gt ( t j ) , p gt ( t j )) for j = 0 , · · · , T .

For the sake of brevity in the ensuing proof, we denote z gt ( t j ) by z gt j , z fwd ( t j ) by z fwd j and z rev ( t j ) by z rev j , and we will use Mathematical Induction to prove the theorem.

A.3.1 Reconstruction Loss ( L pred ) Analysis.

First, we bound the forward loss ∑ T j =0 ∥ z fwd j -z gt j ∥ 2 2 . Since our method models the momentum and position of the system, we can write the following Taylor expansion of the forward process, where

for any 0 ≤ j < T :

and for the ground truth process, we also have from Taylor expansion that

With these, we aim to prove that for any k = 0 , 1 , · · · , T , the following hold :

where C fwd 1 and C fwd 2 are constants.

Base Case k = 0 : Based on the initialization rules, it is obvious that ∥ ∥ q fwd 0 -q gt 0 ∥ ∥ 2 = 0 and ∥ ∥ p fwd 0 -p gt 0 ∥ ∥ 2 = 0 , thus (18a) and (18b) both hold for k = 0 .

Inductive Hypothesis: Assume (18a) and (18b) hold for k = j , which means:

Inductive Proof: We need to prove (18a) and (18b) hold for k = j +1 .

First, using (16c) and (17c), we have

where we iterate through j, j -1 , · · · , 0 in the second equality. Then using (17b) and (16b), we get for j +1 that

where the first inequality uses the triangle inequality, and in the second inequality we use (19b) as well as (20). We can see there exists C fwd 1 such that the final expression above is upper bounded by C fwd 1 ( j +1)∆ t , with which the claim holds for j +1 .

Next for (18a), using (17a) and (16a), we get for any j that

where the first inequality uses the triangle inequality, and in the second inequality we use (19a) and (19b) as well as (20). Thus with an appropriate C fwd 2 , we have the final expression above is upper bounded by C fwd 2 ( j +1) 2 ∆ t 2 , and so the claim holds for j +1 .

Since both the base case and the inductive step have been proven, by the principle of mathematical induction, (18a) and (18b) holds for all k = 0 , 1 , · · · , T .

With this, we finish the forward proof by plugging (18a) and (18b) into the loss function:

A.3.2 Reversal Loss ( L reverse ) Analysis.

Next we analyze the reversal loss ∑ T j =0 ∥ R ( z rev j ) -z fwd j ∥ 2 2 . For this, we need to refine the Taylor expansion residual terms for a more in-depth analysis.

First reconsider the forward process. Since the process is generated from the learned network, we may assume that for some constants c 1 , c 2 , and c 3 , the states satisfy the following for any 0 ≤ j < T :

where the remaining terms ∥ ∥ rem fwd ,i j ∥ ∥ 2 ≤ c i ∆ t i for i = 1 , 2 , 3 . Similarly, we have approximate Taylor expansions for the reverse process:

where ∥ ∥ rem rev ,i j ∥ ∥ 2 ≤ c i ∆ t i for i = 1 , 2 , 3 .

We will prove via induction that for k = T, T -1 , · · · , 0 ,

where C rev 1 , C rev 2 and C rev 3 are constants.

The entire proof process is analogous to the previous analysis of Reconstruction Loss.

Base Case k = T : Since the reverse process is initialized by the forward process variables at k = T , it is obvious that ∥ ∥ q fwd T -q ev T ∥ ∥ 2 = ∥ ∥ p fwd T -p rev T ∥ ∥ 2 = ∥ ∥ ˙ p fwd T -˙ p rev T ∥ ∥ 2 = 0 . Thus (23a), (23b) and (23c) all hold for k = 0 .

Inductive Hypothesis: Assume the inequalities (23b), (23a) and (23c) hold for k = j +1 , which means:

Inductive Proof : We need to prove (23b) (23a) and (23c) holds for k = j .

First, for (23c), using (21c) and (22c), we get for any j that

where the first inequality uses the triangle inequality, and the second inequality plugs in (24c). Thus taking C rev 1 = 2 c 1 , the above is upped bounded by C rev 1 ( T -j )∆ t , and (23b) holds for j .

Second, for (24b), using (21b) and (22b), we get

where the first inequality uses the triangle inequality, and in the second inequality we use (24a) and (24b). Thus taking C rev 2 = max { C rev 1 / 2 , 2 c 2 } , we have the final expression above is upper bounded by C rev 2 ( T -j ) 2 ∆ t 2 , and so the claim holds for j .

Finally, for (24a), we use (21a) and (22a) to get

where the first inequality uses the triangle inequality, and in the second inequality we use (24a), (24b) and (24c). Thus taking C rev 3 = max { C rev 2 / 3 m,C rev 1 / 6 m, 2 c 3 } , we have the final expression above is upper bounded by C rev 3 ( T -j ) 3 ∆ t 3 , and so the claim holds for j .

Since both the base case and the inductive step have been proven, by the principle of mathematical induction, (23b), (23a) and (23c) hold for all k = T, T -1 , · · · , 0 .

With this we finish the proof by plugging (23b) and (23a) into the loss function:

A.4 Proof of Lemma 3.2

b

a

b

a

𝒚 $

𝒊

𝐟𝐰𝐝

𝟎

= 𝒚

𝒊

(𝟎)

𝒚

𝒊

(𝟏)

𝒚 $

𝒊

𝐫𝐞𝐯𝟐

𝟎

= 𝒚

𝒊

𝟎

= 𝒚

$

𝒊

𝐟𝐰𝐝

𝟎

𝒚

𝒊

(𝟏)

𝒚 $

𝒊

𝐟𝐰𝐝

𝟏

𝒚 $

𝒊

𝐟𝐰𝐝

𝟏

𝒚 $

𝒊

𝐫𝐞𝐯𝟐

-𝟏

𝒚 $

𝒊

𝐫𝐞𝐯𝟐

𝟏

𝒚 $

𝒊

𝐫𝐞𝐯

𝟎

𝒚 $

𝒊

𝐫𝐞𝐯

-𝟏

= 𝑹(𝒚

$

𝒊

𝐟𝐰𝐝

𝟏 )

𝑀𝑎𝑥𝐸𝑟𝑟𝑜𝑟

!"#$!

= max	 {𝑎, 𝑏}

𝑀𝑎𝑥𝐸𝑟𝑟𝑜𝑟

!"%-'(#)

= 𝑎 + 𝑏

ℒ

!"#$

=

𝒚 $

𝒊

𝐟𝐰𝐝

𝟏

-𝒚

𝒊

(𝟏)

)

)

∶= 𝑎

ℒ

"#*#"+#

=

𝑅(	𝒚

$

𝒊

𝐫𝐞𝐯

𝟎 ) - 𝒚

$

𝒊

𝐟𝐰𝐝

𝟎

)

)

≔𝑏

ℒ

!"#$

=

𝒚 $

𝒊

𝐟𝐰𝐝

𝟏

-𝒚

𝒊

(𝟏)

)

)

∶= 𝑎

ℒ

"#*#"+#)

=

𝒚 $

𝒊

𝐫𝐞𝐯𝟐

𝟏

-𝒚 $

𝒊

𝐟𝐰𝐝

𝟏

)

)

≔𝑏

Ground truth Trajectory

Forward Trajectory

Reverse Trajectory

TREAT:	𝑅 ∘ 𝜙

!

∘ 𝑅 ∘ 𝜙

!

= 𝐼

TRS-ODEN:	𝑅 ∘ 𝜙

!

= 𝜙

"!

∘ 𝑅

𝒚

$

𝒊

𝐫𝐞𝐯𝟐

𝟏

= 𝑹(	𝒚

$

𝒊

𝐫𝐞𝐯𝟐

-𝟏 )

We expect an ideal model to align both the predicted forward and reverse trajectories with the ground truth. As shown in Figure 8, we integrate one step from the initial state ˆ y fwd i (0) (which is the same as y i (0) ) and reach the state ˆ y fwd i (1) .

The first reverse loss implementation (ours) follows Lemma 2.1 as R ◦ Φ t ◦ R ◦ Φ t = I, which means when we evolve forward and reach the state ˆ y fwd i (1) we reverse it into ˆ y rev i ( -1) = R ( ˆ y fwd i (1)) and go back to reach ˆ y rev i (0) , then reverse it to get R ( ˆ y rev i (0)) , which ideally should be the same as ˆ y fwd i (0) .

The second reverse loss implementation follows Eqn 5as R ◦ Φ t = Φ -t ◦ R , which means we first reverse the initial state as ˆ y rev 2 i (0) = R ( y i (0)) , then evolve the reverse trajectory in the opposite direction to reach ˆ y rev 2 i ( -1) , and then perform a symmetric operation to reach ˆ y rev 2 i (1) , aligning it with the forward trajectory.

We assume the two reconstruction losses L pred = ∥ ˆ y fwd i (1) -y i (1) ∥ 2 2 := a are the same. For the time-reversal losses, we also assume they have reached the same value b :

As shown in Figure 8 where we illustrate the worst case scenario MaxError gt _ rev = max k ∈ [ K ] ∥ y i ( t k ) -ˆ y rev i ( t ′ K -k ) ∥ 2 of TREAT and TRS-ODEN, we can see that in our implementation the worst error is the maximum of two loss, while the TRS-ODEN's implementation has the risk of accumulating the error together, making the worst error being the sum of both:

So it is obvious that MaxError TREAT made by TREAT is smaller., which means our model achieves a smaller error of the maximum distance between the reversal and ground truth trajectory.

B Example of varying dynamical systems

We illustrate the energy conservation and time reversal of the three n-body spring systems used in our experiments. We use the Hamiltonian formalism of systems under classical mechanics to describe their dynamics and verify their energy conservation and time-reversibility characteristics.

The scalar function that describes a system's motion is called the Hamiltonian, H , and is typically equal to the total energy of the system, that is, the potential energy plus the kinetic energy (North, 2021). It describes the phase space equations of motion by following two first-order ODEs called Hamilton's equations:

where q ∈ R n , p ∈ R n , and H : R 2 n ↦→ R are positions, momenta, and Hamiltonian of the system.

, and the time-reversal symmetry

Under this formalism, energy conservative is defined by d H /dt = 0 is defined by H ( q, p, t ) = H ( q, -p, -t ) (Lamb and Roberts, 1998).

B.1 Conservative and reversible systems.

A simple example is the isolated n-body spring system, which can be described by :

where q = ( q 1 , q 2 , · · · , q N ) is a set of positions of each object , p = ( p 1 , p 2 , · · · , p N ) is a set of momenta of each object, m i is mass of each object, k is spring constant.

The Hamilton's equations are:

Hence, we can obtain the Hamiltonian through the integration of the above equation.

Verify the systems' energy conservation

So it is conservative.

Verify the systems' time-reversal symmetry We do the transformation R : ( q , p , t ) ↦→ ( q , -p , -t ) .

It is obvious H ( q , p ) = H ( q , -p ) , so it is reversible

B.2 Non-conservative and reversible systems.

A simple example is a n-body spring system with periodical external force, which can be described by:

The Hamilton's equations are:

Hence, we can obtain the Hamiltonian through the integration of the above equation:

Verify the systems' energy conservation

̸

So it is non-conservative.

Verify the systems' time-reversal symmetry We do the transformation R : ( q , p , t ) ↦→ ( q , -p , -t ) .

It is obvious H ( q , p , t ) = H ( q , -p , t ) , so it is reversible

B.3 Non-conservative and irreversible systems.

A simple example is an n-body spring system with frictions proportional to its velocity, γ is the coefficient of friction, which can be described by:

The Hamilton's equations are:

Hence, we can obtain the Hamiltonian through the integration of the above equation:

Verify the systems' energy conservation

̸

So it is non-conservative.

Verify the systems' time-reversal symmetry We do the transformation R : ( q , p , t ) ↦→ ( q , -p , -t ) .

̸

It is obvious H ( q , p , t ) = H ( q , -p , t ) , so it is irreversible

C Dataset

In our experiments, all datasets are synthesized from ground-truth physical law via sumulation. We generate five simulated datasets: three n -body spring systems under damping, periodic, or no external force, one chaotic tripe pendulum dataset with three sequentially connected stiff sticks that form and a chaotic strange attractor. We name the first three as Sipmle Spring , Forced Spring , and Damped Spring respectively. For multi-agent systems, all n -body spring systems contain 5 interacting balls, with varying connectivities. Each Pendulum system contains 3 connected stiff sticks. For single-agent systems, all spring systems contain only one ball. For the chaotic single Attractor , we follow the setting of (Huh et al., 2020).

For the n -body spring system, we randomly sample whether a pair of objects are connected, and model their interaction via forces defined by Hooke's law. In the Damped spring , the objects have an additional friction force that is opposite to their moving direction and whose magnitude is proportional to their speed. In the Forced spring , all objects have the same external force that changes direction periodically. We show in Figure 1(a), the energy variation in both of the Damped spring and Forced spring is significant. For the chaotic triple Pendulum , the equations governing the motion are inherently nonlinear. Although this system is deterministic, it is also highly sensitive to the initial condition and numerical errors (Shinbrot et al., 1992; Awrejcewicz et al., 2008; Stachowiak and Okada, 2006). This property is often referred to as the "butterfly effect", as depicted in Figure 9. Unlike for n -body spring systems, where the forces and equations of motion can be easily articulated, for the Pendulum , the explicit forces cannot be directly defined, and the motion of objects can only be described through Lagrangian formulations (North, 2021), making the modeling highly complex and raising challenges for accurate learning. We simulate the trajectories by using Euler's method for

0

100

200

300

400

500

600

700

Time steps

10

0

10

20

30

40

Joint

Original initial condition:

0

w/ 1e-3 perturbation:

0

w/ 1e-2 perturbation:

0

1

1

1

2

2

2

n -body spring systems and using the 4th order Runge-Kutta (RK4) method for the Pendulum and Attractor . For all spring systems and Pendulum , We integrate with a fixed step size and subsample every 100 steps. For training, we use a total of 6000 forward steps. To generate irregularly sampled partial observations, we follow (Huang et al., 2020) and sample the number of observations n from a uniform distribution U(40, 52) and draw the n observations uniformly for each object. For testing, we additionally sample 40 observations following the same procedure from PDE steps [6000, 12000], besides generating observations from steps [1, 6000]. The above sampling procedure is conducted independently for each object. We generate 20k training samples and 5k testing samples for each dataset. For Attractor , we integrate a total of 600 forward steps for training and subsample every 10 steps. For testing, we additionally sample 40 observations from step [600,1200].The irregularly sampled partial observations generation is the same as above. We generate 1000 training samples and 50 testing samples following (Huh et al., 2020). Therefore, for all datasets, condition length is

60 steps and prediction length is 40s steps. The features (position/velocity) are normalized to the maximum absolute value of 1 across training and testing datasets.

We also compute the Maximum Lyapunov Exponent (MLE) to assess the chaos level of the systems, using the formula:

We set fixed initial values for each dataset and generate 10 trajectories by perturbing the initial values with random noise (0, 0.0001). We calculate the Maximum Lyapunov Exponent (MLE) between any two trajectories. Finally, we compute the average and std of MLE from all pairs to gauge the chaotic behavior of each dataset. The data is presented in the table below:

From the table, it's evident that the order of MLE values is: Pendulum » three Spring datasets. This observation is consistent with the evaluation results based on MSE presented in our previous responses in Table 3 which indicates that as the prediction length(steps*step size) increases, there is a more significant performance degradation of all models on Pendulum dataset.

In the following subsections, we show the dynamical equations of each dataset in detail.

C.1 Spring Systems

C.1.1 Simple Spring

The dynamical equations of simple spring are as follows:

where where q = ( q 1 , q 2 , · · · , q N ) is a set of positions of each object , p = ( p 1 , p 2 , · · · , p N ) is a set of momenta of each object. We set the mass of each object m = 1 , the spring constant k = 0 . 1 .

C.1.2 Damped Spring

The dynamical equations of damped spring are as follows:

where where q = ( q 1 , q 2 , · · · , q N ) is a set of positions of each object, p = ( p 1 , p 2 , · · · , p N ) is a set of momenta of each object, We set the mass of each object m = 1 , the spring constant k = 0 . 1 , the coefficient of friction γ = 10 .

C.1.3 Forced Spring

The dynamical equations of forced spring system are as follows:

where where q = ( q 1 , q 2 , · · · , q N ) is a set of positions of each object , p = ( p 1 , p 2 , · · · , p N ) is a set of momenta of each object. We set the mass of each object m = 1 , the spring constant k = 0 . 1 , the external strength k 1 = 10 and the frequency of variation ω = 1

We simulate the positions and momentums of three spring systems by using Euler methods as follows:

where d q i dt and d p i dt were defined as above for each datasets, and ∆ t = 0 . 001 is the integration steps.

C.2 Chaotic Pendulum

In this section, we demonstrate how to derive the dynamics equations for a chaotic triple pendulum using the Lagrangian formalism.

The moment of inertia of each stick about the centroid is

The position of the center of gravity of each stick is as follows:

The change in the center of gravity of each stick is:

The Lagrangian L of this triple pendulum system is:

The Lagrangian equation is defined as follows:

and we also have:

where p is the Angular Momentum.

We can list the equations for each of the three sticks separately:

Finally, we have :

We simulate the angular of the three sticks by using the Runge-Kutta 4th Order Method as follows:

where ˙ θ was defined as above , and ∆ t = 0 . 0001 is the integration steps.

C.3 Chaotic Strange Attractor

The dynamical equations of this reversible strange attractor are as follows:

The above equations can be presented as ( ˙ x ( t ) , ˙ y ( t ) , ˙ z ( t )) = Dynamic ( x ( t ) , y ( t ) , z ( t )) . We simulate K ( t ) = ( x ( t ) , y ( t ) , z ( t )) by using the Runge-Kutta 4th Order Method as follows:

We sampling z ( t 0 ) randomly from uniform distribution [1, 3] while fixing x ( t 0 ) = y ( t 0 ) = 0 . We set the trajectory lengths of both training and test dataset to 600, with regular time-step size ∆ t = 0 . 03 and the sample frequency of 10. We add Gaussian noise 0 . 05 n , n ∼ N (0 , 1) to training trajectories.

C.4 Human Motion

For the real-world motion capture dataset(Carnegie Mellon University, 2003), we focus on the walking sequences of subject 35. Each sample in this dataset is represented by 31 trajectories, each corresponding to the movement of a single joint. For each joint, we first randomly sample the number of observations from a uniform distribution U (30 , 42) and then sample uniformly from the first 50 frames for training and validation trajectories. For testing, we additionally sampled 40 observations from frames [51 , 99] .We split different walking sequences into training (15 trials) and test sets (7 trials). For each walking sequence, we further split it into several non-overlapping small sequences with maximum length 50 for training, and maximum length 100 for testing. In this way, we generate total 120 training samples and 27 testing samples. We normalize all features (position/velocity) to maximum absolute value of 1 across training and testing datasets.

D Model Details

In the following we introduce in details how we implement our model and each baseline.

D.1 Initial State Encoder

For multi-agent systems, the initial state encoder computes the latent node initial states z i ( t ) for all agents simultaneously considering their mutual interaction. Specifically, it first fuses all observations into a temporal graph and conducts dynamic node representation through a spatial-temporal GNN as in (Huang et al., 2020):

̸

where || denotes concatenation; σ ( · ) is a non-linear activation function; d is the dimension of node embeddings. The node representation is computed as a weighted summation over its neighbors plus residual connection where the attention score is a transformer-based (Vaswani et al., 2017) dot-product of node representations by the use of value, key, query projection matrices W v , W k , W q . Here h l j ( t ) is the representation of agent j at time t in the l -th layer. i ( t ′ ) is the general index for neighbors connected by temporal edges (where t ′ = t ) and spatial edges (where t = t ′ and i = j ). The temporal encoding (Hu et al., 2020) is added to a neighborhood node representation in order to distinguish its message delivered via spatial and temporal edges. Then, we stack L layers to get the final representation for each observation node: h t i = h L i ( t ) . Finally, we employ a self-attention mechanism to generate the sequence representation u i for each agent as their latent initial states:

̸

where a i is the average of observation representations with a nonlinear transformation W a and ˆ h t i = h t i + TE ( t ) . K is the number of observations for each trajectory. Compared with recurrent models such as RNN, LSTM (Sepp and Jürgen, 1997), it offers better parallelization for accelerating training speed and in the meanwhile alleviates the vanishing/exploding gradient problem brought by long sequences. For single-agent Systems, there only left the self-attention mechanism component.

Given the latent initial states, the dynamics of the whole system are determined by the ODE function g which we parametrize as a GNN as in (Huang et al., 2020) for Multi-Agent Systems to capture the continuous interaction among agents. For single-agent systems, we only include self-loop edges in the graph G = ( V , E ) , which makes the ODE function g a simple MLP.

We then employ Multilayer Perceptron (MLP) as a decoder to predict the trajectories ˆ y i ( t ) from the latent states z i ( t ) .

D.2 Implementation Details

TREAT

For multi-agent systems, our implementation of TREAT follows GraphODE pipeline. We implement the initial state encoder using a 2-layer GNN with a hidden dimension of 64 across all datasets. We use ReLU for nonlinear activation. For the sequence self-attention module, we set the output dimension to 128. The encoder's output dimension is set to 16, and we add 64 additional dimensions initialized with all zeros to the latent states z i ( t ) to stabilize the training processes as in (Huang et al., 2021). The GNN ODE function is implemented with a single-layer GNN from (Kipf et al., 2018) with hidden dimension 128. For single-agent systems, we only include self-loop edges in the graph G = ( V , E ) , which makes the ODE function g a simple MLP. To compute trajectories, we use the Runge-Kutta method from torchdiffeq python package s(Chen et al., 2021) as the ODE solver and a one-layer MLP as the decoder.

We implement our model in pytorch. Encoder, generative model, and the decoder parameters are jointly optimized with AdamW optimizer (Loshchilov and Hutter, 2019) using a learning rate of 0.0001 for spring datasets and 0.00001 for Pendulum . The batch size for all datasets is set to 512.

TREAT L rev = gt-rev and TREAT L rev = rev2 share the same architecture and hyparameters as TREAT, with different implementations of the loss function. In TREAT L rev = gt-rev, instead of comparing forward and reverse trajectories, we look at the L2 distance between the ground truth and reverse trajectories when computing the reversal loss.

For TREAT L rev = rev2, we implement the reversal loss following (Huh et al., 2020) with one difference: we do not apply the reverse operation to the momentum portion of the initial state to the ODE function. This is because the initial hidden state is an output of the encoder that mixes position and momentum information. Note that we also remove the additional dimensions to the latent state that TREAT has. To reproduce our model's results, we provide our code implementation link here.

LatentODE

We implement the Latent ODE sequence to sequence model as specified in (Rubanova et al., 2019). We use a 4-layer ODE function in the recognition ODE, and a 2-layer ODE function in the generative ODE. The recognition and generative ODEs use Euler and Dopri5 as solvers (Chen et al., 2021), respectively. The number of units per layer is 1000 in the ODE functions and 50 in GRU update networks. The dimension of the recognition model is set to 100. The model is trained with a learning rate of 0.001 with an exponential decay rate of 0.999 across different experiments. Note that since latentODE is a single-agent model, we compute the trajectory of each object independently when applying it to multi-agent systems.

HODEN

To adapt HODEN, which requires full initial states of all objects, to systems with partial observations, we compute each object's initial state via linear spline interpolation if it is missing. Following the setup in (Huh et al., 2020), we have two 2-layer linear networks with Tanh activation in between as ODE functions, in order to model both positions and momenta. Each network has a 1000-unit layer followed by a single-unit layer. The model is trained with a learning rate of 0.00001 using a cosine scheduler.HODEN is a single-agent model, we compute the trajectory of each object independently when applying it to multi-agent systems.

TRS-ODEN

Similar to HODEN, we compute each object's initial state via linear spline interpolation if it is missing. As in (Huh et al., 2020), we use a 2-layer linear network with Tanh activation in between as the ODE functions, and the Leapfrog method for solving ODEs. The network has 1000 hidden units and is trained with a learning rate of 0.00001 using a cosine scheduler. TRS-ODEN is a single-agent model, we compute the trajectory of each object independently when applying it to multi-agent systems.

TRS-ODENGNN

For TRSODENGNN, we substitute the ODE function in TRS-ODEN with a GraphODE network. The GraphODE generative model is implemented with a single-layer GNN with hidden dimension 128. As in HODEN and TRS-ODEN, we compute each object's missing initial state via linear spline interpolation and the Leapfrog method for solving ODE. For all datasets, we use 0.5 as the coefficient for the reversal loss in (Huh et al., 2020), and 0.0002 as the learning rate under cosine scheduling.

LGODE

Our implementation follows (Huang et al., 2020) except we remove the Variational Autoencoder (VAE) from the initial state encoder. Instead of using the output from the encoder GNN as the mean and std of the VAE, we directly use it as the latent initial state. That is, the initial states are deterministic instead of being sampled from a distribution. We use the same architecture as in TREAT and train the model using an AdamW optimizer with a learning rate of 0.0001 across all datasets.

E Additional Experiments

E.1 Comparison of different solvers

We next show our model's sensitivity regarding solvers with different precisions. Specifically, we compare against Euler and Runge-Kutta (RK4) where the latter is a higher-precision solver. We show the comparison against LGODE and TREAT in Table 3.

We can firstly observe that TREAT consistently outperforms LGODE, which is our strongest baseline across different solvers and datasets, indicating the effectiveness of the proposed time-reversal symmetry loss. Secondly, we compute the improvement ratio as LGODE -TREAT LGODE . We can see that the improvement ratios get larger when using RK4 over Euler. This can be understood as our reversal loss is minimizing higher-order Tayler expansion terms (Theoreom 3.1) thus compensating numerical errors brought by ODE solvers.

E.2 Evaluation across observation ratios.

For LG-ODE and TREAT, the encoder computes the initial states from observed trajectories. To show models' sensitivity towards data sparsity, we randomly mask out 40 % and 80 % historical observations and compare model performance. As shown in Table 4, when changing the ratios from 80 % to 40% , we observe that TREAT has a smaller performance drop compared with LG-ODE, especially on the more complex Pendulum dataset (LG-ODE decreases 22.04 % while TREAT decreases 1.62 % ). This indicates that TREAT is less sensitive toward data sparsity.

F Discussion about Reversible Neural Networks

In literature, there is another line of research about building reversible neural networks (NNs). For example, (Chang et al., 2018) formulates three architectures for reversible neural networks to address

the stability issue and achieve arbitrary deep lengths, motivated by dynamical system modeling. (Liu et al., 2019) employs normalizing flow to create a generative model of graph structures. They all propose novel architectures to construct reversible NN where intermediate states across layer depths do not need to be stored, thus improving memory efficiency.

However, we'd like to clarify that reversible NNs (RevNet) do not resolve the time-reversal symmetry problem that we're studying. The core of RevNet is that input can be recovered from output via a reversible operation (which is another operator), similar as any linear operator W ( · ) have a reversed projector W -1 ( · ) . In the contrary, what we want to study is that the same operator can be used for both forward and backward prediction over time, and keep the trajectory the same. That being said, to generate the forward and backward trajectories, we are using the same g ( · ) , instead of g ( · ) , g -1 ( · ) respectively.

In summary, though both reversible NN and time-reversal symmetry share similar insights and intuition, they're talking about different things: reversible NNs make every operator g ( · ) having a g -1 ( · ) , while time-reversible assume the trajectory get from ˆ z fwd = g ( z ) and ˆ z bwd = -g ( z ) to be closer. Making g to be reversible cannot make the system to be time-reversible.

G Impact Statement

This paper presents work whose goal is to advance the field of Machine Learning. TREAT is trained upon physical simulation data (e.g., , spring and pendulum) and implemented by public libraries in PyTorch. During the modeling, we neither introduces any social/ethical bias nor amplify any bias in the data. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.

NeurIPS Paper Checklist

The checklist is designed to encourage best practices for responsible machine learning research, addressing issues of reproducibility, transparency, research ethics, and societal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and follow the (optional) supplemental material. The checklist does NOT count towards the page limit.

Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:

You should answer [Yes] , [No] , or [NA] .

[NA] means either that the question is Not Applicable for that particular paper or the relevant information is Not Available.

Please provide a short (1-2 sentence) justification right after your answer (even for NA).

The checklist answers are an integral part of your paper submission. They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include it (after eventual revisions) with the final version of your paper, and its final version will be published with the paper.

The reviewers of your paper will be asked to use the checklist as one of the factors in their evaluation. While "[Yes] " is generally preferable to "[No] ", it is perfectly acceptable to answer "[No] " provided a proper justification is given (e.g., "error bars are not reported because it would be too computationally expensive" or "we were unable to find the license for the dataset we used"). In general, answering "[No] " or "[NA] " is not grounds for rejection. While the questions are phrased in a binary way, we acknowledge that the true answer is often more nuanced, so please just use your best judgment and write a justification to elaborate. All supporting evidence can appear either in the main paper or the supplemental material, provided in appendix. If you answer [Yes] to a question, in the justification please point to the section(s) where related material for the question can be found.

IMPORTANT, please:

Delete this instruction block, but keep the section heading 'NeurIPS paper checklist" ,

Keep the checklist subsection headings, questions/answers and guidelines below.

Do not modify the questions and only use the provided macros for your answers .

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification:

[TODO]

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: Limitations are discussed in Appendix 6

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: Proofs are in Appendix A.3, A.2and A.4.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: The Datasets, Task Setup, Baselines describtion are in Sec. 4. Pseudo code for the implementation of the Time-Reversal Symmetry Loss is in Appendix A.1. More Model Details and Implementation Details are in Appendix D.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: The Datasets, Task Setup, Baselines description are in Sec. 4. Pseudo code for the implementation of the Time-Reversal Symmetry Loss is in Appendix A.1. More Dataset descriptions are in Appendix C. More Model Details and Implementation Details are in Appendix D.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: The Datasets, Task Setup, Baselines description are in Sec. 4. More Implementation Details are in Appendix D.2.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: In Appendix ??

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification:

[TODO]

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: [TODO]

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [Yes]

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification:

[TODO]

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification:

[TODO]

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: [TODO]

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification:

[TODO]

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification:

[TODO]

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper introduces a framework named TREAT (Time-Reversal Symmetry ODE) aimed at improving dynamical system modeling through a physics-informed approach. It incorporates Time-Reversal Symmetry (TRS) as a regularization term to enhance model precision. This approach is shown to preserve energy in conservative systems and provide strong inductive bias for non-conservative, reversible systems. Theoretical proofs are provided to demonstrate the numerical benefitd∂. The framework's effectiveness is validated through experiments on nine diverse datasets, showing improvement over existing models.\n\nSTRENGTHS:\nThe paper proposed the TRS as a regularization term for dynamical system modeling. this regularization intuitively makes sense, by augmenting the reverse trajectory as additional data for training.\n\nThe paper provides theories that TRS minimizes higher-order Taylor terms, enhancing modeling accuracy.\n\nThe paper conducted extensive experiments on nine datasets show that TREAT outperforms other neuralODE models.\n\nWEAKNESSES:\n1.\tThe paper only uses graphODE as backbone, from my understanding the TRS regularization can also fit with other neuralODE methods, it would be interesting to add experiments to show its compatibility (e.g. vanilla neuralODE). \n\n2.\tThe theoretical analysis in theorem 3.1 bounds the error in the embedding space z, however, the loss function in eqn 10 describes the error from original space y. it is worth mentioning the difference due to encoder/decoder error.\n\nQUESTIONS:\n1.\tFigure 1 shows the time reserve symmetry system is inclusive in the general classical mechanics. I’m confused regarding the TRS definition in physics and the numerical TRS used in this paper. In my understanding, any dynamical system with deterministic ode description can reverse trace the initial points by solving it from the end point numerically? therefore, the method should be generally applicable?\n\n2.\tIn figure 4a, do you have any insights why the TREAT does not do well compared with others with shorter prediction length.\n\nLIMITATIONS:\nn/a", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-04"}

comments: SUMMARY:
The paper proposes a novel Time-Reversal Symmetry (TRS) graph neural ODE, where the TRS is introduced as a soft regularisation term.

STRENGTHS:
The paper presents a novel method to combat numerical errors for graph neural ODEs.

WEAKNESSES:
The paper does not introduce early enough that the time-reversal symmetry (TRS) method has been applied to other DL methods such as the TRS-ODEN.

QUESTIONS:
Are there cases where the time-reversal symmetry expectation is violated?

A comparison to a numerical solution could serve as an additional benchmark to show how TRS decreases the numerical errors and improves the prediction. Also, why considering only MSE as a performance metric?

The code is anonymised, but it really needs cleaning and proper documentation to the very least docstrings and comments; I have not tested running the code.

The work will benefit from additional proofreading.

LIMITATIONS:
The authors do not provide a separate Limitations section, and do not discuss limitations of the work in the Conclusions section.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-08

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a regularization term to enforce Time-Reversal Symmetry (TRS) for modeling dynamical systems. The method helps preserve energies for conservative systems while serving as a strong inductive bias for non-conservative reversible systems. They also prove that TRS loss can universally improve modeling accuracy by minimizing higher-order Taylor terms in ODE integration. Numerical examples show that integrating the TRS loss within neural ordinary differential equation models demonstrates superior performance on diverse physical systems.\n\nSTRENGTHS:\nThis paper focuses on learning complex physical dynamics from data and incorporating physics priors into learning models, which is an important area. This paper proposes a regularization term to enforce Time-Reversal Symmetry (TRS), which is an important physics prior.\nAlso, the paper numerically shows that integrating the TRS loss within neural ordinary differential equation models demonstrates superior performance on diverse physical systems.\n\nWEAKNESSES:\nWhile the paper on time-reversal symmetry demonstrates several strengths, there are also some potential weaknesses that could be addressed:\n\n1.  Time-Reversal Symmetric ODE Network has been proposed (https://arxiv.org/pdf/2007.11362). This article is also cited in this paper. However, this paper only points out that two articles use different properties (Equation 5 and Lemma 2.1) to sign the regularization term, the difference in the final regularization term is not fully explained. \n\n2. The numerical experiments also demonstrated that the method proposed in this paper outperforms existing time-reversal symmetry regularization (https://arxiv.org/pdf/2007.11362). However, the reason why this method is better in this paper is not explained in detail.\n\n3. Some descriptions in the text are not clear enough, please refer to the question section below.\n\nQUESTIONS:\n1. In the proof of Lemma 2.1, it is required that R be an involution operator. However, such assumption is not stated in Lemma 2.1. Is there any evidence to prove that this assumption is reasonable？\n\n2. The definition of R does not appear in the paper. Can the method proposed in this paper hold for any R?\n\n3. How is regularization term 9 obtained based on Lemma 2, as R did not appear in the final regularization term 9？\n\n4. Does Theorem 3.1 hold for any untrained NN? \n\n5. The definition of $L_{reverse2}$ does not appear in the paper.  Can the condition $L_{reverse2}=L_{reverse}$ be satisfied?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-11"}

comments: SUMMARY:
This paper proposes a method to enhance neural ordinary differential equations (ODEs) by enforcing approximate Time-Reversal Symmetry. A self-supervised regularization term is introduced to align forward and backward trajectories predicted by a neural network, promoting energy conservation and stability in the system. Experiments are conducted on 9 datasets, comprising both real-world and simulated systems. The proposed approach outperforms other baseline methods, demonstrating its effectiveness in improving neural ODEs.

STRENGTHS:
* The idea of imposing time-reversal symmetry via a regularization term is novel and well-motivated.
* The implementation of the regularization term is simple yet the effect is visible.
* Overall the article is well written, and technical details are explained clearly.
* The results show significant improvement over other baseline models.

WEAKNESSES:
* It would be more interesting to see if the model works for other real-world examples/experiments.
* The article lacks a more comprehensive review on other  energy-preserving / time-reversal neural ODE solvers.

QUESTIONS:
Does the method work well with models other than GraphODE?

LIMITATIONS:
The authors briefly touched upon the possibility of incorporating properties in the spatial aspect such as translation and rotation equivariance.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-13