Non-stationary Equivariant Graph Neural Networks for Physical Dynamics Simulation

Thank you for your thorough and constructive feedback. We appreciate your time and insights, and will address all points below. Corresponding revisions will be incorporated in the manuscript.

Q1: In Section 2.2, why not directly set n=3?

We acknowledge that setting n=3 in Section 2.2 would be appropriate to maintain consistency with the E(3) group mentioned earlier. We will clarify this choice in the revised manuscript.

Q2: What is difference between $s$, $\vec{s}$ and 𝑆, and what are their corresponding dimensions?

We apologize for the notational confusion. For clarity:

• $\vec{s} \in \mathbb{R}^{n\times T \times 3}$: Frequency information extracted by DFT.
• $S \in \mathbb{R}^{n\times r \times 3}$: Frequency information extracted by PFT.
• $s \in \mathbb{R}^{n\times r}$: Normalized version of $S$.

We realize it may be more appropriate to change $S$ as $\vec{S}$ to avoid misunderstanding, and this will be reflected in our revised manuscript.

Q3: In Equ (7), why does the overall time remain 𝑇 after applying the windowing function?

In signal processing, instead of directly cutting off the signal, windowing functions (e.g., Hamming, Hann) attenuate signal discontinuities at segment boundaries without altering the temporal length $T$. This preserves the time dimension while minimizing spectral leakage, consistent with standard signal processing practice. We will add a brief clarification to the manuscript.

Q4: Please check line 142 and 149, both sentences appear to be incomplete.

Thank you for catching these oversights:

• Line 142: Corrected to "$\bar{\vec{X}}_{i}(t)$ is the average 142 position of all nodes in the graph."
• Line 149: Revised to "However, directly cutting off parts of the trajectory can result in spectral leakage."

W1: The methodology section lacks clarity

• Symbol usage should be carefully reviewed for consistency We will make the correponding changes mentioned above.

• How the PFT module is integrated with the NS pooling module We combine these modules by the flow of Equation 3-5. We first extract the invariant features from PFT, and using EGNN to update the coordinates $\hat{X}$, then NS-Pooling is applied to make the final predication.

W1: The methodology section lacks clarity We will improve clarity by addressing both aspects raised:

1. Symbol consistency: We will implement all notational adjustments discussed previously (particularly in Q2), ensuring uniform symbol usage throughout.
2. Module integration: The PFT and NS-pooling modules connect through(Equation 3-5):
   • PFT first extracts invariant features
   • These features update coordinates ($\hat{X}$) via EGNN
   • NS-pooling then generates final predictions We will add a clear pipeline description in Section 3.3 to explicitly articulate this workflow.

W2: An ablation study is necessary to validate the individual contributions of the PFT and NS pooling modules.

We presented ablation studies in:

• Table 6 (PFT): Replacing PFT with standard DFT increases relative MSE by 50.14%, confirming PFT’s necessity.
• Appendix C.1, Table 11 (NS-pooling): Substituting ESTAG’s pooling (equivariant but stationary) improves relative MSE by 24.36%, validating the importance of NS-pooling.

Notably, both ablated variants still outperform all baselines, underscoring the framework’s robustness.

W3: The reported experimental results for ESTAG are weaker than those in the original paper. The authors should clarify this discrepancy.

The discrepancy of ESTAG stems from different experimental settings. Our main results use a more challenging multi-step forecasting setting than ESTAG’s original work. For direct comparison, we include results under identical settings in Appendix C.2 (Table 12), where our implementation matches ESTAG’s reported performance. In this setting, NS-EGNN still outperforms ESTAG in 7 out of 8 cases, and has 10.3% relative performance improvement in average. We will emphasize this distinction in the revised version.

We sincerely appreciate your valuable suggestions, which have strengthened our work. All changes will be implemented in the revised manuscript.

Thank you for your feedback and suggestions. We appreciate the opportunity to address your concerns. Below are our point-by-point responses:

W1: Limited Conceptual Novelty

Our work makes two key conceptual contributions:

1. We identify and formally address the previously overlooked non-stationarity in physical dynamical systems;
2. Due to the non-stationarity, a naive FT is not sufficient to capture the comprehensive frequency information. We introduce windowed FT to capture the dynamic frequency.
3. We establish novel theoretical foundations (Section 3.1.3) for integrating windowed FT within equivariant frameworks, which is a non-trivial extension requiring rigorous proof.

We will further clarify these contributions in our manuscript revision.

W2.1: Incomplete baselines

We appreciate this valuable suggestion. To the best of our knowledge, PG-ODE is not currently open-sourced (we would be happy to include it as a baseline if codes are available), which makes direct comparison challenging. Recognizing the importance of non-stationary baselines, we have included another well-known non-stationary method, Non-stationary Transformer [1], for comparison. The results demonstrate NS-EGNN's significant advantages over existing non-stationary approach:

Note: All values are MSE (lower is better)

This performance gap occurs because standard normalization breaks physical symmetry priors, inducing severe overfitting. We will include this baseline and emphasize this analysis in our revision. Additionally, we would also be pleased to include additional baselines within this realm.

W2.2: Issue of AGL-STAN implemetation

We confirm our AGL-STAN implementation was directly adapted from the official ESTAG repository without any modification. As noted in the ESTAG paper, AGL-STAN underperforms equivariant methods due to its non-equivariant architecture. The significant deviation in AGL-STAN's results arises because it lacks equivariance considerations, making it prone to overfitting in physical dynamics scenarios. We will clarify this methodological distinction in the Experiment section.

Limitation Discussion

Thank you for highlighting this. We have discussed: (1) Broader limitations in Appendix D, and (2) NS-EGNN's specific constraints regarding complex graph relations (Section 3.2). In revision, we will expand these discussions and explicitly note that advanced graph backbones may be required for highly intricate systems. Furthermore, we maintain that these limitations do not diminish our core contribution of addressing non-stationarity in physical dynamic systems.

We thank you again for your valuable insights. All suggested improvements will be reflected in the revised manuscript.

[1]. Non-stationary Transformers: Exploring the Stationarity in Time Series Forecasting

We sincerely thank the reviewer for the thorough evaluation and constructive feedback. We have carefully addressed each point below and will incorporate all changes in the final manuscript.

W1: The paper presentation is not great.

We apologize for any typographical errors in the current draft. We will conduct thorough proofreading to enhance readability and ensure all notations are consistent in the revised version.

W2: The description of the method in Section 3 is not very clear

We regret any confusion caused by the notation. To clarify:

• Equation 2: $\vec{X}$ will be updated to $\vec{X}(t)$.
• Equations 3-5 will be revised for consistency:

$s = \text{PFT}(\{{\vec{\mathbf{X}}(t)}\}_{t=0}^{T})$

$(h, s, \{\vec{\mathbf{X}}(t)^{(L)}\}\_{t=0}^{T}) = \text{EGNN}(h, s, \{\vec{\mathbf{X}}(t)\}\_{t=0}^{T})$

$\vec{\mathbf{X}}^* = \text{NS-Pooling}(\{{\vec{\mathbf{X}}(t)}\}_{t=0}^{T})$

Here, ${\mathbf{X}}(t)^{(L)}$ denotes the $L$-th EGNN layer output, and $\vec{\mathbf{X}}^*$ is the final pooled representation.

• NS-Pooling remains conceptual as represented in Equation 10.
• $\vec{X}^*$ is further utilized to calculate the MSE loss (Section 3.1, line 135).

We are grateful for this feedback, which has strengthened our methodological exposition.

W3: The motivation of the design choices:

Since there are several sub-questions in this part. We carefully address each question individually.

W3.1 What is the performance if we skip the FT entirely and just use encoded patches?

We further conduct the ablation study that completely omitting the FT. Please refer to our response of Q3.2 below.

W3.2 Could we use something temporally localized as well like wavelet transform?

No, we cannot directly employ wavelet transform in our framework. While wavelet transforms are valuable for localized temporal analysis, they present challenges in equivariant settings:

• The calculation of wavelet basis requires the operations of scaling and translation, disrupting the E(3) equivariance.
• Subsampling in discrete wavelets breaks translation equivariance.

Our PFT approach maintains strict E(3) equivariance required for physical dynamics. That said, we recognize wavelet transforms as a powerful alternative for non-stationary signal analysis. Future research directions could investigate modified wavelet constructions that maintain equivariance properties, and we will explicitly include this promising avenue in our future work discussion.

W3.3 How much frequency information is preserved?

Preservation depends critically on:

1. Hop length reduction to increase window overlap [1]
2. Diverse window lengths to capture multi-resolution spectra [2]

Practically, high overlap combined with multi-scale sampling enables comprehensive frequency retention.

W3.4 Could we use other norms directly on the original signal?

Yes, we include an equivariant normalization method as a baseline in Q2.

W3.5 Is the Multi-scale PFT new or an existing idea?

Yes, multi-scale PFT is a new idea. While windowed Fourier analysis exists in signal processing field (e.g., wavelet transforms use a similar ideology), our contributions are:

1. First adaptation to E(3)-equivariant machine learning scenarios for non-stationary dynamics
2. Providing the method to integrate PFT with geometric GNNs while keeping equivariance of the framework
3. Flexible window-length mechanism surpassing wavelet constraints

W3.6 Is there theoretical support for multi-scale PFT?

Regarding the equivariance property of the multi-scale PFT, we provide a rigorous proof of Theorem 3.2. From the perspective from spectral retention, empirically, we observe that NS-EGNN with diverse window lengths and high overlap rates captures more comprehensive spectral features. However, while we can experimentally demonstrate multi-scale PFT's effectiveness, we currently lack a theoretical framework to quantitatively analyze spectral information retention.

W3.7 Does it correspond in some way to the scale-translation group?

No, PFT does not demonstrate scale equivariance but E(3) equivariance (including translation). As established in Theorem 3.2, PFT is E(3)-equivariant (translation/rotation). Scale equivariance doesn't hold: scaling trajectories proportionally scales $s$, breaking EGNN equivariance.

From an application perspective, we argue that NS-EGNN should not inherently possess scale equivariance. This is because changes in scale (manifested as altered edge lengths) typically correspond to meaningful physical property changes in the system.

W4: Visualize the meaning of p-value and ADF Mean

To give a brief understanding of ADF test, we here provide two simplest examples to illustrate the calculation of ADF value and p-value. However, regretfully, since we are not able to upload image, we would include the illustration of example processes and detailed calculation in the Appendix.

1. Process Generation:
   • Stationary: $X_t = 0.6X_{t-1} + e_t, e_t \in \mathcal{N}(0, 1)$
   • Non-stationary: $X_t = X_{t-1} + e_t, e_t \in \mathcal{N}(0, 1)$
2. ADF Test:
   • Regression: $\Delta X_t = \alpha + \beta X_{t-1} + \gamma \Delta X_{t-1} + e_t$
   • Test statistic: $\text{ADF} = \hat{\beta}/\text{SE}(\hat{\beta})$
     • Stationary: $\hat{\beta} = -0.4 \rightarrow \text{ADF} = -4.3$
     • Non-stationary: $\hat{\beta} = -0.02 \rightarrow \text{ADF} = -1.1$
3. p-value Interpretation:
   • Stationary: p = 0.001 (strong evidence against null)
   • Non-stationary: p = 0.72 (weak evidence against null)

Q1: 'm not sure about the claim that "most existing model focus on single-step frame-to-frame forecasting." In my experience (in fluids or robotics ,e.g.), many time series prediction method encode many past frames and predict many future frames. Maybe this is just in MD?

Our statement primarily refers to the specific context of geometric graph neural networks for physical dynamics simulation, where many existing approaches (particularly those cited in our related works section) do indeed focus on single-step prediction.

We will clarify the domain specificity in our revised claims.

Q2: Is normalization incompatible with equivariance?

Indeed, as demonstrated by recent work [3], naive normalization method is incompatible with equivariance. However, we also find a more advanced normalization method maintaining equivariance, and we have included it as our baseline [3] (without FT). The results are below:

The results clearly demonstrate that normalization significantly degrades model performance (average error increases from 0.353 to 2.078). This performance deterioration stems from normalization's dual effect: while it helps stationarize the trajectories, it simultaneously removes critical temporal patterns that are essential for accurate predictions.

Q3.1: Do the baselines encode multiple past time steps

Yes, ESTAG employs an attention layer to encode the multiple past time steps. Other baselines are not designed for time-series in their original paper, and these baselines are modified by encoding the past time steps by linear layers.

Q3.2: An ablation where the FT is omitted completely

We further conduct the ablation study that completely omitting the FT. The results can be presented below:

Surprisingly, we find totally remove FT even can outperform EGNN with FT, which indicates FT cannot accurately extract the intrinsic spectral information. The experiments demonstrates this inaccurate spectral feature also harms the convergence of the model.

Q4: How the NS-pooling operation works with the model to make the output stationary?

We apologize for any misundetstanding. We clarify:

1. The output of EGNN is $\{{\vec{\mathbf{X}}(t)}\}_{t=0}^{T}$.
2. NS-Pooling derives stationarized $\Delta \vec{\mathbf{X}}$ and $\Delta^2 \vec{\mathbf{X}}$ from the output of EGNN $\{{\vec{\mathbf{X}}(t)}\}_{t=0}^{T}$ (Equation 9).
3. This stationarization improves NN generalization by providing NN more stationary features.

Please kindly note that our goal is stationarizing inputs, not outputs (predictions actually should be inherently non-stationary).

Q5: On line 177, is the FFT O(N logN)?

Correct. We will update this in the manuscript. Thanks for pointing this typo.

Q6: Ablation of multi-scale PFT

We employ reduce utilizing multiple windows to only utilize single window, the results can be found below:

Multi-scale PFT consistently outperforms single-scale, validating our design.

All discussed modifications will be implemented in the revision. We thank the reviewer again for their insightful feedback that has significantly improved our work.

[1]. Window and Overlap Processing Effects on Power Estimates from Spectra. 2000

[2]. A Simple Adjustable Window Algorithm to Improve FFT Measurements. 2002

[3]. Towards Geometric Normalization Techniques in SE(3) Equivariant Graph Neural Networks for Physical Dynamics Simulations. IJCAI 2024

Thanks for your valuable feedback! Please kindly find the detailed responses below.

Q1.1: Why is the DFT considered as "static frequency," while PFT is regarded as "dynamic frequency"?

Thank you for raising this point. The distinction stems from their temporal resolution:

• DFT computes the global frequency spectrum of an entire signal, implicitly assuming stationarity (hence "static").
• PFT analyzes short overlapping segments, capturing time-localized frequency evolution (thus "dynamic"). We will clarify this terminology in Section 3.1 to avoid ambiguity.

Q1.2: Is there any experimental evidence proving that PFT performs better than DFT?

Yes, in ablation study (Section 4.5), we conduct an experiment to replace PFT with DFT, and we found the average performance decreases, (Averaged prediction error from 0.353 to 0.530).

Q2: What is the additional time consumption introduced by PFT?

We compared training times for 5 epochs across variants. The experiments are conducted on the machine introduced in Appendix B.3. As shown below, PFT adds minimal computational overhead:

Table: Training time (ms) per 5 epochs.

Q3.1: The meaning of Equation (10)

We apologize for any confusion. Equation (10) $\vec{\mathbf{X}}_{i}^{*} = [\Delta \vec{\mathbf{X}}_i, \Delta^2 \vec{\mathbf{X}}_i] \cdot \mathbf{\gamma} + \vec{\mathbf{X}}_{i}^{(L)}(T-1)$ integrates two components:

1. Stationarized features: $\Delta \vec{\mathbf{X}}_i$ (first-order differences from Eq. 9) and $\Delta^2 \vec{\mathbf{X}}_i$ (second-order differences) are projected via a linear layer $\mathbf{\gamma}$.
2. Temporal context: $\vec{\mathbf{X}}_{i}^{(L)}(T-1)$ denotes the EGNN module's final-layer output at $t=T-1$.

This approach reduces non-stationarity by predicting incremental changes rather than absolute states. For non-stationary data, these incremental values may be more predictable. For example, in molecular dynamics, the model can more effectively learn energy changes within the system than directly learn features from absolute energy values.

We will refine the explanation in Section 3.3.

Q3.2: Predicting the final equilibrium conformation has also been proposed in many previous papers, such as eSCN [1] in their OC20 IS2RS [2] experiments. Should these works be discussed?

We appreciate this valuable suggestion. While eSCN advances equivariant representations, its primary focus on equilibrium conformation prediction makes direct adaptation to molecular dynamics simulations non-trivial.

Key Difference: Dynamics vs. Static Structures - MD simulations require modeling molecular motion over time, whereas equilibrium predictions provide only static "snapshot" structures.

Additionally, eSCN analyzes individual frame conformations rather than trajectory-based forecasting. Since our baseline comparisons already include similar high-representation models (ST-TFN, ST-SE(3)-Transformer), we focused on directly comparable approaches. We will expand the discussion of conformation prediction's relation to our task in Section 5.

We sincerely appreciate your valuable suggestions. We will include above discussions into our revised manuscript to enhance clarity.