Equivariant Spatio-Temporal Attentive Graph Networks to Simulate Physical Dynamics
Authors: Liming Wu
Zhichao Hou
Jirui Yuan
Yu Rong
Wenbing Huang
Venue: NeurIPS 2023 poster (2023)
Keywords: Equivariance
Spatio-Temporal GNNs
Physical Dynamics

Abstract: Learning to represent and simulate the dynamics of physical systems is a crucial yet challenging task. Existing equivariant Graph Neural Network (GNN) based methods have encapsulated the symmetry of physics, \emph{e.g.}, translations, rotations, etc, leading to better generalization ability. Nevertheless, their frame-to-frame formulation of the task overlooks the non-Markov property mainly incurred by unobserved dynamics in the environment. In this paper, we reformulate dynamics simulation as a spatio-temporal prediction task, by employing the trajectory in the past period to recover the Non-Markovian interactions. We propose Equivariant Spatio-Temporal Attentive Graph Networks (ESTAG), an equivariant version of spatio-temporal GNNs, to fulfil our purpose. At its core, we design a novel Equivariant Discrete Fourier Transform (EDFT) to extract periodic patterns from the history frames, and then construct an Equivariant Spatial Module (ESM) to accomplish spatial message passing, and an Equivariant Temporal Module (ETM) with the forward attention and equivariant pooling mechanisms to aggregate temporal message. We evaluate our model on three real datasets corresponding to the molecular-, protein- and macro-level. Experimental results verify the effectiveness of ESTAG compared to typical spatio-temporal GNNs and equivariant GNNs.

Equivariant Spatio-Temporal Attentive Graph Networks to Simulate Physical Dynamics
Liming Wu Zhichao Hou Jirui Yuan", Yu Rong", Wenbing Huang 1.2f
Gaoling School of Artificial Intelligence, Renmin University of China
Beijing Key Laboratory of Big Data Management and Analysis Methods, Beijing, China DDepartment of Computer Science, North Carolina State University Institute for AI Industry Research (AIR), Tsinghua University STencent AI Lab zhou4@ncsuedu
{manl hww.,,w.ewwwnbigi.. Savenbing)@ruc.eeee yu.rrong@hotmmil.com yuanj yuanjirui jinijiruueee www.iaii..tsinnuu.et......
Abstract
Learning to represent and simulate the dynamics of physical systems is crucial yet challenging task. Existing equivariant Graph Neural Network (GNN) based methods have encapsulated the symmetry physics, e,g., translations, rotations, etc, leading to better generalization ability. Nevertheless, their frame-to-frame formulation of the task overlooks the non-Markov property mainly incurred by unobserved dynamics in the environment. In this paper, we reformulate dynamics simulation as spatio-temporal prediction task, by employing the trajectory the past period to recover the Non-Markovian interactions. We propose Equivari- ant Spatio-Temporal Attentive Graph Networks (ESTAG), an equivariant version of spatio-temporal GNNs, to fulfill our purpose. At its core, we design a novel Equivariant Discrete Fourier Transform (EDFT) to extract periodic patterns from the history frames, and then construct Equivariant Spatial Module (ESM) accomplish spatial message passing, and an Equivariant Temporal Module (ETM) with the forward attention and equivariant pooling mechanisms to aggregate tem- poral message. We evaluate our model on three real datasets corresponding to the molecular-, protein- and macro-level. Experimental results verify the effectiveness ESTAG compared to typical spatio-temporal GNNs and equivariant GNNs.
Introduction
has been a goal to represent and simulate the dynamics physical systems by making use of
machine learning techniques [37 g [2]. The related studies, once pushed forward, have great
potential to facilitate a variety of downstream scientific tasks including Molecular Dynamic (MD)
simulation [19,, protein structure prediction [I] virtual screening of drugs and materials (30] model-
based robot planning/control (35, and many others.
Plenty of solutions have been proposed, amongst which the usage of Graph Neural Networks (GNNs) [4 becomes one the most desirable directions. GNNs naturally model particles or unit elements as nodes, physical relations edges, and the latent interactions the message passing thereon. More recently, line of researches 38 with generalizing GNNs to fit the symmetry of our physical world. These works, also known as the geometric input GNNs
32 20 have been concerned
cquivariant GNNs [88, ensure that results in the output transformed in the same way. By handcrafting such Euclidean equivariance, the
"qqual contribution. Corresponding author.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

GT-1
X(T)
Go
GT-1
X(T)
(a) Previous Setting
(b) Our Setting
Figure 1: Comparison of the problem setting between previous methods and our paper. Here, we choose the dynamics of the Aspirin molecular with time lag as for illustration.
models are well predictable to scenarios under arbitrary coordinate systems, giving rise to enhanced generalization ability..
In spite of the fruitful progress, existing methods overlook vital point: the observed physical dynamics are almost non-Markovian. In previous methods, they usually take as input the conformation of a system
single temporal frame and predict as output the future conformation after a fixed time interval, forming a frame-to-frame frame forecasting problem. Under the Markovian assumption, this setting is pardonable as future frames are independent of all other past frames given the input one. Nevertheless, the Markovian assumption is rather unrealistic when there are other unobserved objects interacting with the system we are simulating 440.. For instance, when we consider simulating the dynamics of protein that is interacting with unobserved solvent (such water), the Markovian property no longer holds; in other words, even conditional on the current frame, the future dynamics of the protein depends on the current state of the solvent which, however, is influenced by the past states of the protein itself, owing to the interaction between the protein and the solvent. The improper Markovian assumption makes current works immature in dynamics modeling.
To relieve from the Markovian assumption, this paper proposes to employ the states in the past period reflect the latent and unobserved dynamics. In principle, we recover the non-Markovian behavior (e.g. interacting with solvent) if the past period is sufficiently long. We collect period of past system states spatio-temporal graphs, and utilize them the input to formulate spatio- temporal prediction task, other than the frame-to to to-frame problem as usual (see Figure D This motivates to leverage existing Epatio-Tempone which, unfortunately, are unable to conform to the aforementioned Euclidean symmetry and the underlying physical laws. Hence, the equivariant version of STGNNs is no doubt in demand. Another point is that periodic motions are frequently observed in typical physical systems [3. For example, the Aspirin molecular exhibits clear periodic thermal vibration when binding target protein. Under the spatio-temport setting, we are able to model periodicity of dynamics, which, however, is less investigated before.
GNNs (STGNNs) 44 to fulfill our purpose,
Our contributions are summarized as follows:
We reveal the non-Markov behavior in physical dynamics simulation by developing equivari- ant spatio-temporal graph models. The proposed model dubbed Equivariant Spatio-Tempore Attentive Graph network (ESTAG) conforms to Euclidean symmetry and alleviates the limi- tation of the Markovian assumption.
We design a novel Equivariant Discrete Fourier Transform (EDFT) to extract periodic features from the dynamics, and then construct Equivariant Spatial Module (ESM), and an Equivariant Temporal Module (ETM) with forward attention and equivariant pooling, to process spatial and temporal message passing, respectively.
The effectiveness of ESTAG is verified on three real datasets corresponding to the molecular-, protein- and macro-level. We prove that, involving both temporal memory and equivariance are advantageous compared to typical STGNNs and equivariant GNNs with adding trivial spatio-temporal aggregation.
2
Related Work
GNNs for Physical Dynamics Modeling Graph Neural Networks (GNNs) have shown great potential in physical dynamics modeling. IN 22, NRI 233 and HRN 29 are series of works to learn physical

---

Dwilte
15g
0g
Go
G
KAA
(T)
307)
Input Trajecerry
Equivarinat DFT
Equitariant Spotial Madale
Equivariant Temporal Module
Output Pusition
Figure 2: Schematic overview of ESTAG. After inputting historical graph trajectories Go...., Equivariant Discrete Fourier Transform (EDFT) extracts equivariant frequency features f from the trajectory. We process them into the invariant node-wise feature and adjacency matrix A to be adopted for the next stage. Then we stack Equivariant Spatial Module (ESM) and Equivariant Temporal Module (ETM) alternatively for L times explore spatial and temporal dependencies. After the equivariant temporal pooling layer, we obtain the estimated position / (T').
system interaction and evolution. Considering the energy conservation and incorporating physical prior knowledge into GNNs, HNN [13), HOGN [31] leverage ODE and Hamiltonian mechanics to capture the interactions in the systems. However, all of the above-mentioned models don't 't take into account the underlying symmetries of a system. In order to introduce Euclidean equivariance, Tensor-Field networks (TFN) 38 and Trans-Trrn [[1] equip filters with rotation eguivariance by irreducible representation of the SO(3) group. LieTrrasformer [n] LieConv p leverage the Lie group to enforce equivariance. Besides, EGNN [32] utilized a simpler E(n)-equivariant framework which can achieve competitive results without computationally expensive information. Based on EGNN, GMN (20) further proposes an equívariant and constraint-aware architecture by making use of forward kinematics information in physical system. Nevertheless, all of these methods ignore the natural spatiotemporal patterns of physical dynamics and model it as a Frame-to-framt forecasting problem.
22 aim to capture the spatial and temporal dependency simultaneously and are widely investigated various applications like traffic forecasting and human recognition. DCRNN (26) and GaAN 45) are RNN-based methods which filter inputs and hidden states passed to recurrent unit using graph convolutions. However, RNN- based approaches are time-consuming and may suffer from gradient explosion/vanishing problem. CNN-based approaches, such as STGCN 44 and ST-GCN 43 interleave ID-CNN layers with graph convolutional layers tackle spatial temporal graphs in a non-recursive manner. Besides, attention mechanism, an important technique STGCNs, is employed by GaAN [45), AGL-STAN (36) and ASTGCN [16] to learn dynamic dependencies in both space and time domain. The aforementioned approaches are targeted to the applications in 2D graph scenarios such as traffic networks and human skeleton graphs, and may not be well applicable to 3D physical systems in which geometric equivariance is really important property.
Spatio-temporal graph neural networks STGNNs
Equivariant spatio-temporal graph neural networks There few previous works designing spatio-temporal GNNs while maintaining equivariance. Particularly, by using GRU [6] to record the memory of past frames, LoCS (24) additionally incorporates rotation-invariance to improve the model's generalization ability. Different from the recurrent update mechanism used in LoCs, EqMotion [42 distills the history trajectories of each node into miuliidiimnnion vector, by which the spatio-temporal graph is compressed as a spatial graph, then it designs an equivariant module and an interaction reasoning module predict future frames. However, both of LoCS and EqMotion are still defective in exploring the interactions among history trajectories, while in this paper we propose a transformer-alike architecture fully leverage the spatio-temporal interactions based on equivariant attentions.
Notations and Task Definition
The dynamics of physical objects (such as molecules) can formulated with the notion of spatiotem- poral graphs, as shown in Figure (left). In particular, spatiotemporal graph of node number N and
3

---

temporal length T' is denoted as {G; (),, Here, the nodes V2 are shared across time, but each node is assigned with different scalar feature (t) ER, position vector क; (t) R at differ- ent temporal frame I; the edges & are associated with an identical adjacency matrix A E x whose element Aij E Ihmm defines the edge feature between node i and j. We henceforth denote by the matrices (t) E R N and (t) E R N the collection of all nodes in G. Here for simplicity, we specify the time lag as At 1 which could be valued more than 1 in practice. In general. the t-th frame corresponds to time T tåt.
Task Definition This paper is interested in predicting the physical state, particularly the position of each node frame T given the historical graph series {Ge ) In form, we learn the function Φi
(1), X (L), 2 2 X (T).
(I)
Although previous approaches such as EGNN dynamics modeling, they neglect the application of spatiotemporal patterns and only accomplish frame-to-frame prediction, where the input of Φ is reduced to a single frame, e.g., GT in Eq. |
and GMN (0] also claim to tackle physical
Equivariance A crucial constraint on physical dynamics is that the function & should meet the symmetry of our 3D world. In other words, for any translation/rotaiio//eeeeetciii g in the group H(3), 6 satisfies:
((( X(t), A)17 = (T),
(2)
where the group action is instantiated as g X(I) X(1) + b for translation b E R and
g x())= = Ox for rotation/reflection O R333
4 Our approach: ESTAG
Discovering informative spatiotemporal patterns from the input graph series is vital to physical dynamics modeling. In this section, we introduce ESTAG that pursues this goal in a considerate way: We first extract the frequency of each node's trajectory by EDFT ($ 4.1). which captures the node-wise temporal dynamics in global sense and returns important features for the next stage; We then separately characterize the spatial dependency among the nodes of each input graph G, via ESM ($4.2); We finally unveil the temporal dynamics of each node through the attention-based mechanism ETM and output the estimated position of each node in Gr after an equivariant temporal pooling layer ($ 4.3) The overall architecture shown in Figure 2
4.1
Equivariant Discrete Fourier Transform (EDFT)
The Fourier Transform (FT) gives us insight into the wave frequencies contained in the input signal that is usually periodic. With the extracted frequencies, are able view the global behavior of each node i in different frequency domains. Conventional multidimensional FT employs distinct Fourier bases for different input dimensions of the original signals. Here, ensure equivariance, we first translate the signals by the mean position and then adopt the same basis over the spatial dimension. To be specific, we compute equivariant DFT as follows:
f(())
2 kt (z,(t) (t()))
2=
T is the frequency index, Z(t) the average position of all nodes in the t-th frame Gr. and the output (A())) E C3 is complex. The frequencies calculated by Eq. 3 are then utilized to formulate two crucial quantities: the frequency cross-correlation Ai E IR? between node and j. and the frequency amplitude c, Rd of node i.
where, i' is the imaginary unit, 1: 0,1, .......T---11
In signal processing, cross correlaton measures the similarity of two functions f1 and f2. satisfies F(J f2) F(A)- F(f2), where F and * denote the FT and the cross-correlation operator, respectively, and F indicates the complex conjugate of F. Borrowing this idea. we compute the cross-correlation in the frequency domain by Eq. 3 aas
Ak) = wx (h.)wa (h (KF())))))) F(.))
(4)

---

where ( defines the complex inner product. Notably, we have added two learnable parameters wy (h;) and wn. wh(h) dependent on node features, which act like spectral filters of the k-th frequency and enable us to select related frequency for the prediction. In the next subsection, we will apply A;j as the edge feature to capture the relationship between different nodes. We use Aspirin as an example and visualize the A in Figure
Aspirin
400
Figure 3: Visualization of cross-correlation A Aspirin. EDFT not only identify strongly- connected nodes (e.g. Node and Node 11), but also discover latent relationship between two nodes which are disconnected yet may have similar structures functions (e.g. Node 8 and Node 10).
We further compute for node i the amplitude the frequency f((k) along with the parameter wx (h.):
c (k) wx (h,)), (b.
(5)
This term will be used in the update of the hidden features in the next subsection.
promising property of Eq. is that it is translation invariant and rotation/reflection equivariant. Therefore, both Ai and c; are ((3)-invarian)) which will facilitate the design of following modules.
4.2 Equivariant Spatial Module (ESM)
For each graph G. our ESM is proposed to encode its spatial geometry through equivariant message passing. ESM is built upon EGNN (3) which is prevailing kind of equivariant GNNs, but has subtly involved the FT features from the last subsection for enhanced performance beyond EGNN.
The 1-th layer message passing in ESM is as below:
10 (t), IE (D1P, A.)) mij Qm (n (t), h
(6)
(44) (t) = (n'(() + ob
h1 (t), (t), c.. Σ Emi mis
(7)
(t( = [W(i) Σ
1 (f)))(
(8)
JEN(()
((+))()) (t) E (()() (t) + (t).
(9)
where, om computes the message mi, from node j to i, dr updates the hidden representation h., & returns one-dimensional scalar for the update of (t), and all the above functions are Multi- Layer Perceptrons (MLPs); (t) = (t) (t) ) is the relative position and (()) denotes the neighborhoods of node i.
Notably, we leverage the cross-correlation A.j as the edge feature in Eq 6 to evaluate the connection between node i and over the global temporal window, since it is computed from the entire trajectory. We also make use of e; as the input of the update in Eq. The benefit of considering these two terms will be ablated in experiments.
5

---

4.3 Equivariant Temporal Module (ETM)
Forward Temporal Attention Inspired by the great success of Transformer [39 in sequence
modeling, we develop ETM that describes the seff-correspondence of each node's s trajectory based on the foward attention mechanism, and more importantly, in an E(3)-qquivariant way.
In detail, each layer of ETM conducts the following process:
a (1s) exp(: (1) (1)TA K (s))
(10)
E.= exp 9 11)
k (5))
h ((+))()) h1 (t) + Σ
a
0 (ts)
(11)
==0
E 0 (t) (t)= क (t) + Σ a (ts) (t
(12)
where, Ctts is the attention weight between time and 8, computed by the query qe and key ky; the hidden feature (t) is updated as weighted combination of the value Usi the position vector E; (4) is derived from weighted combination one-dimensional scalar (1,())) multiplied with
(1) the temporal displace vector E. (ts) x (t) (!)(ss). Specifically, q (1) (t) 0a h ()).
k (1) dx (e()(() (t) and 0 (t) = 6. h (t) are B(3)-invariant functions. Notably, we
derive particle's 's next position in a forward-looking way, keep physical rationality as the derivation of current state should not be dependent on future positions.
Equivariant Temporal Pooling We alternate one-layer ESM and one-layer ETM over L layers, and finally attain the updated coordinates x!!²) E R for each node Then the predicted coordinates at time 7 is given by the following equivariant linear pooling:
क (T) X,w + (T 1),
(13)
where the parameter E R (T-1) consists of learnable weights, and X; = (0( (L)
x; (T
1), a ( (1)
x (T'- 1), .., (T - 1)] is translated by ( (T' 1) to allow
translation invariance.
We train ESTAG end-to-end via the mean squared error (MSE) loss:
L (T.(T) क (T)
(14)
By the design of EDFT, ESM and ETM, we have the following property of our model ESTAG.
Theorem 4.1. We denote ESTAG as X(T) =6 (((H())).) X(1,, ())T))) then & E(3)- eguivariant.
Proof. See Appendix A
Although we mainly exploit EGNN as the backbone (particularly in ESM), our framework is general and can be easily extended other equivariant GNNs, such as GMN 220, the multi-channel version of EGNN. In general, the extended models deal with multi-channel coordinate Z E mx instead of ē E 3 The most significant feature of these models is to replace the invariant scalar | in the formulations of the message mij (Eq. 6) with the term Z Z. It is easily to prove that this term is equivariant to any orthogonal matrix O, i.c., (Ož) (0Z) = zT Z, 44 E IR3x3 O I. Besides, can be reduced to invariant scalar || when m 1. Empirically, we add the normalization term in order to achieve more stable performance:
zTz__ //' where || || is the Frobenius norm. In the above derivation, we only display the formulation on EGNN, the details of multi-channe ESTAG are shown in Appendix C.I

---

Table 1:
Prediction error (x10->) on MD17 dataset. Results averaged across 3 runs. We not display the standard deviation due its small value.
ASPIRIN BENZENE ETHANOL MALONALDEHYYE NAPHTHALENE SALICYLIC TOLUENE URACIL
PT-8 15.579
4.457
4.332
13.206
8.958
12.256
6.818
10.269
PT-m
9.058
2.536
2.688
6.749
6.918
8.122
5.622
7.257
PT-f
0.715
0.114
0,456
0.596
0.737
0.688
0.688
0.674
EGNN-8
12.056
3.290
2.354
10.635
4.871
8.733
3.154
6.815
EGNN-m
6.237
1.882
1.532
4.842
3.791
4.623
2.516
3.606
EGNN-E
0.625
0.112
0.416
0.513
0.614
0.598
0.577
0.568
ST_ TFN
0.719
0.122
0.432
0.569
0.688
0.684
0.628
0.669
ST GNN
1.014
0.210
0.487
0.664
0.769
0.789
0.713
0.680
ST SE(3YTR
0.669
0.119
0.428
0.550
0.625
0,630
0.591
0.597
ST EGNN
0.735
0.163
0.245
0.427
0.745
0.687
0.553
0.445
EOMOTION
0.721
0.156
0.476
0.600
0.747
0.697
0.691
0.681
STGCN
0.715
0.106
0,456
0.596
0.736
0,682
0.687
0.673
AGL-STAN
0.719
0.106
0,459
0.596
0.601
0,452
0,683
0.515
ESTAG
0.063
0.003
0.099
0.101
0.068
0.047
0.079
0.066
Experiments
Datasets. To verify the superiority of the proposed model. we evaluate our model on three real world datasets: molecular-leve. MD17 (5), 2) protein-level: AdK equilibrium trajectory dataset [34] and 3) macro-level: CMU Motion Capture Databse 77] These datasets involve several continuous long trajectories. Note that all the three datasets contain unobserved dynamics or factors and thus conform to the non-Markovian setting. In particular, the external temperature and pressure are unknown on MD17, the dynamics of water and ions is unobserved on AdK, and the states of the environment not provided Motion Capture. The original datasets are composed of long trajectories. We randomize the start point and extract the following + points with the interval At. We take the first T timestamps as previous observations and last timestamp the future position label.
We regard the previous observation at start/mid/terminal (s/m/t) timepoint as the estimated position at timestamp T directly, 2) EGNN [33 utilizes simple yet efficient framework which transforms the 3D vectors into invariant scalars. We provide EGNN with only one previous position at s/m/t timepoint to predict the future position in frame-to- to-faet manner. 3) STGCN 44 is a spatio-temporal GNN that adopts a "sandwich" structure with two gated sequential convolution layers and on spatial graph convolution layer between. modify its default settings by predicting the residual coordinate between time T and T, since directly predicting the exact coordinate at time T yields much worse performance.4) AGL-STAN (36) leverages adaptive graph learning and self-attention for comprehensive representation of intra-temporal dependencies and inter-spatial interactions. We modify AGL-STAN's setting in the same way as we do with STGCN. 5) Typical GNNs with trivial spatio-temporal aggregation:
Baselines. We compare the performance of ESTAG with several baselines: 1)
implement GNN [13), and other equivariant models EGNN, TFN [38], and (E(3)-Transform) [T for each temporal frame the historical trajectory and then estimate the future position the weighted sum of all past frames, where the weights learnable. All models are denoted with a prefix "ST" and we initialize their node features along with temporal positional encoding. 6) EqMotion 42] is one of equivariant spatio-temporal GNNs, which leverages the temporal information by fusing them for the model's initialization.
5.1 Molecular-level: MD17
Implementation details. MD17 dataset includes the trajectories of small molecules generated by MD simulation. We use the atomic number as the time-independent input node feature (0) The two atoms are 1-hop neighbor if their distance is less than the threshold À and we consider two types of neighbors (i.e. 1-hop neighbor and 2-hop neighbor). Other settings including the hyrre-earsmeters introduced in Appendix D.I
7

---

Ground Truth
MSE+0.065
MSE-0.001
MSE-0.040
MSE-0.060
MSE-0.086
MSE-0
ESTAS
MSE-0.997
MSE-0.135
MSE-0.203
MSE-0.729
MSS-0U654
MSE-0.6E 654
VSE-0,670
ST,EGNN
ASPIRIN
DENZENE
ETHANOL
WALONALDEHYDE
MAPHTHALENE
SALICYLIC
TOLUENE
URACIL
Figure 4: PyMol visualization of the predicted molccules by our ESTAG and ST EGNN, where the MSE (x 10 3) with respect to the ground truth is also shown. As expected, the predicted instances by ESTAG exhibit much smaller MSE than ST EGNN, although the difference is not easy to visualize in some cases. For those obviously mispredicted regions of ST EGNN, highlight them with red rectangles. It is observed that ST EGNN occasionally outputs isolated atoms, which could caused by violation of the bond length tolerance in PyMol.
Results. Table shows the average MSE all models on 8 molecules. We have some interesting observations follows: 1) ESTAG exceeds other models in all cases by significant margin, the proposed ideas. 2) From the Pt-s/m/t,, we observe that
supporting the general effectiveness the point closer to the future point has less prediction error. EGNN-s m//t only takes one frame (s/m/t) as input, and attains slight improvement relative to Pt-s/m/t. 3) Compared with EGNN-I, ST SY_ENNN, although equipped with trivial spatio-temporal aggregation, is unable to obtain consistent improvement, which indicates that how to unveil temporal dynamics appropriately is beyond triviality on this dataset. 4. The con-equivariant methods particularly ST L_NN perform unsatisfactorily in most cases, implying that equivariance is an important property when modeling 3D structures.
Visualization. We have provided some visualizations of the predicted molecules by using the PyMol toolkit. Figure shows that the prdicted MSEs of our model are much lower than ST EGNN and our predicted molecules
_T_
closer the ground-trun -ounn molecules. We have also displayed the learned attentions (Eq. (0) and the temporal weight (Eq. [3). where we find meaningful patterns of the temporal correlation. For clearer comparison, display the molecule URACIL in 3D coordinate system, which can found in Appendix D.5
5.2 Protein-level: Protein Dynamics
Implementation details. We evaluate our model on the AdK cquilibrium trajectory dataset [34] via MDAnalysis toolkit [1]. In order to reduce the data scale, we utilize MDAnalysis to locate the backbone atoms (Ca, C, N, O) of the residues and then regard the residues other than atoms in protein the nodes with 4-channel geometric features. We use the atomic number of four backbone atoms the time-independent input node feature (0)), We connect two atoms via an edge their distance is less than a threshold parameters are introduced in Appendix D.2, Slightly dif- ferent from the model on MD17 dataset, we generalize the single-channe ESTAG into multi-channel version which presented in detail in Appendix [] We do conduct EqMotion. TFN and SE(3)-TR used in the last experiment since it is non-trivial to modify them into multi-channel modeling. The ofate-of-the-ar method GMN 20 is imple- mented by further adding the weighed temporal pooling similar to ST EGNN.
Table 2: Prediction error and training time on Protein dataset. Results aver- aged across runs.
METHOD
MSE
TIME(S)
PT-s
3.260
Pr-m
3.302
Pr-t
2.022
Other settings including the hyper-
EGNN-s EGNN- mm
3.254
1.062
3.278
1.088
EGNN-t
1.983
1.069
ST GNN
1.871
2.769
ST GMN
1.526
4.705
ST EGNN
1.543
4.705
STGCN
1.578
1.840
AGL-STAN
1.671
1.478
ESTAG
1.471
6.876

---

Results. The predicted MSEs are displayed in T Generally, the spatio- tempooal models are better than Pt-s/m/t and EGNN-s/ S/GNNss/m/tt and suggests that applying spatio-temporal clues on this dataset is crucial, particularly given that the protein trajectories are generated under the interactions with external molecules such as water and ions. The equivariant models always outperform the non- equivariant counterparts (for example, ST T_ENNNN vs. ST T_GNN. Overall, our model ESTAG achieves the best performance owing to its elaboration of equivariant spatio-temporal modeling. Additionally, we report the training time averaged over epochs in Table 2 It shows that the computation overhead of ESTAG over its backbone EGNN acceptable given its remarkable performance enhancement, It is expected that the superiority of ESTAG protein dataset is not obvious that on MD17, owing to various kinds of physical interactions between different amino acids, let along each amino acid is composed of certain number of atoms, which makes the dynamics of protein much more complicated than small molecules.
5.3 Macro-level: Motion Capture
Implementation details. We finally adopt CMU Motion Capture Database 7 to evaluate our model. CMU Motion Capture Database involves the trajecto- ries of human motion under several scenarios and we focus walking motion (subject #35) and basketball motion (subject #102, only take trajectories whose length is greater than 170). The input feature of all the joints (nodes) ( are all 1s. The two joints are 1 1-hop neighbor if they connected naturally and we consider two types neighbors (i.e. 1-hop neighbor and 2-hop neighbor). Other settings including the hyper-parameters are introduced in Appendix [D3] Notably, the input of EqMotion only contain node coordinates, as the same as our method and other baselines. We find that EqMotion performs much worse by directly predicting the absolute coordinates. We then modify EqMotion to predict the relative co- ordinates across two adjacent frames and perform zero-mean normalization of node coordinates, for further improvement.
Table 3: Prediction error (x (x10-1) on Motion dataset. Results averaged across runs.
METHOD
WALK
BASKETBALL
PT-s
329.474
886.023
PT-m
127.152
413.306
PT-/
3.831
15.878
EGNN-s
63.540
749.486
EGNN-m
32.016
335.002
EGNN-t
0.786
12.492
ST GNN
0.441
15.336
ST TFN
0.597
13.709
ST ST_SE(3)
0.236
13.851
TG_NNNN
0.538
13.199
EQMOTION
1.011
4.893
STGCN
0.062
4.919
AGL-STAN
0.037
5.734
ESTAG
0.040
0.746
Results. Table 3 summarizes the results of all models on the Motion dataset. The spatio-temporal models are better than Pt-s/m/t and EGNN-s/m//t which again implies the necessity of taking the spatio-temporal history into account. Unexpectedly, the non-equivariant models are even superior the equivariant baselines in walking motion, by, for instance, comparing AGL-STAN with ESTAG. We conjure that the samples of this dataset are usually collected the same orientation, which potentially subdues the effect of rotation equivariance. It is thus not surprising that GNN even outperforms EGNN since GNN involves more flexible form of message passing. But for basketball motion which is more complicated to simulate, ESTAG yields much lower MSE. We also notice that the attention-be based models including our ESTAG, ST ST_ SE(3)-Tr, and AGL-STAN perform promisingly, which probably due to the advantage of using attention to discovery temporal interactions within the trajectories.
5.4 Ablation Studies
Here we conduct several ablation experiments on MD17 to inspect how our proposed components contribute the overall performance and the results are shown Table
1) Without EDFT. We replace the FT-based edge features with the predefined edge features based on the connecting atom types and the distance between them. The simplified model encounters an average of increase in MSE, showcasing the effectiveness of the FT-based feature in modeling the spatial relation in graphs.
2) Without attention. We remove the ETM of ESTAG and observe slight detriment in the model performance, which demonstrates that attention mechanism well capture the temporal dynamics.

---

3) Without equivariance. We construct enuivarivaian ppatio-temporal attentive framework based on vanilla GNN. ESTAG performs much better than the non-equivariant STAG, indicating that Euclidean equivariance is a crucial property when designing model on geometric graph.
3) Without equivariance. We construct non-equivariant spatio-temporal attentive framework based on vanilla GNN. ESTAG performs much better than the non-equivariant STAG, indicating that Euclidean equivariance is a crucial property when designing model on geometric graph.
4) The impact of the number of layers L. We investigate effect of the number of layers on Ethanol dataset. We vary T from 1, 2, 4, 5, 6 and present the results in Table 3 Considering the accuracy and efficiency simultaneously, we choose == 2 for the ESTAG.
More ablation studies will be shown in Appendix
Table 4: Ablation studies (x (x10-3) on MD17 dataset. Results averaged across 3 runs.
Aspirin Benzene Ethanol Malonaldehyde Naphthalene Salicylic Toluene Uracil
ESTAG
0.063
0.003
0.099
0.101
0.068
0.047
0.079
0.066
w/o EDFT
0.079
0.003
0.108
0.148
0.104
0.145
0.102
0.063
w/o Attention
0.087
0.004
0.104
0.112
0.129
0.095
0.097
0.078
w/o Equivariance
0.762
0.114
0.458
0.604
0.738
0.698
0.690
0.680
w/o Temporal
0.084
0.003
0.111
0.139
0.141
0.098
0.153
0.071
Table 5: MSE on Ethanol w.r.t. the number of layers L.
2
5
MSE (x10 4) 1.25 0.990 1.096 1.022 1.042 1.028
6 Conclusion
In this paper, we propose ESTAG, an end-to-end equivariant architecture for physical dynamics modeling. ESTAG first extracts frequency features via novel Equivariant Discrete Fourier Transform (EDFT), and then leverages Equivariant Spatial Module (ESM) and an attentive Equivariant Temporal Module (ETM) refine the coordinate in space and time domain alternatively. Comprehensive experiments over multiple tasks verify the superiority of ESTAG from molecular-level, protein-level, and macro-level. Necessary ablations, visualizations, and analyses are also provided support the validity of our design as well as the generalization of our method. One potential limitation of our model is that we only enforce the E(3) symmetry while other inductive bias like the energy conservation law is also required in physical scenarios.
In the future, we will continue extending benchmark with more tasks and datasets and evaluate more baselines to validate the effectiveness our model. is also promising to extend our model to multi-scale GNN (like SGNN [17], REMuS-GNN (27). BSMS-GNN [4] and MS-MGN (10), which is useful particularly for industrial-level applications involving huge graphs. Besides, it is valuable to employ our simulation method as a basic block for other applications such as drug discovery, material design, robotic control, etc.
7 Acknowledgement
This work was jointly supported by the following projects: the Scientific Innovation 2030 Major Project for New Generation of AI under Grant NO. 2020AAA0107300, Ministry of Science and Technology of the People's Republic of China; the National Natural Science Foundation of China (62006137); Beijing Nova Program (20230484278); the Fundamental Research Funds for the Central Universities, and the Research Funds of Renmin University of China (23XNKJ19); Tencent AI Lab Rhino-Binn Bird Focused Research Program (RBFR2023005); Ant Group through CCF-Ant Research Fund CCCF-AFSG RF20220204); Public Computing Cloud, Renmin University of China.
10

---

References
[1] Bissan Al-Lazikani, Joon Jung, Zhexin Xiang, and Barry Honig. Protein structure prediction. Current opinion in chemical biology. 5(1):11-56, 2001.
[2] Peter Battaglia, Razvan Pascanu, Matthew Lai, Danilo Jimenez Rezende, et al. Interaction networks for learning about objects, relations and physics. Advances in neural information processing systems, 29, 2016.
[3] Hendrik W Broer, George Huitema, and Mikhail B Sevryuk. Quasi-periodic motions in families of dynamical systems: order amidst chaos. Springer, 2009.
[4] Yadi Cao, Menglei Chai, Minchen Li, and Chenfanfu Jiang. Efficient learning of mesh- bessh- physical simulation with bi-stride multi-scale graph neural network. 2023.
[5] Stefan Chmiela, Alexandre Tkatchenko, Huziel E. Sauceda, Igor Poltavsky, Kristof T. Schütt, and Klaus-Robert Müller. Machine learning of accurate energy-conserving molecular force fields. Science Advances, 3(5):e 2017.
[6] Junyoung Chung, Caglar Gulcehre, KyungHyun Cho, and Yoshua Bengio. Empirical evaluation of gated recurrent neural networks on sequence modeling. arXiv preprint 151.:554
2014.
[7] CMU. Carnegie-mellon motion capture database, 2003. URL hwttt:hhooo//pp..... cmu cdu.
[8] Mingyu Ding, Zhenfang Chen, Tao Du, Ping Luo, Josh Tenenbaum, and Chuang Gan. Dynamic visual reasoning by learning differentiable physics models from video and language. Advances In Neural Information Processing Systems, 34:887-899, 2021.
[9] Mare Finzi, Samuel Stanton, Pavel Izmailov, and Andrew Gordon Wilson. Generalizing convolutional neural networks for equivariance to lie groups arbitrary continuous data. In International Conference on Machine Learning, pages 3165-3176. PMLR, 2020.
[10] Meire Fortunato, Tobias Pfaff, Peter Wirnsberger, Alexander Pritzel, and Peter Battaglia. Multiscale meshgraphnets. arXiv preprint arXiv:2210000612, 2022.
[11] Fabian Fuchs, Daniel Worrall, Volker Fischer, and Max Welling. Se (3)-transformers: 3d into-translation equivariant attention networks. Advances in Neural Information Processing Systems, 33:1970-199 1981, 2020.
[12] Chuang Gan, Jeremy Schwartz, Seth Alter, Martin Schrimpf, James Traer, Julian De Freitas, Jonas Kubilius, Abhishek Bhandwaldar, Nick Haber, Megumi Sano, et al. Threedworld: A platform for interactive multi-modal physical simulation. arXiv preprint arXiv:2007, 04954, 2020.
[13] Justin Gilmer, Samuel S Schoenholz, Patrick Riley, Oriol Vinyals, and George E Dahl. Neural message passing for quantum chemistry. In International conference on machine learning, pages 1263-1272. PMLR, 2017.
[14] Richard Gowers, Max Linke, Jonathan Barnoud, Tyler JE Reddy, Manuel Melo, Sean Seyler, Jan Domanski, David L. Dotson, Sébastien Buchoux, Ian M Kenney, et al. Mdanalysis: a python package for the rapid analysis of molecular dynamics simulations. In Proceedings of the 15th python in science conference, volume 98, page 105. SciPy Austin, TX. 2016.
[15] Samuel Greydanus, Misko Dzamba, and Jason Yosinski. Hamiltonian neural networks. Advances in neural information processing systems, 32, 2019.
[16] Shengnan Guo, Youfang Lin, Ning Feng, Chao Song, and Huaiyu Wan. Attention based spatial- temporal graph convolutional networks for traffic flow forecasting. In Proceedings of the AAAI conference on artificial intelligence, volume 33, pages 922-929, 2019.
[17] Jiaqi Han, Wenbing Huang, Hengbo Ma, Jiachen Li, Josh Tenenbaum, and Chuang Gan. Learning physical dynamics with subequivariant graph neural networks. Advances Neural Information Processing Systems, 35:26256-2623 2022.

---

[18]
Jiaqi Han, Yu Rong, Tingyang Xu, and Wenbing Huang, Geometrically equivariant graph neural networks: A survey. arXiv preprint arXiv:2202.07250, 2022.
[19] Tomas Hansson, Chris Oostenbrink, and WilfredF van Gunsteren. Molecular dynamics simula- tions. Current opinion in structural biology, 12(2):190- 19690099999990011 2002.
[20] Wenbing Huang, Jiaqi Han, Yu Rong, Tingyang Xu, Fuchun Sun, and Junzhou Huang. Equiv- ariant graph mechanics networks with constraints. In International Conference on Learning Representations.
[21] Michael J Hutchinson, Charline Le Lan, Sheheryar Zaidi, Emilien Dupont, Yee Whye Teh, and Hyunjik Kim. Lietransformer: Equivariant self-attention for lie groups. In International Conference on Machine Learning, pages 4533-4543. PMLR, 2021.
[22] Guangyin Jin, Yuxuan Liang, Yuchen Fang, Jincai Huang, Junbo Zhang, and Yu Zheng. Spatio- temporal graph neural networks for predictive learning in urban computing: A survey. arXiv preprint arXiv:2303. 14483, 2023.
[23] Thomas Kipf, Ethan Fetaya, Kuan-Chich Wang, Max Welling, and Richard Zemel. Neural relational inference for interacting systems. In International Conference on Machine Learning, pages 1688-2667) PMLR, 2018.
[24] Miltiadis Kofinas, Naveen Nagaraja, and Efstratios Gavves. Roto-translated local coordinate frames for interacting dynamical systems. Advances in Neural Information Processing Systems, 34:6417-6429, 2021.
[25] Langley. Crafting papers on machine learning. In Pat Langley, editor, Proceedings of the 17th International Conference on Machine Learning (ICML 2000), pages 1207-1216, Stanford, CA, 2000. Morgan Kaufmann.
[26] Yaguang Li, Rose Yu, Cyrus Shahabi, and Yan Liu. Diffusion convolutional recurrent neural network: Data-driven traffic forecasting, arXiv preprint arXiv: 2017.
[27] Mario Lino, Stathi Fotiadis, Anil A Bharath, and Chris D Cantwell. Multi-scale rotation- equivariant graph neural networks for unsteady eulerian fluid dynamics. Physics of Fluids, 34(8), 2022.
[28] Elaine C Meng. Thomas D Goddard, Eric F Pettersen, Greg S Couch, Zach J Pearson, John H Morris, and Thomas E Ferrin. Ucsf chimerax: Tools for structure building and analysis. Protein Science, page e4792, 2023.
[29] Damian Mrowca, Chengxu Zhuang, Elias Wang, Nick Haber, Li F Fei-Fei, Josh Tenenbaum, and Daniel L Yamins. Flexible neural representation for physics prediction. Advances in neural information processing systems, 31, 2018.
[30] A Srinivas Reddy, S Priyadarshini Pati, P Praveen Kumar, HN Pradeep, and G Narahari Sastry. Virtual screening in drug discovery-a computational perspective. Current Protein and Peptide Science, 8(4):329-351, 2007.
[31] Alvaro Sanchez-Gonzales Victor Bapst, Kyle Cranmer, and Peter Battaglia. Hamiltonian graph networks with ode integrators. arXiv preprint Stf::::999.2299000 2019.
[32] Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E (n) equivariant graph neural networks. In International conference on machine learning, pages 2332-933.. PMLR, 2021.
[33] Víctor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E(n) equivariant graph neural networks. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 9323 3323-9332. PMLR, 18-24 Jul 2021.
[34] Sean Seyler and Oliver Beckstein. Molecular dynamics trajectory for benchmarking MDAnaly- sis. 2017.
[35] Mark W Spong, Seth Hutchinson, Mathukumalli Vidyasagar, al. Robot modeling and control, volume 3. Wiley New York, 2006.
12

---

[36] Mingjie Sun, Pengyuan Zhou, Hui Tian, Yong Liao, and Haiyong Xie. Spatial-temporae temporal attention network for crime prediction with adaptive graph learning. International Conference on Artificial Neural Networks, pages 656-669. Springer, 2022.
[37] Joshua B Tenenbaum, Charles Kemp, Thomas L Griffiths, and Noah D Goodman. How to grow a mind: Statistics, structure, and abstraction. science, 331(6022): (11022)):229-22888) 2011.
[38] Nathaniel Thomas, Tess Smidt, Steven Kearnes, Lusann Yang. Li Li, Kai Kohlhoff, and Patrick Riley. Tensor field networks: Rotation-and tranllatio-equiivalaan neural networks for 3d point clouds. arXiv preprint arXiv:/802.082219 2018.
[39] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017.
[40] Pantelis R Vlachas, Julija Zavadlav, Matej Praprotnik, and Petros Koumoutsakos. Accelerated simulations of molecular systems through learning of effective dynamics. Journal of Chemical Theory and Computation, 18(1):538-549 2021.
[41] Zonghan Wu, Shirui Pan, Fengwen Chen, Guodong Long, Chengqi Zhang, and S Yu Philip. A comprehensive survey on graph neural networks. IEEE transactions on neural networks and learning systems, 32(1 ):4 24, 2020.
[42] Chenxin Xu, Robby T Tan, Yuhong Tan, Siheng Chen, Yu Guang Wang, Xinchao Wang, and Yanfeng Wang. Eqmotion: Equivariant multi-agent motion prediction with invariant interaction reasoning. Proceedings of the IEEE/CVE Conference on Computer Vision and Pattern Recognition, pages 1410-1420, 2023.
[43] Sijie Yan, Yuanjun Xiong, and Dahua Lin. Spatial temporal graph convolutional networks for skeleton-based action recognition. In Thirty-second AAAI conference on artificial intelligence, 2018.
[44] Bing Yu, Haoteng Yin, and Zhanxing Zhu. Spatio-temports graph convolutional networks: deep learning framework for traffic forecasting. In Proceedings of the 27th International Joint Conference Artificial Intelligence, pages 3634-36400 2018
[45] Jiani Zhang, Xingjian Shi, Junyuan Xie, Hao Ma, Irwin King, and Yan Yeung. Gaan: Gated attention networks for learning on large and spatiotemporal graphs. arXiv preprint arXiv:1803.07294 2018.
13

---

A Proof of ESTAG's Equivariance
Theorem A.1. We denote ESTAG as (TT) @ (((H()).) x(t), A))=d))
then o is E13)
equivariant.
Proof. 1. We firstly prove that EDFT is E3)--qqiiariint..
Of((k) Σ E #k kke
Oz(() +b) (bE(()+b-Oa()++
1=0
(k)))) = we w, (t)
<(k) = w.((b)]l0f(((]])
2. We secondly prove the E3)-eqquivarianee of ESM.
h (t), hg 10
mi = dm
h (t) on h (t), Σ mis
(t)
Oa,(1) = [1(1]
Σ (EN(
Oa (m))))m)))
0a (z((t)(() + Oz (t) (t)+ + Oa! (t).
3. We then prove that ETM is E(3 E(3)_equvvaraant.
(1)() 04 h (0
(t)
k ( dR h (1).
00 ( de (1))
explq: (t) k (s)
af (16()) Σ exp(1 (1) k (s))
h (() (1)(() Σe
a 1
(Is)o (s),.
(t)
b+
a (ts) (z.0((t8 (1) (s)).
4. We finally prove that the linear pooling is equivariant:
(T) + b= OX,W + (T- (T 1) +b.
B Full Algorithm Details
In the main body of the paper, for better readability, we present the implementation details of ESTAG. Here, we combine them into one singe algorithmic flowchart in Algorithm [
14

---

C Extended Models
C.1
Multi-channel ESTAG
For proteins, there are four backbone atoms ( (N, Cao C, O) in residue i, hence the above men- tioned node position vector (t) E R is extended to a 4-channel position matrix X,() E R2x4 Particularly, we denote क. (t), a certain column from (t)) as the position of C. at time t.
EDFT:
T-1 fi(k) = Σ 1=
kt (1))))
A1y(k) (h, (b)))
(k) = (
ESM:
( A.s Au))
my Qm
h
t (t),,
[|(X; (0)) T (t)lee
(t)le
h (() Φh
h (t), (k)) Em,y mis
(1) [N(T]]
15 (t)x((mm))) JEN())
Σ
(())(())
(t) (1)
A,
ETM:
a; (1s) exp(a 10 (1) k: (s))
Σx_ exp(a (1) k (s))
h (() h (1)+ Σ
1) (ts)o (s),
x((+)(()== = (t)
Σ
(1
(ts) X ( (t)))
a
where
q (1) 0 h (1 (t)
(00 du ( (t).
(1) dv (1)
D More Experimental Details and Results
D.1 More Details on MD17
Experiment setup and hyper-parameters. We use the following super-parameters across all experimental evaluations: batch size 100, the number of epochs 500, weight decay 1 x 10-12 12 the number of layers 4 (we consider one ESTAG includes two layers, i.e. ESM and ETM), hidden dim 16, Adam optimizer with learning rate 5 10-3, We set the length of previous time series J. = 10 and the interval between two timestamps Al- 10. The number of training, validation and testing sets are 500, 2000 and 2000, respectively.
15

---

D.2 More Details on Protein
Experiment setup and hyper-parameeers. We use the following spper-parameterss across all experimental evaluations: batch size 100, the number of epochs 500, weight decay 1 x 10-12 the number of layers 4 (we consider one ESTAG includes two layers, i.. ESM and ETM), hidden dim 16, Adam optimizer with learning rate 5 10-5. We divide the whole dataset into training. validation and testing sets by a ratio of 6:2:2, resulting the numbers of the three sets are 2482, 827 and 827 respectively. The number of previous timestamps is T 10 and the interval between timestamps is At 5 frames.
D.3 More Details Motion
Experiment setup and hyper-parameters. We use the following syper-parameters across all experimental evaluations: batch size 100, the number of epochs 500, weight decay 1 x 10 the number of layers 4 (we consider one ESTAG includes two layers, i.e. ESM and ETM), hidden dim 16, Adam optimizer with learning rate 5 10-3. We set the length of previous time series I. 10 and the interval between two timestamps At = 5. The traiiingvalliatiioneesstg sets sizes are 3000/800/800 respectively.
D.4 Long-term recurrent forecasting
We additionally explore the performance of the proposed method in long-term recurrent forecasting. The setting in our current experiment predicts only one frame at time Here we recurrently predict the future frames at time T, AT, 7 + 2AT, T + 10AT (the value of AT follows the setting in the Section 5) in rollout manner, where the the next frame prediction, within a sliding window of length T. Note that the recurrent forecasting task is more challenging than the original scenario, and we need to make some extra improvements to prevent accumulated errors over time. Particularly for our method, we change the forward attention mechanism to be full attention mechanism (namely replacing & with T - 1 in the superscript of the summation of Eq. 10 and 12.12 as we observe that under the recurrent setting, forward attention tends lead to biased predictions. The results are reported in Figure [5] where verify that the rollout version of ESTAG delivers generally smaller MSE than compared methods for all time steps.
frame will be used as the input for
$1,500
14
Figure 5: The rollout-MSE curves on 8 molecules in MD17. Our model generally achieves the lowest MSE.
D.5 More visualization
Visualization of data. To validate that the movement of MD17 molecules is periodical, Figure depict the trajectory of one randomly selected atom in each molecule, from timestep 127947 to timestep 327947. It obvious that almost all molecules move with period.
Visualization of model parameters. Moreover, we display the attention map in ETM of MAL- ONALDEHHDDE'S atom O and temporal pooling weights of all molecules, as shown in Figure
16

---

20 SENAENS
EIHAMOL
MALONALDETFOE
CS.
0.
2000
5330
19000
13300
20008
53300
20066
RACHI HALOME
SAUCYUC
TOLUENS
UNNOI
-07
-31
12930
1990S
20093
5909
19093
19996
27999
Figure 6: The movement trajectory of molecules in MD17.
Attention Map
Temporal Weights
1.0
0.3
0.8
2
Temporal Index
Molecule Index
0.2
3
0.6
4
0.1
5
0.4
6
7
-0.1
0.2
8
-0.2
-0.0
0
1
2
3 4
5
6
1
3
5
6
7
8
Temporal Index
Temporal Index
Figure 7: The visualization of temporal attention Map (left) and temporal pooling weights (right).
Visualization of the prediction results on MD17 dataset. For those non con-obvious cases in Figure the unclarity is mainly caused by 2D printing that is hard to depict 3D conformations such as angle molecule URACIL which not clearly visualized,
deviation. For better visualization, we choose and render its 3D atom coordinates in Figure g ESTAG yields more accurate 3D prediction than ST ST_ENNN, which is consistent with the MSE difference.
Visualization of the prediction results on Protein dataset. We further evaluate the predicted protein by visualizing it alongside the ground truth in Figure P using the ChimeraX software (28). It is obvious that protein predicted by ESTAG aligns closer to the ground truth and ST EGNN fails to predict the alpha helix which highlight in arrow.
Visualization of the prediction results on Motion dataset. Following the visualization manner used for molecules, represent the entities from both walking and basketball motion scenarios within 3D coordinate system, in Figure 10 and Figure respectively. It confirms the superiority of ESTAG, particularly in the significantly more challenging task of simulating baskethall motion.
E
More Ablation Studies
E.1
The impact of the number of past timestamps T
We investigate effect of the number of previous timestamps 7 on the Ethanol dataset. We vary T from 3, 5, 8, 10, 20, 30, 40, 50, 60, 70 and present the results in Table/aand Figure 12] To balance the trade-off between efficiency and performance, we choose T=10 in our experiments.
17

---

0.2
0.2
0.2
0.2
-0.4
-0,4
2
2
-1 22002 2
-2
-2
Figure 8: 3D visualization of URACIL for comparison between ESTAG (left, MSE=5.988 x 10-6) and ST EGNN (right, MSE=6.393 are blue. The molecule predicted by ESTAG is better aligned with the ground truth, particularly for the regions highlighted by circles.
10 10-4). The ground truths in red while the predicted states
Alpht helte
Alpha helts
Alpha hette
Alpin
Figure 9: Comparison on Protein dataset between ESTAG (left, MSE=1.014) and ST_EGNN (right, MSE=1.388), The ground truths are in red while the predicted states are in blue.
E.2 The impact of the learnable parameter we in EDFT module
In section We have noticed that wy acts like spectral filters of the k-th frequency and enables us to select related frequency for the prediction. Here, we conduct ablation study MD17 dataset to analyse how significant the parameter is. From results in Table 7 wv is indeed capable of improving the prediction of physical dynamics.
E.3
Ablation studies on Protein and Motion datasets
Here we conduct several ablations on Protein and Motion dataset similar to MD17 and show the results in Table and Table respectively.
Table 6:
MSE on Ethanol w.r.f. the number of previous timestamps T
3
8
30
50
60
MSE (x 10 4)
1.760
1.395
1.064
0,990
0.924
1,017
0.962
1.009
0.992
1.019
18

---

30
30
25
25
20
20
15
15
10
10
5
7.5
7.5
5,0
-2
0.0
2
0.0
2
5.00
2.5
5500
-2.5
Figure 10: Comparison on Motion walk subject between ESTAG (left, MSE=0.0048) and ST EGNN (right, MSE=0.0811) 1). The ground truths are in red while the predicted states are in bluc.
20
15
15
10
5
4200 9
-2
0
4
Figure 11: Comparison Motion basketball subject between ESTAG (left, MSE=0.0749) and
ST ST_ENNN (right, MSE=2.6380). ground truths are in red while predicted states in blue.
ETHANOL
1.8
MSE(x100>)
1.4
1.2
35
Figure 12: MSE on Ethanol w.r.t. the number of previous timestamps T
Table Ablation studies (x10-3) of learnable parameter wy on MD17 dataset. Results averaged
across 3 runs.
Aspirin Benzene Ethanol Malonaldehyde Naphthalene Salicylic Toluene Uracil
ESTAG
0.063
0.003
0.099
0.101
0.068
0.047
0.079
0.066
w/o wi
0.071
0.003
0.102
0.104
0.081
0.081
0.079
0.069
19

---

Table 8: Ablation study Protein dataset.
Table 9: Ablation study Motion dataset.
Model
MSE
Model
MSE (x 10-
ESTAG 1.471
ESTAG
0.040
w/o EDFT
1,490
w/o EDFT
0.053
w/o Attention
1.493
w/o Attention
0.047
w/o Equivariance
2.011
w/o Equivariance
0.702
w/o Temporal
1.472
w/o Temporal
0.044
F Codes
The codes of ESTAG are available at: https//hhthbbboomm/MmiioW////T
20

---

Algorithm 1 Equivariant Tentti- Attentive Graph Networks (ESTAG)
Input: Initial historical graph series (HH) (t), X ( (t), (t)
for i = to N do
Equivariant Discrete Fourier Transform (EDFT):
T-1
fi(k) = Σ t=
P' 7
(at()- - a(t)
(15)
(k) = wx (h,) ) (h2)) (f((() F, (1)))
(16)
((k) = ww (x.b)))l
(17)
end for for 1 = to L do
for 0 to - do
Equivariant Spatial Module (ESM):
(1=11((). (1). (1). (00)P, (:s)
mij = 2m
hn
(18)
11 1 (1) = Φh
(1)
h
(k)), Σ mis
(19)
a!!-o.5(() WO(1 Σ x!! (11-0.5) (1))
(20)
jEN(i)
0.5) (1) = x (1( (() + a
(3-..5) (I).
(21)
end for
for i = 1 to N do Equivariant Temporal Module (ETM)
cxp(a." l1 (0050 (t) ki (((-0.0)s)) (s))
(- (ts) "' (-.5 (t) k (1- (1-0.)) (s))
d
(22)
h (t) = h (- a
(1 0-
(ts) 11 (s),
(s),
(23)
20((1 = क! = (1) + Σ 0 (00.5) (1s) (1,)) (1s) (2)) (1) ((-0..)
(s)),
(24)
where
d 0-0.5) (t) = vo
h
11 [00.5) (t)
(25)
k aa-0.5) (t) = + h2
(1-0.5) (t)
(26)
v ((-as)(1) = 6 (n! 100.5) (1)
(27)
end for end for Equivariant linear pooling:
(T) N + E where x;= - (1) Output:X..1T (T)
(T 1),
(28)
(1) - (T 1), ... a (T- 2) - xT 1)).
21


Reviews:
[{"review_id": "wzKILW49vP", "reviewer": "Reviewer_iyxR", "summary": "This paper aims to simulate the physical dynamics with a spatio-temporal attentive graph network. The major contribution is to integrate the concept of spatio-temporal graph neural network with DFT to capture the data dependencies. The proposed model is evaluated on three datasets regarding the molecular-, protein- and macro-level prediciton.", "questions": "Please see the weaknesses.", "limitations": "N/A", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The paper is well-written and easy to follow.\n2. The proposed model is technically solid.\n3. The experiments cover three real-world datasets.", "weaknesses": "1. My primary concern is technical novelty. This paper adapts STGNN for physical dynamic simulation without significant contributions or novel designs. The paper claims that it is important to relieve from the Markovian assumption, however, it has been widely explored in the literature on STGNNs. \n2. The related work should be further investigated. As far as I know, there are many advanced STGNNs that can achieve much higher accuracy than STGCN (which was proposed in 2018). For more details, please refer to a recent survey [1].\n3. More powerful baselines should be discussed and considered as baselines [1]. This paper only compares ESTAN with STGCN in the line of STGNNs, which is not convincing enough. For example, comparing ESTAG with an existing STGNN published in 2022 is not investigated.\n\nReference:\n[1] Jin, Guangyin, et al. \"Spatio-temporal graph neural networks for predictive learning in urban computing: A survey.\" arXiv preprint arXiv:2303.14483 (2023).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your comments!  We provide the following responses to your concerns:\n>Q1: My primary concern is technical novelty. This paper adapts STGNN for physical dynamic simulation without significant contributions or novel designs. The paper claims that it is important to relieve from the Markovian assumption, however, it has been widely explored in the literature on STGNNs.\n\nThe reviewer probably misunderstood our contributions. We strongly disagree that our paper adapts STGNN without significant contributions or novel designs. We provide the reasons below.\n\n1. **Equivariance is not well explored in most STGNNs.**\nThe main focus of our paper is on the task of 3D physical dynamics simulation, while most spatio-temporal GNNs discussed in the mentioned survey [A] are developed for uban computing. The unique challenge of physical dynamics simulation compared to uban computing is that, the model for physical dynamics simulation should obey E(3) equivariance: transforming the input coordinates by any translation/rotation/reflection will result in the output transformed in the same way. Most STGNNs, unfortunately, do not satify this crucial symmetry, as already pointed out in Line 50-52 in the paper. Exploring equivariant GNNs is currently a popular and challenging topic in machine learning [24, 16]. Our paper moves a step forward by investigating equivariant spatio-temporal GNNs, which exhibit particular difficulty and untriviality; for example, we should make both the spatial and temporal message passing equivariant. \nWe understand that the non-Markovian property can be modeled in STGNNs, but our claim of relieving from the Markovian assumption is made in comparison with those equivariant GNNs without spatio-temporal modeling (such as EGNN). Overall, we propose an equivariant version of spatio-temporal GNNs, which can not only encode the  non-Markovian property but also ensure E(3) equivariance for physical simulation. \n\n2. **The proposed equivariant model is novel.**\nThe entire architecture we design is novel, which consists of three equivariant modules: Equivariant Discrete Fourier Transform (EDFT),  Equivariant Spatial Module (ESM), and Equivariant Temporal Module (ETM). Particularly, to the best of our knowledge, there is no previous attempt to develop equivariant DFT; in this paper, we achieve this by first translating the signals by the mean position and then adopt the same basis over the spatial dimension (Eq.3). The extracted invariant frequencies are then embedded into ESM to better leverage periodicity patterns. For ETM, the equivariant attention-based mechanism is also novel and carefully designed to ensure equivariance. \n\n>Q2: The related work should be further investigated. As far as I know, there are many advanced STGNNs that can achieve much higher accuracy than STGCN (which was proposed in 2018). For more details, please refer to a recent survey [A].\n\nThank you for your suggestion. We will cite the mentioned survey in Related Work. Indeed, besides STGCN [34], we did investigate ST-GCN [33], GaAN [35], and ASTGCN [13] in Line 85-95. We will discuss more advanced methods as suggested. It is worth mentioning that most spatio-temporal GNNs mentioned in [A] focus on urban computing, while our paper is concerned with the task of 3D physical dynamics simulation. \n\n>Q3: More powerful baselines should be discussed and considered as baselines [A]. This paper only compares ESTAG with STGCN in the line of STGNNs, which is not convincing enough. For example, comparing ESTAG with an existing STGNN published in 2022 is not investigated.\n\nThank you for your suggestion. We have carefully read the survey you mentioned and choose AGL-STAN published in 2022 for comparison, since the encoder it used is transformer-based and competitive in performance. We find that AGL-STAN is originally for crime prediction, and it performs much worse for our task by directly predicting the absolute coordinates (i.e. AGL-STAN (abs)). We then modify AGL-STAN to predict the relative coordinates across two adjacent frames (i.e. AGL-STAN (rel)). The results become better and are tabulated below:\n|MD17|ASPIRIN|BENZENE|ETHANOL|MALONALDEHYDE|NAPHTHALENE|SALICYLIC|TOLUENE|URACIL|\n|-|-|-|-|-|-|-|-|-|\n|AGL_STAN (abs)|4.084|1.651|1.358|1.135|0.938|1.003|1.502|0.784|\n|AGL_STAN (rel)|0.719|0.106|0.459|0.596|0.601|0.452|0.683|0.515|\n| ESTAG | 0.063 | 0.003 |0.099 | 0.101 | 0.068 | 0.047 | 0.079 | 0.066 |\n\n\n|Protein|MSE|\n|-|-|\n|AGL_STAN(abs)|1.859|\n|AGL_STAN(rel)|1.671|\n| ESTAG |1.471 |\n\n\n|Motion|walk($\\times 10^{-1}$)|basketball($\\times 10^{-1}$)|\n|-|-|-|\n|AGL_STAN (abs)|1.675|189.082|\n|AGL_STAN (rel)|0.037|5.734|\n| ESTAG | 0.040 |0.746|\n\nWe still observe that our ESTAG generally outperforms AGL-STAN. AGL-STAN (rel) performs slightly better than ESTAG on Motion-Walk, but for Motion-Basketball which is more complicated to simulate, ESTAG yields a much lower MSE than AGL-STAN (0.746 vs 5.734).  Again, we highlight that AGL-STAN is not equivariant, and it could fail if we steer the input via E(3) transformation. The above discussions will be added into the revised paper to address the reviewer's concern.\n\n[A]  Jin, Guangyin, et al. \"Spatio-temporal graph neural networks for predictive learning in urban computing: A survey.\""}, {"review_id": "4LisAvnDDv", "reviewer": "Reviewer_Hnbu", "summary": "This work proposes a novel architecture for predicting physical dynamics. This architecture first extracts the frequency feature of the input dynamics using a new technique. The frequency feature is then processed by spatial-temporal networks to generate predictions for the future dynamic. Through evaluating the novel architecture on multiple datasets ranging from molecular to macro levels, the authors show that their architecture outperform existing models significantly, especially on the molecular benchmark. The authors also present ablation studies to show which module is the most critical in yielding better performance.", "questions": "See the weakness.", "limitations": "yes", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The architecture proposed in this work (ESTAG) is novel. It is also shown to significantly outperform earlier works in this work on multiple datasets. The paper is well written and easy to read. The ablation studies conducted in this work also shows that the newly proposed frequency computation technique seems to be the most important module in the model. The result on the molecular dataset (MD17) is especially convincing that ESTAG is better than other models.", "weaknesses": "There are some details in the evaluation that are not clear. In the MD-17 dataset, the visualization seems to be really hard to differentiate the ESTAG model and the STEGNN model. But why is the MSE difference so big? Is the observed MSE difference in fact important for whatever downstream task that’s important? Or the STEGNN’s result is already good enough? The other two datasets (Protein and Motion datasets) show that ESTAG is still better than other models, but the gap is much smaller than the MD-17 dataset. It is unclear to me why ESTAG on MD-17 yields such a good result but not on others, can the authors provide an explanation on this? Moreover, on the motion dataset, only one person’s trajectory is used. Why only using one person’s trajectory but not all subjects’?  Why visualizations on the other two datasets are not provided?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are grateful for your positive and constructive comments, and provide the answers to your questions below. \n\n>Q1：In the MD17 dataset, the visualization seems to be really hard to differentiate the ESTAG model and the STEGNN model. But why is the MSE difference so big? Is the observed MSE difference in fact important for whatever downstream task that’s important? Or the STEGNN’s result is already good enough?\n\nSorry for the unclear visualization showing the difference between ESTAG and STEGNN in Figure 4. Actually, for the molecules (ASPIRIN, NAPHTHALENE, SALICYLIC) highlighted by red rectangles, the difference between  ESTAG and STEGNN is obvious.  For those nonobious cases, the unclarity is mainly caused by 2D printing that is hard to depict 3D conformations such as angle deviation. For better visualization, we choose the molecule URACIL which is not clearly visualized in the paper, and redisplay its 3D atom coordinates in Figure A4 (General Response). ESTAG yields more accurate 3D prediction than STEGNN, which is consistent with the MSE difference. We will add the new visualizations to the revised paper. \n\n>Q2: The other two datasets (Protein and Motion datasets) show that ESTAG is still better than other models, but the gap is much smaller than the MD17 dataset. It is unclear to me why ESTAG on MD17 yields such a good result but not on others, can the authors provide an explanation on this?\n\nThank you for this nice observation. Here, we provide some potential explanations on why ESTAG on MD17 yields such a good result but not on the other two datasets:\n1. On Protein dataset, the dynamics of a protein is much more complicated than those small molecules on MD17, owing to various kinds of physical interaction between different amino acids, let along each amino acid compose of a certain number of atoms.  We conjecture that our ESTAG is still hard to reveal sufficiently accurate dynamical patterns for proteins, even though its performance is already better than other methods.\n2. On Motion dataset, ESTAG performs much better than other methods except STGCN. We suspect that the simulation of the walking motion is not that challenging. When we follow the reviewer's suggestion in Q3 and additionally conduct evaluation on a more challenging and complicated task: the basketball motion, the gap between ESTAG and STGCN is significant (see the table below), which indicates the effectiveness of ESTAG in broad and practical cases.\n\n| Motion_basketball | MSE ($\\times 10^{-1}$) |\n|:---:|:---:|\n| PT-s | 886.023 |\n| PT-m | 413.306 |\n| PT-t | 15.878 |\n| Baseline-s | 749.486 |\n| Baseline-m | 335.002 |\n| Baseline-t | 12.492 |\n| GNN | 15.336 |\n| EGNN | 13.199 |\n| TFN | 13.709 |\n| SE3 | 13.851 |\n| STGCN | 4.919 |\n| ESTAG | 0.746 | \n\nWe will add the above explanastions to the revised paper. \n\n>Q3: Moreover, on the motion dataset, only one person’s trajectory is used. Why only using one person’s trajectory but not all subjects’? \n\nThank you for this comment. Indeed, there are more than one trajectories in one subject. The reason why we only use one subject (#subject 35) for evaluation is following the setting of GMN [16] which is the initial work to explore equivariant dynamics simulation on this motion dataset. To better demonstrate the effectiveness of our method in more cases, we additionally carry out experiments on the basketball motion (#subject 102) which is more challenging to simulate. Notably, for basketball motion data, we focus on the trajectories whose length is greater than 170. The results are reported when answering Q2, where we can observe that our ESTAG outperforms other methods remarkably. Conducting evaluations on all subjects (the total number is 144) is expensive and could be better left for future exploraion. \n\n>Q4: Why visualizations on the other two datasets are not provided?\n\nWe apologize for the missed visualization in the paper. Here, we redisplay the visualization on Protein and Motion Cature in Figure A1-A3 (General Response). It can be seen that ESTAG achieves better prediction accuracy. "}, {"review_id": "XoKOeWB3Fx", "reviewer": "Reviewer_Gpx3", "summary": "This paper studies the non-Markovian dynamics that often appear in physical systems and proposes a spatio-temporal E(3) equivariant graph network that moves beyond the simple frame-to-frame prediction task.\nThe authors introduce an equivariant feature extraction method based on Fourier Transform, as well as separable equivariant spatial and temporal modules to process spatio-temporal information.\nThey evaluate the proposed method on different benchmarks and vastly outperform frame-to-frame equivariant methods and non-equivariant spatio-temporal GNNs.", "questions": "1. Following the weaknesses above, a comparison with other spatio-temporal equivariant graph networks would further enhance the credibility of the proposed method.\n\n2. How important is the use of $w_k$, given that these features are further processed during message passing?\n\n3. Since this work focuses on non-Markovian dynamics, an ablation study on the optimal number of past timesteps would be beneficial and insightful.", "limitations": "The authors have adequately addressed the limitations.", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "The paper studies the non-Markovian dynamics in physical systems, an often overlooked yet very important property.\nIt is well-written and easy to follow. The novelties are clear, and the ablation studies support their usefulness.\nThe quantitative results show a massive performance gain from incorporating equivariance and sequence dynamics.\n", "weaknesses": "The paper claims that \"we are the first to use equivariant spatio-temporal graph models for physical dynamics simulation\", yet it is missing out on a few related works, namely LoCS [1], and more recently, EqMotion [2]. Both works propose equivariant graph networks and focus on sequence-to-sequence prediction for physical systems. Hence, the authors should do a more thorough literature review and adjust their claims.\n\nMany of the neural network modules used in the proposed method are not adequately described in the manuscript, and their significance is not tested with an ablation study. For example, the learnable parameters $w_k$ are only briefly described, and their exact form, as well as their usefulness, are unclear.\n\n#### References\n[1] Kofinas, Miltiadis et al. Roto-translated Local Coordinate Frames for Interacting Dynamical Systems. NeurIPS 2021.\n\n[2] Xu, Chenxin et al. EqMotion: Equivariant Multi-agent Motion Prediction with Invariant Interaction Reasoning. CVPR 2023.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are grateful for your positive and constructive comments, and provide the answers to your questions below. \n>Q1: Following the weaknesses above, a comparison with other spatio-temporal equivariant graph networks would further enhance the credibility of the proposed method.\n\nThank you very much for raising these two papers: LoCS and EqMotion, which will be definitely cited and discussed in Related Work.  LoCS proposes two versions: the Markovian one and the non-Markovian one; the non-Markovian version is related to our method and it resorts to GRU units to record the memory of past frames, and predict the next frame conditional on the current frame in an auto-regressive manner, which can be regarded as the RNN style. On the contrary, our method resorts to the spatio-temporal setting and employs all past frames as the input to predict the target one, which can be regarded as the Transformer style. As for EqMotion, it first distills the input trajectory of each node into one multi-dimension vector, by which the spatio-temporal graph is compressed as one single spatial graph. By contrast, our method retains all input frames within both the spatial and temporal modules, such that it is able to better capture spatio-temporal correlations in a more elaborate way.  \n\nHere, we additionally implement  EqMotion on MD17 and Motion datasets for evaluation, since its code is more friendly to be adjusted for these two tasks. For fair comparisons, the input of EqMotion only contain node coordinates, as the same as our method and other baselines.  We find that EqMotion performs much worse by directly predicting the absolute coordinates. We then modify EqMotion to predict the relative coordinates across two adjacent frames. Besides, we perform zero-mean normalization by subtracting all coordinate vectors of all nodes and all graphs with their mean, for further improvement of EqMotion on Motion dataset. The results are reported as follows, where the clear superiority of our ESTAG is still observed. \n\n|MD17 | ASPIRIN | BENZENE | ETHANOL | MALONALDEHYDE | NAPHTHALENE | SALICYLIC | TOLUENE | URACIL |\n|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|\n|EqMotion  | 0.721 | 0.156 | 0.476 | 0.600 | 0.747 | 0.697 | 0.691 | 0.681 |\n| ESTAG | 0.063 | 0.003 | 0.099 |0.101 | 0.068 | 0.047 | 0.079 | 0.066|\n\n\n| Motion | walk ($\\times 10^{-1}$) | basketball ($\\times 10^{-1}$) |\n|:---:|:---:|:---:|\n| EqMotion | 201.008 | 1362.900 |\n| EqMotion (zero-mean) | 1.011 | 4.893 |\n| ESTAG | 0.040 | 0.746|\n\nWe will add the above discussions and adjust our claims accordingly.\n\n>Q2: How important is the use of $w_k$ given that these features are further processed during message passing?\n\nSorry for the currently insufficient explanation. As mentioned Line 140-142, $w_k$ acts like spectral filters of the $k$-th frequency and enables us to select related frequency for the prediction. It is calculated as $\\omega_k=f(\\mathbf{h})$,\nwhere $w_k$ is a scalar ranging from 0 to 1, $f$ is implemented as an MLP plus a Sigmoid output, and $\\mathbf{h}$ is the input feature. Here, we conduct an ablation study to evaluate the effect of $w_k$:\n|MD17 | ASPIRIN | BENZENE | ETHANOL | MALONALDEHYDE | NAPHTHALENE | SALICYLIC | TOLUENE | URACIL |\n|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|\n| ESTAG |0.063 | 0.003 | 0.099 | 0.101 | 0.068|0.047 | 0.079 |0.066 |\n|w/o $w_k$ | 0.071 | 0.003 | 0.102 | 0.104 | 0.081 | 0.081 |0.079| 0.069 |\n\nIn a general sense, the use of $w_k$ can further promote the performance, although the improvement is not so remarkable.  We will include the above explanations into the revised paper. \n\n>Q3: Since this work focuses on non-Markovian dynamics, an ablation study on the optimal number of past timesteps would be beneficial and insightful.\n\nThoughtful viewpoint! The number of past timesteps is indeed a significant factor we need to focus on and we did conduct an ablation study in Table 1 and Figure 3 in the supplementary material. We find that extending the number of past timesteps from 3 to 10 is able to generally improve the performance, which exhibits the necessity of non-Markovian modeling. We will highlight this point in the main paper. "}, {"review_id": "0P4KrMlkEv", "reviewer": "Reviewer_pTt6", "summary": "This paper addresses the Markov limitation of previous methods in simulating physical dynamics by treating it as a spatio-temporal prediction task. The authors propose Equivariant Graph Neural Networks (GNNs) to account for the non-Markovian nature of the systems. Additionally, they design three components to extract spatio-temporal features while preserving equivariance. The experiments conducted on three real datasets demonstrate that the proposed method surpasses previous approaches and validate the effectiveness of the three designed components.", "questions": "Please refer to weaknesses.", "limitations": "The authors do not address the limitations. Please refer to weaknesses.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Identification of Markov limitation: The paper recognizes the Markov limitation present in previous methods and appropriately considers the non-Markovian nature of the systems. This approach is well-founded.\n2. Equivariant property preservation: The designed components successfully maintain the equivariant property while extracting spatio-temporal features. The paper also provides theoretical evidence of the proposed Equivariant Spatio-Temporal Attentive Graph (ESTAG) being E(3)-equivariant.\n3. Comprehensive experimental validation: The paper includes a substantial number of experiments to substantiate the proposed methods. The results consistently demonstrate the superiority of the proposed approach over alternative methods and confirm the effectiveness of the designed Equivariant Discrete Fourie rTransform (EDFT), Equivariant Spatial Module (ESM), and Equivariant Temporal Module (ETM).", "weaknesses": "1. Lack of clarity regarding EDFT: The paper does not provide sufficient explanation of how EDFT improves prediction accuracy. It is important to clarify the underlying mechanisms and intuition behind the proposed component.\n2. Need for additional experiments: It would be valuable to conduct further experiments to explore the performance of the proposed method in long-term recurrent forecasting. Providing results in such scenarios would enhance the understanding of the model's capabilities.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are grateful for your positive and constructive comments, and provide the answers to your questions below. \n\n>Q1: Lack of clarity regarding EDFT: The paper does not provide sufficient explanation of how EDFT improves prediction accuracy. It is important to clarify the underlying mechanisms and intuition behind the proposed component.\n\nSorry for the insufficient clarity. Our use of EDFT is inspired by the observation of molecular trajectories on MD17. As already visualized in Figure 1 in the supplementary material, the molecular dynamics exhibit certain periodicity in terms of different frequencies. This motivates us to first transform the trajectories from time domain to frequency domain, and then compute the inner-product between different frequencies in Eq.4. In signal processing, a well-known theorem is that the Fourier transform of cross-correlation between two signals is equal to the inner-product between the Fourier transform of two signals. In this sense, Eq.4 is able to measure the cross-correlation (thus similarity) between any two trajectories in the frequency domain, and thus can be regarded as the adjacent value between nodes in the message passing in Eq. 6. Similarly, Eq. 5 computes the amplitude which is regarded as a node feature in the message passing in Eq.7. We will include the above explanations into the revised version.\n\n>Q2: Need for additional experiments: It would be valuable to conduct further experiments to explore the performance of the proposed method in long-term recurrent forecasting. Providing results in such scenarios would enhance the understanding of the model's capabilities.\n\nNice suggestion!  We additionally explore the performance of the proposed method in long-term recurrent forecasting as suggested. The setting in our current paper predicts only one frame at time $T$. Here we recurrently predict the future frames at time $T, T+\\Delta t, T+2\\Delta t, \\cdots, T+10\\Delta t $ (the value of $\\Delta t$ follows the setting in the paper) in a rollout manner, where the currently-predicted frame will be used as the input for the prediction of the next frame, within a sliding window of length $T$. Note that the recurrent forecasting task is more challenging than the original scenario, and we need to make some extra improvement to prevent accumulated errors over time. Particularly for our method, we change the forward attention mechanism to be full attention mechanism (namely replacing $t$ with $T-1$ in the superscript of the summation of Eq.10 and Eq.12), as we find that the foward attention is prone to biased prediction under the recurrent setting. \nThe results are reported in Figure A5 (General Response), where we verify that the rollout version of ESTAG delivers generally smaller MSE than all compared methods for all time steps.  We will contain the evaluation of the long-term recurrent forecasting in the revised paper. "}, {"review_id": "BwW3NXOyXx", "reviewer": "Reviewer_TWyG", "summary": "This paper introduces an E(3)-invariant temporal attention scheme, calculated with the help of discrete Fourier transform, within the E(3)-equivariant GNN framework. The overall idea of considering higher-order temporal effects in physics is sound, and the formulation appears to be correct.\n\nThere are a few typos that do not affect the overall scoring, I would recommend the authors do a full proofreading.\n\nThere may be missing references and potentially a missing benchmark to compare with. I recommend adding these (in **Weaknesses**).\n\nOverall, this work is solid, and I recommend accepting it for Neural IPS 2023.", "questions": "Please refer to the **Weaknesses** section.", "limitations": "Please see the comments regarding \"combining equivariance with multi-scale\" in the **Weaknesses** section. The attention module for higher-order temporal relationships will significantly increase complexity, which may limit the application in industrial scenarios with huge graphs. The author should either acknowledge this limit, and/or analyze potential remedies for this overhead.", "rating": 7, "confidence": 5, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "The idea presented in the paper is novel, although not groundbreaking. It fills a gap in the existing framework and the direction is practical and meaningful.\n\nThere are empirical improvements.\n\nThe illustrations are very easy to follow.", "weaknesses": "There are some typos in the paper, such as the missing year in reference [16]. I recommend thorough proofreading.\n\n—Lack of previous SOTA for comparison—\n\n[1] Chen, Runfa and Han, Jiaqi and Sun, Fuchun and Huang, Wenbing. \"Subequivariant Graph Reinforcement Learning in 3D Environments\". Link: https://arxiv.org/abs/2305.18951\n\n—Lack of reference for future improvements—\n\nOne future direction I have, which has already been used in [1], is combining equivariance with multi-scale (MS) GNN, as most industrial-level applications involve huge graphs. Therefore, the following papers should be cited as future works. Note that [2] also combines equivariance with MS, similar to [1]. However, since there are significant differences in graph type and application needs between this paper and [2], it is not suggested to compare them directly (but they should still be cited).\n\n[2] Lino, Mario and Fotiadis, Stathi and Bharath, Anil A and Cantwell, Chris D. “Multi-scale rotation-equivariant graph neural networks for unsteady Eulerian fluid dynamics”. Link: https://pubs.aip.org/aip/pof/article/34/8/087110/2847850\n\n[3] Cao, Yadi, Menglei Chai, Minchen Li, and Chenfanfu Jiang. \"Efficient learning of mesh-based physical simulation with bi-stride multi-scale graph neural network.\". Link: https://openreview.net/forum?id=2Mbo7IEtZW\n\n[4] Meire Fortunato, Tobias Pfaff, Peter Wirnsberger, Alexander Pritzel, Peter Battaglia. “MultiScale MeshGraphNets”. Link: https://arxiv.org/abs/2210.00612", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks! Your feedback is instrumental in strengthening our paper.\n\n>Q1: There are some typos in the paper, such as the missing year in reference [16]. I recommend thorough proofreading.\n\nThank you very much for the mentioned typos, and we will fix them and proofread our paper carefully. \n\n>Q2: —Lack of previous SOTA for comparison—\n\nThank you for raising the related paper [A], which will be cited in the revised paper. Notably, the reference [A] clearly differs from our paper in two aspects:  Firstly, [A] mainly incorporates equivariance into Reinforcement Learning (RL) for morphology-agnostic locomotion learning, whereas our paper aims at equivariant dynamics simulation. Secondly, both the policy and Q functions used in [A] still belong to the frame-to-frame prediction paradigm and are constructed under the Markovian assumption, while our model is of the spatio-temporal form to pursue non-Markovian modeling. We will add the above discussions to the revised paper. \n\n>Q3: —Lack of reference for future improvements—\n\nNice suggestion! We agree that combining equivariance with multi-scale (MS) GNN is valuable particuarly for  industrial-level applications involve huge graphs. We are willing to cite and discuss the mentioned papers [B-D], and consider equivariant MS GNN as a future exploration direction. \n\n>Q4: On limitations. \n\nThanks for the suggestion. We will discuss the efficiency issue for industrial scenarios with huge graphs. As suggested by the reviewer, exploring multi-scale architectures upon our model is potential to reduce the complexity overhead, which will be acknowledged in the revised paper. \n\n\n[A] Chen, Runfa and Han, Jiaqi and Sun, Fuchun and Huang, Wenbing. \"Subequivariant Graph Reinforcement Learning in 3D Environments\".\n\n[B] Lino, Mario and Fotiadis, Stathi and Bharath, Anil A and Cantwell, Chris D. “Multi-scale rotation-equivariant graph neural networks for unsteady Eulerian fluid dynamics”.\n\n[C] Cao, Yadi, Menglei Chai, Minchen Li, and Chenfanfu Jiang. \"Efficient learning of mesh-based physical simulation with bi-stride multi-scale graph neural network.\".\n\n[D] Meire Fortunato, Tobias Pfaff, Peter Wirnsberger, Alexander Pritzel, Peter Battaglia. “MultiScale MeshGraphNets”."}]

Author rebuttal:
## General Response

We sincerely thank all reviewers and ACs for their time and efforts on reviewing the paper. We are glad that the reviewers recognized the contributions of our paper, which we briefly summarize as follows.
- **Novelty**. "The architecture proposed in this work (ESTAG) is novel" (Hnbu); "The novelties are clear" (Gpx3); "This approach is well-founded" (pTt6); "The idea presented in the paper is novel" "It fills a gap in the existing framework and the direction is practical and meaningful" (TWyG). 
- **Presentation**. "The paper is well-written and easy to follow" (iyxR); "The paper is well written and easy to read" (Hnbu); "It is well-written and easy to follow" (Gpx3);“The illustrations are very easy to follow” (TWyG).
- **Experiment**. "The experiments cover three real-world datasets" (iyxR);  "It is also shown to significantly outperform earlier works in this work on multiple datasets" (Hnbu); "The quantitative results show a massive performance gain from incorporating equivariance and sequence dynamics" (Gpx3);  "Comprehensive experimental validation" (pTt6); “There are empirical improvements” (TWyG). 

We also appreciate the reviewers for their thoughtful comments and concerns, and provide additional visualizations and experment results in the attached PDF file for more details. We summarize the extra contents as follows.

- **Figure A1 (to Reviewer Hnbu)** visualizes the difference between ESTAG and ST-EGNN on the walk subject on Motion dataset, showing that ESTAG ahieves clearly better prediction accuracy.
- **Figure A2 (to Reviewer Hnbu)** visualizes the difference between ESTAG and ST-EGNN on the basketball subject on Motion dataset, again showing that ESTAG ahieves clearly better prediction accuracy.
- **Figure A3 (to Reviewer Hnbu)** visualizes the difference between ESTAG and ST-EGNN on Protein dataset, where the protein predicted by ESTAG aligns closer to the ground truth while ST-EGNN fails to predict the alpha helix. 
- **Figure A4 (to Reviewer Hnbu)** visualizes  the 3D conformation of the molecular URACIL on MD17, where our ESTAG yields closer prediction to the ground truth. 
- **Figure A5 (to Reviewer pTt6)** visualizes rollout MSE on MD17, verifying the effectiveness of ESTAG in the recurrent forecasting scenarios. 

Meta-review: The paper discusses the problem of modelling non-markovian dynamics in dynamical systems, and specifically by making sure the model preserves equivariance. All reviewers agree that this is a sufficiently novel paper with a good set of experiments. All criticisms were well addressed and reviewers even raised their scores. All in all, I think it is a clear accept.
Decision: Accept (poster)