SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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 + \Delta T, T + 2\Delta T, \dots, T + 10\Delta T$  (the value of  $\Delta 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.
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, 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.
Figure 6: The movement trajectory of 8 molecules in MD17.
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-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.
Figure 8: 3D visualization of URACIL for comparison between ESTAG (left, MSE=5.983 × 10 $^{-5}$ ) and ST_EGNN (right, MSE=6.393 × 10 $^{-4}$ ). The ground truths are in red while the predicted states are in blue. The molecule predicted by ESTAG is better aligned with the ground truth, particularly for the regions highlighted by circles.
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  $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.
Table 6: MSE on Ethanol w.r.t. the number of previous timestamps  $T$
T	3	5	8	10	20	30	40	50	60	70
MSE (×10-4)	1.760	1.395	1.064	0.990	0.924	1.017	0.962	1.009	0.992	1.019
Figure 10: Comparison on Motion walk subject between ESTAG (left, MSE=0.0048) and ST_EGNN (right, MSE=0.0811). The ground truths are in red while the predicted states are in blue.
Figure 11: Comparison on Motion basketball subject between ESTAG (left, MSE=0.0749) and ST_EGNN (right, MSE=2.6380). The ground truths are in red while the predicted states are in blue.
Figure 12: MSE on Ethanol w.r.t. the number of previous timestamps  $T$
Table 7: Ablation studies  $(\times 10^{-3})$  of learnable parameter  $w_{k}$  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 wk	0.071	0.003	0.102	0.104	0.081	0.081	0.079	0.069
Table 8: Ablation study on Protein dataset. Table 9: Ablation study on Motion dataset.
Model	MSE
ESTAG	1.471
w/o EDFT	1.490
w/o Attention	1.493
w/o Equivariance	2.011
w/o Temporal	1.472
Model	MSE (×10-1)
ESTAG	0.040
w/o EDFT	0.053
w/o Attention	0.047
w/o Equivariance	0.702
w/o Temporal	0.044
F Codes
The codes of ESTAG are available at: https://github.com/ManlioWu/ESTAG.
Input: Initial historical graph series  $\{(H^{(0)}(t),\vec{X}^{(0)}(t),A\}_{t = 0}^{T - 1}$
for  $i = 1$  to  $N$  do
Equivariant Discrete Fourier Transform (EDFT):
$$
\vec {\boldsymbol {f}} _ {i} (k) = \sum_ {t = 0} ^ {T - 1} e ^ {- i ^ {\prime} \frac {2 \pi}{T} k t} \left(\vec {\boldsymbol {x}} _ {i} (t) - \overline {{\vec {\boldsymbol {x}} (t)}}\right), \tag {15}
$$
$$
\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 |, \tag {16}
$$
$$
\boldsymbol {c} _ {i} (k) = w _ {k} \left(\boldsymbol {h} _ {i}\right) \| \vec {\boldsymbol {f}} _ {i} ^ {*} (k) \| ^ {2}. \tag {17}
$$
for  $l = 1$  to  $L$  do
Equivariant Spatial Module (ESM):
$$
\boldsymbol {m} _ {i j} = \phi_ {m} \left(\boldsymbol {h} _ {i} ^ {(l - 1)} (t), \boldsymbol {h} _ {j} ^ {(l - 1)} (t), \| \vec {\boldsymbol {x}} _ {i j} ^ {(l - 1)} (t) \| ^ {2}, \boldsymbol {A} _ {i j}\right), \tag {18}
$$
$$
\boldsymbol {h} _ {i} ^ {(l - 0. 5)} (t) = \phi_ {h} \left(\boldsymbol {h} _ {i} ^ {(l - 1)} (t), \boldsymbol {c} _ {i} (k), \sum_ {j \neq i} \boldsymbol {m} _ {i j}\right), \tag {19}
$$
$$
\vec {\boldsymbol {a}} _ {i} ^ {(l - 0. 5)} (t) = \frac {1}{| \mathcal {N} (i) |} \sum_ {j \in \mathcal {N} (i)} \vec {\boldsymbol {x}} _ {i j} ^ {(l - 0. 5)} (t) \phi_ {x} (\boldsymbol {m} _ {i j}), \tag {20}
$$
$$
\vec {\boldsymbol {x}} _ {i} ^ {(l - 0. 5)} (t) = \vec {\boldsymbol {x}} _ {i} ^ {(l - 1)} (t) + \vec {\boldsymbol {a}} _ {i} ^ {(l - 0. 5)} (t). \tag {21}
$$
Equivariant Temporal Module (ETM)
$$
\alpha_ {i} ^ {(l - 0. 5)} (t s) = \frac {\exp \left(\boldsymbol {q} _ {i} ^ {(l - 0 . 5)} (t) ^ {\top} \boldsymbol {k} _ {i} ^ {(l - 0 . 5)} (s)\right)}{\sum_ {s = 0} ^ {t} \exp \left(\boldsymbol {q} _ {i} ^ {(l - 0 . 5)} (t) ^ {\top} \boldsymbol {k} _ {i} ^ {(l - 0 . 5)} (s)\right)}, \tag {22}
$$
$$
\boldsymbol {h} _ {i} ^ {(l)} (t) = \boldsymbol {h} _ {i} ^ {(l - 0. 5)} (t) + \sum_ {s = 0} ^ {t} \alpha_ {i} ^ {(l - 0. 5)} (t s) \boldsymbol {v} _ {i} ^ {(l - 0. 5)} (s), \tag {23}
$$
$$
\vec {\boldsymbol {x}} _ {i} ^ {(l)} (t) = \vec {\boldsymbol {x}} _ {i} ^ {(l - 0. 5)} (t) + \sum_ {s = 0} ^ {t} \alpha_ {i} ^ {(l - 0. 5)} (t s) \vec {\boldsymbol {x}} _ {i} ^ {(l - 0. 5)} (t s) \phi_ {x} \left(\boldsymbol {v} _ {i} ^ {(l - 0. 5)} (s)\right), \tag {24}
$$
$$
\boldsymbol {q} _ {i} ^ {(l - 0. 5)} (t) = \phi_ {q} \left(\boldsymbol {h} _ {i} ^ {(l - 0. 5)} (t)\right), \tag {25}
$$
$$
\boldsymbol {k} _ {i} ^ {(l - 0. 5)} (t) = \phi_ {k} \left(\boldsymbol {h} _ {i} ^ {(l - 0. 5)} (t)\right), \tag {26}
$$
$$
\boldsymbol {v} _ {i} ^ {(l - 0. 5)} (t) = \phi_ {v} \left(\boldsymbol {h} _ {i} ^ {(l - 0. 5)} (t)\right). \tag {27}
$$
Equivariant linear pooling:
$$
\vec {\boldsymbol {x}} _ {i} ^ {*} (T) = \hat {\boldsymbol {X}} _ {i} \boldsymbol {w} + \vec {\boldsymbol {x}} _ {i} ^ {(L)} (T - 1), \tag {28}
$$
where  $\hat{\pmb{X}}_i = [\vec{\pmb{x}}_i^{(L)}(0) - \vec{\pmb{x}}_i^{(L)}(T - 1),\vec{\pmb{x}}_i^{(L)}(1) - \vec{\pmb{x}}_i^{(L)}(T - 1),\dots ,\vec{\pmb{x}}_i^{(L)}(T - 2) - \vec{\pmb{x}}_i^{(L)}(T - 1)]$
Output:  $\vec{X}^* (T)$
Algorithm 1 Equivariant Spatio-Temporal Attentive Graph Networks (ESTAG)