SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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 Database [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  $\Delta 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 MD simulation. 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  $\lambda$  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.
Figure 4: PyMol visualization of the predicted molecules by our ESTAG and ST_EGNN, where the MSE  $(\times 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, we highlight them with red rectangles. It is observed that ST_EGNN occasionally outputs isolated atoms, which could be caused by violation of the bond length tolerance in PyMol.
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  $\mathrm{Pt - s / m / t}$ , we observe that 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  $\mathrm{Pt - s / m / t}$ . 3) Compared with EGNN- $t$ , 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 predicted 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_{\alpha}, 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  $\lambda$ . 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.
Table 2: Prediction error and training time on Protein dataset. Results averaged across 3 runs.
METHOD	MSE	TIME(S)
PT-s	3.260	-
PT-m	3.302	-
PT-t	2.022	-
EGNN-s	3.254	1.062
EGNN-m	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 Table 2. Generally, the spatio-temporal models are better than  $\mathrm{Pt - s / m / t}$  and EGNN-  $s / 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_i^{(0)}$  are all 1s. 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.
Table 3: Prediction error  $(\times 10^{-1})$  on Motion dataset. Results averaged across 3 runs.
METHOD	WALK	BASKETBALL
PT-s	329.474	886.023
PT-m	127.152	413.306
PT-t	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_SE(3)TR	0.236	13.851
ST_EGNN	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  $\mathrm{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 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.
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.
Table 4: Ablation studies  $(\times 10^{-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 Equivalence	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$  .
L	1	2	3	4	5	6
MSE (×10-4)	1.25	0.990	1.096	1.022	1.042	1.028
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.
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-Bird Focused Research Program (RBFR2023005); Ant Group through CCF-Ant Research Fund (CCF-AFSG RF20220204); Public Computing Cloud, Renmin University of China.
[36] 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.
[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):1279-1285, 2011.
[38] 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.
[39] Ashish Vavwani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is 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. In Proceedings of the IEEE/CVF 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-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.
[45] 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  $\vec{X}(T) = \phi\left(\{(H(t), g \cdot \vec{X}(t), A)\}_{t=0}^{T-1}\right)$ , then  $\phi$  is  $E(3)$ -equivariant.
Proof. 1. We firstly prove that EDFT is E(3)-equivariant.
$$
\boldsymbol {O} \vec {\boldsymbol {f}} _ {i} (k) = \sum_ {t = 0} ^ {T - 1} e ^ {- i ^ {\prime} \frac {2 \pi}{T} k t} \left(\boldsymbol {O} \vec {\boldsymbol {x}} _ {i} (t) + \boldsymbol {b} - \overline {{\boldsymbol {O} \vec {\boldsymbol {x}} (t) + \boldsymbol {b}}}\right),
$$
$$
\boldsymbol {A} _ {i j} (k) = w _ {k} (\boldsymbol {h} _ {i}) w _ {k} (\boldsymbol {h} _ {j}) | \langle \boldsymbol {O} \vec {\boldsymbol {f}} _ {i} (k), \boldsymbol {O} \vec {\boldsymbol {f}} _ {j} (k) \rangle |,
$$
$$
\boldsymbol {c} _ {i} (k) = w _ {k} (\boldsymbol {h} _ {i}) \| \boldsymbol {O} \vec {\boldsymbol {f}} _ {i} (k) \| ^ {2}.
$$
2. We secondly prove the E(3)-equivariance of ESM.
$$
\pmb {m} _ {i j} = \phi_ {m} \left(\pmb {h} _ {i} ^ {(l)} (t), \pmb {h} _ {j} ^ {(l)} (t), \| \pmb {O} \vec {\pmb {x}} _ {i j} ^ {(l)} (t) \| ^ {2}, \pmb {A} _ {i j}\right),
$$
$$
\boldsymbol {h} _ {i} ^ {(l + 1)} (t) = \phi_ {h} \left(\boldsymbol {h} _ {i} ^ {(l)} (t), \boldsymbol {c} _ {i} (k), \sum_ {j \neq i} \boldsymbol {m} _ {i j}\right),
$$
$$
\boldsymbol {O} \vec {\boldsymbol {a}} _ {i} (t) = \frac {1}{| \mathcal {N} (i) |} \sum_ {j \in \mathcal {N} (i)} \boldsymbol {O} \vec {\boldsymbol {x}} _ {i j} ^ {(l)} (t) \phi_ {x} (\boldsymbol {m} _ {i j}),
$$
$$
\boldsymbol {O} \vec {\boldsymbol {x}} _ {i} ^ {(l + 1)} (t) + \boldsymbol {b} = \boldsymbol {O} \vec {\boldsymbol {x}} _ {i} ^ {(l)} (t) + \boldsymbol {b} + \boldsymbol {O} \vec {\boldsymbol {a}} _ {i} ^ {(l + 1)} (t).
$$
3. We then prove that ETM is E(3)-equivariant.
$$
\boldsymbol {q} _ {i} ^ {(l)} (t) = \phi_ {q} \left(\boldsymbol {h} _ {i} ^ {(l)} (t)\right),
$$
$$
\boldsymbol {k} _ {i} ^ {(l)} (t) = \phi_ {k} \left(\boldsymbol {h} _ {i} ^ {(l)} (t)\right),
$$
$$
\boldsymbol {v} _ {i} ^ {(l)} (t) = \phi_ {v} \left(\boldsymbol {h} _ {i} ^ {(l)} (t)\right),
$$
$$
\alpha_ {i} ^ {(l)} (t s) = \frac {\exp (\boldsymbol {q} _ {i} ^ {(l)} (t) ^ {\top} \boldsymbol {k} _ {i} ^ {(l)} (s))}{\sum_ {s = 0} ^ {t} \exp (\boldsymbol {q} _ {i} ^ {(l)} (t) ^ {\top} \boldsymbol {k} _ {i} ^ {(l)} (s))},
$$
$$
\boldsymbol {h} _ {i} ^ {(l + 1)} (t) = \boldsymbol {h} _ {i} ^ {(l)} (t) + \sum_ {s = 0} ^ {t} \alpha_ {i} ^ {(l)} (t s) \boldsymbol {v} _ {i} ^ {(l)} (s),
$$
$$
\boldsymbol {O} \vec {\boldsymbol {x}} _ {i} ^ {(l + 1)} (t) + \boldsymbol {b} = \boldsymbol {O} \vec {\boldsymbol {x}} _ {i} ^ {(l)} (t) + \boldsymbol {b} + \sum_ {s = 0} ^ {t} \alpha_ {i} ^ {(l)} (t s) \boldsymbol {O} \vec {\boldsymbol {x}} _ {i} ^ {(l)} (t s) \phi_ {x} (\boldsymbol {v} _ {i} ^ {(l)} (s)).
$$
4. We finally prove that the linear pooling is equivariant:
$$
\boldsymbol {O} \vec {\boldsymbol {x}} _ {i} ^ {*} (T) + \boldsymbol {b} = \boldsymbol {O} \hat {\boldsymbol {X}} _ {i} \boldsymbol {w} + \boldsymbol {O} \vec {\boldsymbol {x}} _ {i} ^ {(L)} (T - 1) + \boldsymbol {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 single algorithmic flowchart in Algorithm 1.
C Extended Models
C.1 Multi-channel ESTAG
For proteins, there are four backbone atoms ( N,  $\mathbf{C}_{\alpha}$ , C, O) in residue  $i$ , hence the above mentioned node position vector  $\pmb{x}_i(t) \in \mathbb{R}^3$  is extended to a 4-channel position matrix  $X_{i}(t) \in \mathbb{R}^{3\times 4}$ . Particularly, we denote  $\vec{\pmb{x}}_i^\alpha (t)$ , a certain column from  $X_{i}(t)$  as the position of  $\mathbf{C}_{\alpha}$  at time  $t$ .
EDFT:
$$
\vec {\pmb {f}} _ {i} (k) = \sum_ {t = 0} ^ {T - 1} e ^ {- i ^ {\prime} \frac {2 \pi}{T} k t} \left(\vec {\pmb {x}} _ {i} ^ {\alpha} (t) - \overline {{\vec {\pmb {x}} ^ {\alpha} (t)}}\right),
$$
$$
\boldsymbol {A} _ {i j} (k) = w _ {k} \left(\boldsymbol {h} _ {i}\right) w _ {k} \left(\boldsymbol {h} _ {j}\right) | \langle \vec {\boldsymbol {f}} _ {i} (k), \vec {\boldsymbol {f}} _ {j} (k) \rangle |,
$$
$$
\boldsymbol {c} _ {i} (k) = w _ {k} \left(\boldsymbol {h} _ {i}\right) \| \vec {\boldsymbol {f}} _ {i} ^ {\prime} (k) \| ^ {2}.
$$
ESM:
$$
\boldsymbol {m} _ {i j} = \phi_ {m} \left(\boldsymbol {h} _ {i} ^ {(l)} (t), \boldsymbol {h} _ {j} ^ {(l)} (t), \frac {(\vec {\boldsymbol {X}} _ {i j} ^ {(l)} (t)) ^ {\top} \vec {\boldsymbol {X}} _ {i j} ^ {(l)} (t)}{\| (\vec {\boldsymbol {X}} _ {i j} ^ {(l)} (t)) ^ {\top} \vec {\boldsymbol {X}} _ {i j} ^ {(l)} (t) \| _ {F}}, \boldsymbol {A} _ {i j}\right),
$$
$$
\vec {\boldsymbol {A}} _ {i} ^ {(l)} (t) = \frac {1}{| \mathcal {N} (i) |} \sum_ {j \in \mathcal {N} (i)} \vec {\boldsymbol {X}} _ {i j} ^ {(l)} (t) \phi_ {\boldsymbol {X}} (\boldsymbol {m} _ {i j}),
$$
$$
\vec {X} _ {i} ^ {(l + 1)} (t) = \vec {X} _ {i} ^ {(l)} (t) + \vec {A} _ {i} ^ {(l)} (t).
$$
ETM:
$$
\vec {\boldsymbol {X}} _ {i} ^ {(l + 1)} (t) = \vec {\boldsymbol {X}} _ {i} ^ {(l)} (t) + \sum_ {s = 0} ^ {t} \alpha_ {i} ^ {(l)} (t s) \vec {\boldsymbol {X}} _ {i} ^ {(l)} (t s) \phi_ {\boldsymbol {X}} (\boldsymbol {v} _ {i} ^ {(l)} (s)),
$$
$$
\boldsymbol {v} _ {i} ^ {(l)} (t) = \phi_ {v} \left(\boldsymbol {h} _ {i} ^ {(l)} (t)\right).
$$
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 \times 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 \times 10^{-3}$ . We set the length of previous time series  $L = 10$  and the interval between two timestamps  $\Delta 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 \times 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 \times 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  $\Delta 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 \times 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 \times 10^{-3}$ . We set the length of previous time series  $L = 10$  and the interval between two timestamps  $\Delta t = 5$ . The training/validation/testing sets sizes are 3000/800/800 respectively.
D.4 Long-term recurrent forecasting