Title: Equivariant Spatio-Temporal Attentive Graph Networks to Simulate Physical Dynamics
Venue: NeurIPS 2023
Accepted: yes
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

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

1 Introduction

It has been a goal to represent and simulate the dynamics of physical systems by making use of machine learning techniques [37, 8, 12]. 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 [1], virtual screening of drugs and materials [30], modelbased robot planning/control [35], and many others.

Plenty of solutions have been proposed, amongst which the usage of Graph Neural Networks (GNNs) [41] becomes one of the most desirable directions. GNNs naturally model particles or unit elements as nodes, physical relations as edges, and the latent interactions as the message passing thereon. More recently, a line of researches [38, 11, 21, 9, 32, 20] have been concerned with generalizing GNNs to fit the symmetry of our physical world. These works, also known as equivariant GNNs [18], ensure that translating/rotating/reflecting the geometric input of GNNs results in the output transformed in the same way. By handcrafting such Euclidean equivariance, the

GLYPH(cmap:d835)

GLYPH(cmap:d835)

GLYPH(cmap:d835)

GLYPH(cmap:d835)-GLYPH(cmap:d835)

GLYPH(cmap:d835)(GLYPH(cmap:d835))

GLYPH(cmap:d835)

GLYPH(cmap:d835)-GLYPH(cmap:d835)

GLYPH(cmap:d835)(GLYPH(cmap:d835))

(GLYPH(cmap:d835) )

(GLYPH(cmap:d835) )

Our Setting

GLYPH(cmap:d835)

GLYPH(cmap:d835)

8

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 a vital point: the observed physical dynamics are almost non-Markovian . In previous methods, they usually take as input the conformation of a system at a single temporal frame and predict as output the future conformation after a fixed time interval, forming a frame-to-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 [40]. For instance, when we consider simulating the dynamics of a protein that is interacting with an unobserved solvent (such as 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 to reflect the latent and unobserved dynamics. In principle, we can recover the non-Markovian behavior ( e.g. interacting with a solvent) if the past period is sufficiently long. We collect a period of past system states as spatio-temporal graphs, and utilize them as the input to formulate a spatiotemporal prediction task, other than the frame-to-frame problem as usual (see Figure 1). This motivates us to leverage existing Spatio-Temporal GNNs (STGNNs) [44] to fulfill our purpose, 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 to a target protein. Under the spatio-temporal setting, we are able to model periodicity of dynamics, which, however, is less investigated before.

Our contributions are summarized as follows:

We reveal the non-Markov behavior in physical dynamics simulation by developing equivariant spatio-temporal graph models. The proposed model dubbed Equivariant Spatio-Temporal Attentive Graph network (ESTAG) conforms to Euclidean symmetry and alleviates the limitation of the Markovian assumption.

We design a novel Equivariant Discrete Fourier Transform (EDFT) to extract periodic features from the dynamics, and then construct an 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 [2], NRI [23], and HRN [29] are a series of works to learn physical

!

!"#,%

!

%,&

"

'

"

(

"

)"*

#

"

$

…

…

!

!

!

"

"

"!

#

!

#

"

$

#

%

+,

%

+-

%

+.

!

$

!

%

!

&

!

'

ℎ

!

ℎ

(

ℎ

)

&

/

&

0

$

*

$

+,-

…

…

'

*

∗

())

'

2

∗

())

'

3

∗

())

'

∗

())

…

…

'())

ℒ

+

'

+

(

+

)"*

"

'

"

(

"

)"*

0

,

) 1

…

…

-

Equivariant Temporal Module

Equivariant DFT

Input Trajectory

Output Position

×/

Equivariant Linear Pooling

ℎ

"

&

/

system interaction and evolution. Considering the energy conservation and incorporating physical prior knowledge into GNNs, HNN [15], HOGN [31] leverage ODE and Hamiltonian mechanics to capture the interactions in the systems. However, all of the above-mentioned models don't take into account the underlying symmetries of a system. In order to introduce Euclidean equivariance, Tensor-Field networks (TFN) [38] and SE(3)-Transformer [11] equip filters with rotation equivariance by irreducible representation of the SO(3) group. LieTransformer [21]and LieConv [9] 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 equivariant 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-frame forecasting problem.

Spatio-temporal graph neural networks STGNNs [22] aim to capture the spatial and temporal dependency simultaneously and are widely investigated in 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 1D-CNN layers with graph convolutional layers to 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 a really important property.

Equivariant spatio-temporal graph neural networks There are 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 a multi-dimension vector, by which the spatio-temporal graph is compressed as a spatial graph, then it designs an equivariant module and an interaction reasoning module to 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 to fully leverage the spatio-temporal interactions based on equivariant attentions.

3 Notations and Task Definition

The dynamics of physical objects (such as molecules) can be formulated with the notion of spatiotemporal graphs, as shown in Figure 1 (left). In particular, a spatiotemporal graph of node number N and

temporal length T is denoted as {G t = ( V t , E ) } T -1 t =0 . Here, the nodes V t are shared across time, but each node i is assigned with different scalar feature h i ( t ) ∈ R c , position vector ⃗ x i ( t ) ∈ R 3 at different temporal frame t ; the edges E are associated with an identical adjacency matrix A ∈ R N × N × m whose element A ij ∈ R m defines the edge feature between node i and j . We henceforth denote by the matrices H ( t ) ∈ R c × N and ⃗ X ( t ) ∈ R 3 × N the collection of all nodes in G t . Here for simplicity, we specify the time lag as ∆ t = 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 at frame T given the historical graph series {G t } T -1 t =0 . In form, we learn the function ϕ :

Although previous approaches such as EGNN [32] and GMN [20] also claim to tackle physical 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. , G T -1 in Eq. 1.

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/rotation/reflection g in the group E(3), ϕ satisfies:

where the group action · is instantiated as g · ⃗ X ( t ) := ⃗ X ( t ) + b for translation b ∈ R 3 and g · ⃗ X ( t ) := O ⃗ X ( t ) for rotation/reflection O ∈ R 3 × 3 .

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 a 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 t 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 G T after an equivariant temporal pooling layer (§ 4.3). The overall architecture is 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, we are able to 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, to 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:

where, i ′ is the imaginary unit, k = 0 , 1 , · · · , T -1 is the frequency index, ⃗ x ( t ) is the average position of all nodes in the t -th frame G t , and the output ⃗ f i ( k ) ∈ C 3 is complex. The frequencies calculated by Eq. 3 are then utilized to formulate two crucial quantities: the frequency cross-correlation A ij ∈ R T between node i and j , and the frequency amplitude c i ∈ R T of node i .

In signal processing, cross-correlation measures the similarity of two functions f 1 and f 2 . It satisfies F{ f 1 ⋆ f 2 } = F{ f 1 } · F{ f 2 } , 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 as:

where ⟨· , ·⟩ defines the complex inner product. Notably, we have added two learnable parameters w k ( h i ) and w k ( h j ) 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 ij as the edge feature to capture the relationship between different nodes. We use Aspirin as an example and visualize the A in Figure 3.

4

3

1

6

7

2

11

10

8

13

12

9

5

Aspirin

…

𝑨(𝑘)

𝑨(0)

…

𝑨(								)

𝑇 1

-

m

12

11

A22

10

400

0-16

Node Index

c0

0.35

0.13

寸

m

1:04

Node Index

6

8

910111213

0.22

0.0

We further compute for node i the amplitude of the frequency ⃗ f i ( k ) along with the parameter w k ( h i ) :

This term will be used in the update of the hidden features in the next subsection.

A promising property of Eq. 3 is that it is translation invariant and rotation/reflection equivariant. Therefore, both A ij and c i are E(3)-invariant, which will facilitate the design of following modules.

4.2 Equivariant Spatial Module (ESM)

For each graph G t , our ESM is proposed to encode its spatial geometry through equivariant message passing. ESM is built upon EGNN [33] which is a prevailing kind of equivariant GNNs, but it has subtly involved the FT features from the last subsection for enhanced performance beyond EGNN.

The l -th layer message passing in ESM is as below:

̸

where, ϕ m computes the message m ij from node j to i , ϕ h updates the hidden representation h i , ϕ x returns a one-dimensional scalar for the update of ⃗ a i ( t ) , and all the above functions are MultiLayer Perceptrons (MLPs); ⃗ x ij ( t ) = ⃗ x i ( t ) -⃗ x j ( t ) is the relative position and N ( i ) denotes the neighborhoods of node i .

Notably, we leverage the cross-correlation A ij as the edge feature in Eq. 6 to evaluate the connection between node i and j over the global temporal window, since it is computed from the entire trajectory. We also make use of c i as the input of the update in Eq. 7. The benefit of considering these two terms will be ablated in our experiments.

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 self-correspondence of each node's trajectory based on the foward attention mechanism, and more importantly, in an E(3)-equivariant way.

In detail, each layer of ETM conducts the following process:

where, α ts is the attention weight between time t and s , computed by the query q t and key k s ; the hidden feature h i ( t ) is updated as a weighted combination of the value v s ; the position vector ⃗ x i ( t ) is derived from a weighted combination of a one-dimensional scalar ϕ x ( v s ) multiplied with the temporal displace vector ⃗ x ( l ) i ( ts ) = ⃗ x ( l ) i ( t ) -⃗ x ( l ) i ( s ) . Specifically, q ( l ) i ( t ) = ϕ q ( h ( l ) i ( t ) ) , k ( l ) i ( t ) = ϕ k ( h ( l ) i ( t ) ) and v ( l ) i ( t ) = ϕ v ( h ( l ) i ( t ) ) are all E( 3 )-invariant functions. Notably, we derive a particle's next position in a forward-looking way, to 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 ( L ) i ∈ R T × 3 for each node i . Then the predicted coordinates at time T is given by the following equivariant linear pooling:

where the parameter w ∈ R ( T -1) consists of learnable weights, and ˆ X i = [ ⃗ x ( L ) i (0) -⃗ x ( L ) i ( T -1) , ⃗ x ( L ) i (1) -⃗ x ( L ) i ( T -1) , · · · , ⃗ x ( L ) i ( T -2) -⃗ x ( L ) i ( T -1)] is translated by ⃗ x ( L ) i ( T -1) to allow translation invariance.

We train ESTAG end-to-end via the mean squared error (MSE) loss:

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 ) = ϕ ( { ( H ( t ) , g · ⃗ X ( t ) , A ) } T -1 t =0 ) , then ϕ is E( 3 )-equivariant.

Although we mainly exploit EGNN as the backbone (particularly in ESM), our framework is general and can be easily extended to other equivariant GNNs, such as GMN [20], the multi-channel version of EGNN. In general, the extended models deal with multi-channel coordinate ⃗ Z ∈ R 3 × m instead of ⃗ x ∈ R 3 . The most significant feature of these models is to replace the invariant scalar ∥ ⃗ x ∥ 2 in the formulations of the message m ij (Eq. 6) with the term ⃗ Z ⊤ ⃗ Z . It is easily to prove that this term is equivariant to any orthogonal matrix O , i.e. , ( O ⃗ Z ) ⊤ ( O ⃗ Z ) = ⃗ Z ⊤ ⃗ Z , ∀ O ∈ R 3 × 3 , O ⊤ O = I . Besides, it can be reduced to invariant scalar ∥ ⃗ x ∥ 2 when m = 1 . Empirically, we add the normalization term in order to achieve more stable performance: ⃗ Z ⊤ ⃗ Z ∥ ⃗ Z ⊤ ⃗ Z ∥ F , where ∥ · ∥ F is the Frobenius norm. In the above derivation, we only display the formulation on EGNN, the details of multi-channel ESTAG are shown in Appendix C.1.

5 Experiments

Datasets. To verify the superiority of the proposed model, we evaluate our model on three real world datasets: 1) molecular-level: MD17 [5], 2) protein-level: AdK equilibrium trajectory dataset [34] and 3) macro-level: CMU Motion Capture Databse [7]. 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 are not provided on Motion Capture. The original datasets are composed of long trajectories. We randomize the start point and extract the following T +1 points with the interval ∆ t . We take the first T timestamps as previous observations and last timestamp as the future position label.

Baselines. We compare the performance of ESTAG with several baselines: 1) 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 a 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 a frame-to-frame 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 in between. We modify its default settings by predicting the residual coordinate between time T -1 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 a 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: we implement GNN [13], and other equivariant models EGNN, TFN [38], and SE(3)-Transformer [11] for each temporal frame in the historical trajectory and then estimate the future position as the weighted sum of all past frames, where the weights are 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 8 small molecules generated by MDsimulation. We use the atomic number as the time-independent input node feature h (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 hyper-parameters are introduced in Appendix D.1.

Ground Truth

ESTAG

ST_EGNN

SALICYLIC

ASPIRIN

ETHANOL

MALONALDEHYDE

NAPHTHALENE

TOLUENE

URACIL

BENZENE

MSE=0.987

MSE=0.001

MSE=0.135

MSE=0.040

MSE=0.241

MSE=0.203

MSE=0.060

MSE=0.729

MSE=0.088

MSE=0.654

MSE=0.084

MSE=0.654

MSE=0.084

MSE=0.670

Results. Table 1 shows the average MSE of all models on 8 molecules. We have some interesting observations as follows: 1) ESTAG exceeds other models in all cases by a significant margin, supporting the general effectiveness of the proposed ideas. 2) From the Pts/m/t , we observe that the point closer to the future point has less prediction error. EGNNs/m/t only takes one frame ( s/m/t ) as input, and attains slight improvement relative to Pts/m/t . 3) Compared with EGNNt , ST_EGNN, 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 non-equivariant methods particularly ST_GNN 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 4 shows that the prdicted MSEs of our model are much lower than ST_EGNN and our predicted molecules are closer to the ground-truth molecules. We have also displayed the learned attentions (Eq. 10) and the temporal weight (Eq. 13), where we find meaningful patterns of the temporal correlation. For clearer comparison, we display the molecule URACIL in 3D coordinate system, which can be found in Appendix D.5.

5.2 Protein-level: Protein Dynamics

Implementation details. We evaluate our model on the AdK equilibrium trajectory dataset [34] via MDAnalysis toolkit [14]. In order to reduce the data scale, we utilize MDAnalysis to locate the backbone atoms ( C α , C, N, O ) of the residues and then regard the residues other than atoms in protein as the nodes with 4 -channel geometric features. We use the atomic number of four backbone atoms as the time-independent input node feature h (0) . We connect two atoms via an edge if their distance is less than a threshold λ . Other settings including the hyperparameters are introduced in Appendix D.2. Slightly different from the model on MD17 dataset, we generalize the single-channel ESTAG into multi-channel version which is presented in detail in Appendix C. We do not 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 state-of-the-art method GMN [20] is implemented by further adding the weighed temporal pooling similar to ST_EGNN.

Results. The predicted MSEs are displayed in Table 2. Generally, the spatio-temporal models are better than Pts/m/t and EGNNs/m/t , and it 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 nonequivariant counterparts (for example, ST_EGNN vs. ST_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 is acceptable given its remarkable performance enhancement. It is expected that the superiority of ESTAG on protein dataset is not as obvious as that on MD17, owing to various kinds of physical interactions between different amino acids, let along each amino acid is composed of a certain number of atoms, which makes the dynamics of a 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 trajectories of human motion under several scenarios and we focus on 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) h (0) i are all 1 s. The two joints are 1-hop neighbor if they are connected naturally and we consider two types of neighbors (i.e. 1-hop neighbor and 2-hop neighbor). Other settings including the hyper-parameters are introduced in Appendix D.3. 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 coordinates across two adjacent frames and perform zero-mean normalization of node coordinates, for further improvement.

Results. Table 3 summarizes the results of all models on the Motion dataset. The spatio-temporal models are better than Pts/m/t and EGNNs/m/t , which again implies the necessity of taking the spatio-temporal history into account. Unexpectedly, the non-equivariant models are even superior to 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 in 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 a much lower MSE. We also notice that the attention-based models including our ESTAG, 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 to the overall performance and the results are shown in Table 4.

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 can well capture the temporal dynamics.

3) Without equivariance. We construct a 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.

3) Without equivariance. We construct a 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 L on Ethanol dataset. We vary T from 1, 2, 3, 4, 5, 6 and present the results in Table 5. Considering the accuracy and efficiency simultaneously, we choose L = 2 for the ESTAG.

More ablation studies will be shown in Appendix E.

6 Conclusion

In this paper, we propose ESTAG, an end-to-end equivariant architecture for physical dynamics modeling. ESTAG first extracts frequency features via a novel Equivariant Discrete Fourier Transform (EDFT), and then leverages Equivariant Spatial Module (ESM) and an attentive Equivariant Temporal Module (ETM) to 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 to macro-level. Necessary ablations, visualizations, and analyses are also provided to 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 our benchmark with more tasks and datasets and evaluate more baselines to validate the effectiveness of our model. It 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.

References

Bissan Al-Lazikani, Joon Jung, Zhexin Xiang, and Barry Honig. Protein structure prediction. Current opinion in chemical biology , 5(1):51-56, 2001.

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.

Hendrik W Broer, George B Huitema, and Mikhail B Sevryuk. Quasi-periodic motions in families of dynamical systems: order amidst chaos . Springer, 2009.

Yadi Cao, Menglei Chai, Minchen Li, and Chenfanfu Jiang. Efficient learning of mesh-based physical simulation with bi-stride multi-scale graph neural network. 2023.

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):e1603015, 2017.

Junyoung Chung, Caglar Gulcehre, KyungHyun Cho, and Yoshua Bengio. Empirical evaluation of gated recurrent neural networks on sequence modeling. arXiv preprint arXiv:1412.3555 , 2014.

CMU. Carnegie-mellon motion capture database, 2003. URL Lhttp://mocap.cs.cmu.edu.

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.

Marc Finzi, Samuel Stanton, Pavel Izmailov, and Andrew Gordon Wilson. Generalizing convolutional neural networks for equivariance to lie groups on arbitrary continuous data. In International Conference on Machine Learning , pages 3165-3176. PMLR, 2020.

Meire Fortunato, Tobias Pfaff, Peter Wirnsberger, Alexander Pritzel, and Peter Battaglia. Multiscale meshgraphnets. arXiv preprint arXiv:2210.00612 , 2022.

Fabian Fuchs, Daniel Worrall, Volker Fischer, and Max Welling. Se (3)-transformers: 3d roto-translation equivariant attention networks. Advances in Neural Information Processing Systems , 33:1970-1981, 2020.

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.

Justin Gilmer, Samuel S Schoenholz, Patrick F Riley, Oriol Vinyals, and George E Dahl. Neural message passing for quantum chemistry. In International conference on machine learning , pages 1263-1272. PMLR, 2017.

Richard J Gowers, Max Linke, Jonathan Barnoud, Tyler JE Reddy, Manuel N Melo, Sean L 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.

Samuel Greydanus, Misko Dzamba, and Jason Yosinski. Hamiltonian neural networks. Advances in neural information processing systems , 32, 2019.

Shengnan Guo, Youfang Lin, Ning Feng, Chao Song, and Huaiyu Wan. Attention based spatialtemporal graph convolutional networks for traffic flow forecasting. In Proceedings of the AAAI conference on artificial intelligence , volume 33, pages 922-929, 2019.

Jiaqi Han, Wenbing Huang, Hengbo Ma, Jiachen Li, Josh Tenenbaum, and Chuang Gan. Learning physical dynamics with subequivariant graph neural networks. Advances in Neural Information Processing Systems , 35:26256-26268, 2022.

Jiaqi Han, Yu Rong, Tingyang Xu, and Wenbing Huang. Geometrically equivariant graph neural networks: A survey. arXiv preprint arXiv:2202.07230 , 2022.

Tomas Hansson, Chris Oostenbrink, and WilfredF van Gunsteren. Molecular dynamics simulations. Current opinion in structural biology , 12(2):190-196, 2002.

Wenbing Huang, Jiaqi Han, Yu Rong, Tingyang Xu, Fuchun Sun, and Junzhou Huang. Equivariant graph mechanics networks with constraints. In International Conference on Learning Representations .

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.

Guangyin Jin, Yuxuan Liang, Yuchen Fang, Jincai Huang, Junbo Zhang, and Yu Zheng. Spatiotemporal graph neural networks for predictive learning in urban computing: A survey. arXiv preprint arXiv:2303.14483 , 2023.

Thomas Kipf, Ethan Fetaya, Kuan-Chieh Wang, Max Welling, and Richard Zemel. Neural relational inference for interacting systems. In International Conference on Machine Learning , pages 2688-2697. PMLR, 2018.

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.

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

Yaguang Li, Rose Yu, Cyrus Shahabi, and Yan Liu. Diffusion convolutional recurrent neural network: Data-driven traffic forecasting. arXiv preprint arXiv:1707.01926 , 2017.

Mario Lino, Stathi Fotiadis, Anil A Bharath, and Chris D Cantwell. Multi-scale rotationequivariant graph neural networks for unsteady eulerian fluid dynamics. Physics of Fluids , 34(8), 2022.

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.

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.

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.

Alvaro Sanchez-Gonzalez, Victor Bapst, Kyle Cranmer, and Peter Battaglia. Hamiltonian graph networks with ode integrators. arXiv preprint arXiv:1909.12790 , 2019.

Vıctor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E (n) equivariant graph neural networks. In International conference on machine learning , pages 9323-9332. PMLR, 2021.

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-9332. PMLR, 18-24 Jul 2021.

Sean Seyler and Oliver Beckstein. Molecular dynamics trajectory for benchmarking MDAnalysis. 6 2017.

Mark W Spong, Seth Hutchinson, Mathukumalli Vidyasagar, et al. Robot modeling and control , volume 3. Wiley New York, 2006.

Mingjie Sun, Pengyuan Zhou, Hui Tian, Yong Liao, and Haiyong Xie. Spatial-temporal attention network for crime prediction with adaptive graph learning. In International Conference on Artificial Neural Networks , pages 656-669. Springer, 2022.

Joshua B Tenenbaum, Charles Kemp, Thomas L Griffiths, and Noah D Goodman. How to grow a mind: Statistics, structure, and abstraction. science , 331(6022):1279-1285, 2011.

Nathaniel Thomas, Tess Smidt, Steven Kearnes, Lusann Yang, Li Li, Kai Kohlhoff, and Patrick Riley. Tensor field networks: Rotation-and translation-equivariant neural networks for 3d point clouds. arXiv preprint arXiv:1802.08219 , 2018.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems , 30, 2017.

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.

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.

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. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 1410-1420, 2023.

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.

Bing Yu, Haoteng Yin, and Zhanxing Zhu. Spatio-temporal graph convolutional networks: a deep learning framework for traffic forecasting. In Proceedings of the 27th International Joint Conference on Artificial Intelligence , pages 3634-3640, 2018.

Jiani Zhang, Xingjian Shi, Junyuan Xie, Hao Ma, Irwin King, and Dit-Yan Yeung. Gaan: Gated attention networks for learning on large and spatiotemporal graphs. arXiv preprint arXiv:1803.07294 , 2018.

A Proof of ESTAG's Equivariance

Theorem A.1. We denote ESTAG as ⃗ X ( T ) = ϕ ( { ( H ( t ) , g · ⃗ X ( t ) , A ) } T -1 t =0 ) , then ϕ is E( 3 )-equivariant.

Proof. 1. We firstly prove that EDFT is E( 3 )-equivariant.

We secondly prove the E(3)-equivariance of ESM.

We then prove that ETM is E( 3 )-equivariant.

We finally prove that the linear pooling is equivariant:

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

̸

C Extended Models

C.1 Multi-channel ESTAG

For proteins, there are four backbone atoms ( N, C α , C, O) in residue i , hence the above mentioned node position vector x i ( t ) ∈ R 3 is extended to a 4-channel position matrix X i ( t ) ∈ R 3 × 4 . Particularly, we denote ⃗ x α i ( t ) , a certain column from X i ( t ) as the position of C α at time t .

EDFT:

ESM:

ETM:

where

D More Experimental Details and Results

D.1 More Details on MD17

Experiment setup and hyper-parameters. We use the following hyper-parameters across all experimental evaluations: batch size 100, the number of epochs 500, weight decay 1 × 10 -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 L = 10 and the interval between two timestamps ∆ t = 10 . The number of training, validation and testing sets are 500, 2000 and 2000, respectively.

̸

D.2 More Details on Protein

Experiment setup and hyper-parameters. We use the following hyper-parameters across all experimental evaluations: batch size 100, the number of epochs 500, weight decay 1 × 10 -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 -5 . We divide the whole dataset into training, validation and testing sets by a ratio of 6:2:2, resulting in 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 ∆ t = 5 frames.

D.3 More Details on Motion

Experiment setup and hyper-parameters. We use the following hyper-parameters across all experimental evaluations: batch size 100, the number of epochs 500, weight decay 1 × 10 -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 L = 10 and the interval between two timestamps ∆ t = 5 . The training/validation/testing 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 a time . Here we recurrently predict the future frames at time T, T +∆ T, T +2∆ T, · · · , T +10∆ T (the value of ∆ T follows the setting in the Section 5) in a rollout manner, where the currently-predicted frame will be used as the input for 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 t with T -1 in the superscript of the summation of Eq.10 and Eq.12 10), as we observe that under the recurrent setting, forward attention tends to lead to biased predictions. The results are reported in Figure 5, where we verify that the rollout version of ESTAG delivers generally smaller MSE than all compared methods for all time steps.

0

2

4

6

8

Step

10

1

10 0

10 1

MSE (×10

3

)

ASPIRIN

ESTAG

GNN

STGCN

ST_EGNN

0

2

4

6

8

Step

10

3

10

2

10

1

10 0

10 1

10 2

MSE (×10

3

)

BENZENE

ESTAG

GNN

STGCN

ST_EGNN

0

2

4

6

8

Step

10

1

10 0

10 1

MSE (×10

3

)

ETHANOL

ESTAG

GNN

STGCN

ST_EGNN

0

2

4

6

8

Step

10

1

10 0

10 1

MSE (×10

3

)

MALONALDEHYDE

ESTAG

GNN

STGCN

ST_EGNN

0

2

4

6

8

Step

10

1

10 0

10 1

MSE (×10

3

)

NAPHTHALENE

ESTAG

GNN

STGCN

ST_EGNN

0

2

4

6

8

Step

10

1

10 0

10 1

MSE (×10

3

)

SALICYLIC

ESTAG

GNN

STGCN

ST_EGNN

0

2

4

6

8

Step

10 0

10 1

MSE (×10

3

)

TOLUENE

ESTAG

GNN

STGCN

ST_EGNN

0

2

4

6

8

Step

10

1

10 0

10 1

MSE (×10

3

)

URACIL

ESTAG

GNN

STGCN

ST_EGNN

D.5 More visualization

Visualization of data. To validate that the movement of MD17 molecules is periodical, in Figure 6 we depict the trajectory of one randomly selected atom in each molecule, from timestep 127947 to timestep 327947. It is obvious that almost all molecules move with period.

Visualization of model parameters. Moreover, we display the attention map in ETM of MALONALDEHYDE's atom O and temporal pooling weights of all molecules, as shown in Figure 7.

0

20000

40000

60000

80000

t

1.0

0.8

0.6

0.4

0.2

0.0

0.2

coordinates

ASPIRIN

0

50000

100000

150000

200000

t

0.5

0.6

0.7

0.8

0.9

1.0

coordinates

BENZENE

0

50000

100000

150000

200000

t

0.4

0.5

0.6

0.7

0.8

0.9

1.0

coordinates

ETHANOL

0

50000

100000

150000

200000

t

1.0

0.5

0.0

0.5

1.0

coordinates

MALONALDEHYDE

0

50000

100000

150000

200000

t

1.0

0.8

0.6

0.4

0.2

0.0

coordinates

NAPHTHALENE

0

50000

100000

150000

200000

t

1.0

0.8

0.6

0.4

0.2

0.0

0.2

0.4

coordinates

SALICYLIC

0

50000

100000

150000

200000

t

1.0

0.9

0.8

0.7

0.6

coordinates

TOLUENE

0

1000

2000

3000

4000

5000

6000

t

1.0

0.8

0.6

0.4

0.2

0.0

0.2

coordinates

URACIL

0

1

2

3

4

5

6

7

8

9

Temporal Index

0

1

2

3

4

5

6

7

8

9

Temporal Index

Attention Map

0

1

2

3

4

5

6

7

8

Temporal Index

0

1

2

3

4

5

6

7

Molecule Index

Temporal Weights

0.0

0.2

0.4

0.6

0.8

1.0

0.2

0.1

0.0

0.1

0.2

0.3

0.4

Visualization of the prediction results on MD17 dataset. For those non-obvious cases in Figure 4, 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, and render its 3D atom coordinates in Figure 8. ESTAG yields more accurate 3D prediction than ST_EGNN, 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 9, 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 we highlight in arrow.

Visualization of the prediction results on Motion dataset. Following the visualization manner used for molecules, we represent the entities from both walking and basketball motion scenarios within a 3D coordinate system, in Figure 10 and Figure 11 respectively. It confirms the superiority of ESTAG, particularly in the significantly more challenging task of simulating basketball 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 T on the Ethanol dataset. We vary T from 3, 5, 8, 10, 20, 30, 40, 50, 60, 70 and present the results in Table 6 and Figure 12. To balance the trade-off between efficiency and performance, we choose T=10 in our experiments.

0.2

0.2

0.0

0.0

-0.2

-0.2

-0.4

-0.4

1

0

0

0

1

1

2

2

2

Alpha helix

Alpha helix

Alpha helix

Alpha helix

E.2 The impact of the learnable parameter w k in EDFT module

In section 4.1, We have noticed that w k acts like spectral filters of the k -th frequency and enables us to select related frequency for the prediction. Here, we conduct ablation study on MD17 dataset to analyse how significant the parameter is. From the results in Table 7, w k 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 8 and Table 9, respectively.

7.5

5.0

2.5

-4

-2

0.0

0

-2.5

2

4

-5.0

7.5

5.0

-4

-2

0.0

2.5

0

-2.5

2

4

-5.0

20

20

15

15

10

10

5

5

0

0

6

4

4

一2

-4-20

0

2

4

4

2

4

4

6

-6

6

-6

35 810

20

30

40

50

60

70

T

1.0

1.2

1.4

1.6

1.8

5

)

ETHANOL

F Codes

The codes of ESTAG are available at: https://github.com/ManlioWu/ESTAG.

Algorithm 1 Equivariant Spatio-Temporal Attentive Graph Networks (ESTAG)

Input: Initial historical graph series { ( H (0) ( t ) , ⃗ X (0) ( t ) , A } T -1 t =0 .

for i = 1 to N do

Equivariant Discrete Fourier Transform (EDFT):

end for for l = 1 to L do

̸

end for for i = 1 to N do

Equivariant Temporal Module (ETM)

where

end for end for

Equivariant linear pooling:

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nThere 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\n\nQUESTIONS:\nPlease refer to the **Weaknesses** section.\n\nLIMITATIONS:\nPlease 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. 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).\n\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\nPlease refer to weaknesses.\n\nLIMITATIONS:\nThe authors do not address the limitations. Please refer to weaknesses.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThe 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\nWEAKNESSES:\nThe 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.\n\nQUESTIONS:\n1. 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.\n\nLIMITATIONS:\nThe authors have adequately addressed the limitations.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nThere 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?\n\nQUESTIONS:\nSee the weakness.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

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

STRENGTHS:
1. The paper is well-written and easy to follow.
2. The proposed model is technically solid.
3. 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. 
2. 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].
3. 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.

Reference:
[1] Jin, Guangyin, et al. "Spatio-temporal graph neural networks for predictive learning in urban computing: A survey." arXiv preprint arXiv:2303.14483 (2023).

QUESTIONS:
Please see the weaknesses.

LIMITATIONS:
N/A
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-09