SPARK: Physics-Guided Quantitative Augmentation for Dynamical System Modeling

Dear Reviewer pH3K,

We sincerely appreciate the time you’ve dedicated to reviewing our paper, as well as your valuable insights and support. Your positive feedback is highly motivating for us. Below, we address your primary concern and offer further clarification.

---

Q1. The motivation for using each component in SPARK can be further clarified. The paper will benefit from discussing the interconnection between each network component.

A1. Thank you for your comment. In this work, we design a unified plugin SPARK with three basic modules:

• Physics-incorporated data compression. We integrate physical parameters, boundary information and input observations into latent space through position encoding and channel attention.

• Memory bank construction. We then pre-train a discrete memory bank through GNN-based reconstruction and discrete quantization mechanism.

• Downstream augmentation and prediction. We froze the memory bank's weights and use it to realize augmentation, which introduces diversity into the training set. Further, we design a fourier-enhanced graph ODE to accurately predict complex dynamics.

These three designs contribute jointly to the high accuracy and strong generalization ability in various environments. We will add disscussions about components' interconnection in the revised manuscript.

Q2. It would also be good to have ablation studies on incorporated physics. The authors may consider reducing physical information (i.e., boundary information and physical parameters) for pre-training. Then, we can see the contribution of each physical component.

A2. Thank you for your valuable feedback. We add ablation experiments on incorporated physics by reducing different physical information. The experiments are conducted on OOD scenarios, and the results are shown below. The results demonstrate the effectiveness of each physical component. We will include it in our revised manuscript.

Dear Reviewer bjTD,

We sincerely appreciate the time you’ve dedicated to reviewing our paper, as well as your valuable insights and support. Below, we address your primary concern and offer further clarification.

---

Q1. The symbols and formulas appear to be somewhat disorganized, which makes it difficult for readers to understand the meaning. Clear definitions and a more structured presentation of the equations would greatly enhance the paper's accessibility and overall readability.

A1. Thank you for your valuable feedback. To better facilitate the understanding of our paper, we have made the following modifications:

• Problem definition. We have refined the problem definition and ensured consistency throughout the manuscript. As follows:

"Given a dynamical system governed by physical laws such as PDEs, we aim to enhance prediction using autoencoder reconstruction and discrete quantization. We have $N$ observation points in the domain $\Omega$, located at $\mathbf{s} = \{\mathbf{s}\_1, \cdots, \mathbf{s}\_N\}$, where $\mathbf{s}\_i \in \mathbb{R}^{d\_s}$. At time step $t$, the observations are $\mathcal{X}^t = \{\mathcal{X}\_1^t, \cdots, \mathcal{X}\_N^t\}$, where $\mathcal{X}\_i^t \in \mathbb{R}^{d}$ and $d$ represents the number of observation channels. Boundary information and physical parameters affect the dynamical system, leading to different conditions and distribution shifts. We first employ reconstruction model and construct a discrete memory bank to compress and store physical prior information. Then, given historical observation sequences {$\{\mathcal{X}\_i^{-T\_0+1:0}\}$}$\_{i=1}^N$, our goal is to use the pre-trained memory bank for data augmentation and predict future observations {$\{\mathcal{Y}\_i^{1:T}\}$}$\_{i=1}^N$ at each observation point."

• Symbols and formulas. We thoroughly review all symbols and formulas in the manuscript to ensure their meanings are precise and clear. For instance, we make the following modifications:

$\boldsymbol{u}_i = \text{Proj} \left( \mathcal{X}_i , \boldsymbol{p}^{rel}_i \right) \quad \text{with} \quad \boldsymbol{p}^{rel}_i = \phi\left( \mathbf{s}_i, \boldsymbol{p}^{boun}_i \right), \quad (1)$

$\mathcal{L}\_{pre}=\frac{1}{T N} \sum\_{t=1}^T \sum\_{i=1}^N\left(\hat{\mathcal{X}}\_{i}^{t}-\mathcal{X}\_{i}^{t}\right)^2+ \frac{1}{T N} \sum\_{t=1}^T \sum\_{i=1}^N\left (\mu \left\|\boldsymbol{h}\_{i}^{t}-\mathbf{s g}[\boldsymbol{e}]\right\|\_2^2+\gamma\left\|\mathbf{s g}\left[\boldsymbol{h}\_{i}^{t}\right]-\boldsymbol{e}\right\|\_2^2\right ), \quad (6)$

$\boldsymbol{q}\_{i}=\frac1{T\_0}\sum\_{t=1}^{T\_0}\delta(\alpha\_{i}^{t} \cdot \boldsymbol{v}\_{i}^{t}),\quad \alpha\_{i}^{t}= \left(\boldsymbol{v}\_{i}^{t} \right )^{T} \cdot \mathrm{tanh}\left(\left(\frac{1}{T\_0} \sum\_{t=1}^{T\_0} \boldsymbol{v}\_{i}^{t}\right)W\_{\alpha}\right), \quad (8)$

$\mathcal{L}\_{\text{dyn}} = \frac{1}{TN} \sum\_{i=1}^T \sum\_{i=1}^N \|\hat{\mathcal{Y}}\_{i}^{t} - \mathcal{Y}\_{i}^{t}\|\_2^2 + \lambda\_{\text{reg}} \mathcal{R}(\theta). \quad (11)$

Dear Reviewer 89BR,

Thank you for your valuable feedback on our manuscript! We have taken your comments seriously and have made the necessary revisions and additions to address the concerns raised.

---

Q1. Problem about originality.

A1. Thank you for your comment. We think our SPARK differs significantly from DGODE[1].

• Difference with discrete coding. DGODE is an end-to-end framework. It utilizes discrete "codebank" for disentangling environment features and minimizing their impact on node representations. While, SPARK is a upstream-downstream paradigm. In upstream phase, except observations, we incorporate physical parameters and boundary information to train a discrete, physics-rich memory bank. In downstream phase, we forze memory bank's weights and use it to enable targeted augmentation.

• Difference with graph ODE. The use of graph ODEs is motivated by their ability to handle irregular data and efficiently perform multi-step temporal extrapolation. However, SPARK and DGODE differ significantly in their details: (1) DGODE uses an RNN architecture to compress historical states, while we use an attention mechanism to compress them into an initial state; (2) we incorporate Fourier blocks into the Graph ODE to enhance the capture of global spectral features, improving spatiotemporal generalization.

To further address your concerns, we run DGODE's open-source code and conduct comparative experiments in both non-OOD (ID) and OOD scenarios. The results shown below indicate that SPARK performs better. We will include these in our revised version.

---

Q2. The algorithm is not clearly presented.

A2. Thank you for your detailed comments. We have taken steps to address these issues in the revised manuscript.

• About boundary representation. In our paper, the "real boundary" ($\boldsymbol{p}^{boun}_i$) refers to the relative positional relationship between node $i$ and the nearest boundary point. We combine this with the node $i$'s coordinate information $\mathbf{s}_i$ to obtain the position encoding $\boldsymbol{p}^{rel}_i$. We have made detailed corrections to Equation (1), as shown below. For a more intuitive understanding, we use the ERA5 dataset as an example and visualize the boundary information in Appendix I.

$\boldsymbol{u}_i = \text{Proj} \left( \mathcal{X}_i , \boldsymbol{p}^{rel}_i \right) \quad \text{with} \quad \boldsymbol{p}^{rel}_i = \phi\left( \mathbf{s}_i, \boldsymbol{p}^{boun}_i \right).$

• Time index. We consistently follow the notation in the Problem Definition. The length of history observations is $T_0$, and the length of future predictions is $T$. We have revised Equations (8) and (11) accordingly, as shown below.

\alpha_{i}^{t}= \left(\mathcal{v}\_{i}^{t} \right )^{T} \cdot \mathrm{tanh}\left(\left(\frac{1}{T\_0} \sum\{t=1}^{T\_0} \mathcal{v}\_{i}^{t}\right)W\_{\alpha}\right), \quad (8)

$\mathcal{L}\_{\text{dyn}} = \frac{1}{TN} \sum\_{i=1}^T \sum\_{i=1}^N \|\hat{\mathcal{Y}}\_{i}^{t} - \mathcal{Y}\_{i}^{t}\|\_2^2 + \lambda\_{\text{reg}} \mathcal{R}(\theta). \quad (11)$

• Index notation. We have revised the index notation in the pretraining loss equation (6) to ensure it is accurate and meaningful, which is shown below.

$\mathcal{L}\_{pre}=\frac{1}{T N} \sum\_{t=1}^T \sum\_{i=1}^N\left(\hat{\mathcal{X}}\_{i}^{t}-\mathcal{X}\_{i}^{t}\right)^2+ \frac{1}{T N} \sum\_{t=1}^T \sum\_{i=1}^N\left (\mu \left\|\mathcal{h}\_{i}^{t}-\mathbf{s g}[\mathcal{e}]\right\|\_2^2+\gamma\left\|\mathbf{s g}\left[\mathcal{h}\_{i}^{t}\right]-\mathcal{e}\right\|\_2^2\right ). \quad (6)$

Further, we have thoroughly reviewed the entire manuscript and standardized the writing to ensure consistency in symbols and formulas.

---

[1] Wu H, et al. "Prometheus: Out-of-distribution Fluid Dynamics Modeling with Disentangled Graph ODE." ICML2024.

Q3. It's unclear how the discrete memory bank is built. What are the e_i in E and how are they constructed?

A3. Thank you for your feedback. SPARK builds its discrete memory bank using a training strategy inspired by VQ-VAE[1]. Each $e_i$ in the memory bank $E =${$\{e_1, e_2, \dots, e_M\}$}$\in \mathbb{R}^{M \times D}$ represents an embedding vector, where $M$ is the fixed size of the memory bank, manually set as a hyperparameter. These embeddings are initialized randomly at the start of training.

Q4. How are the with and without OOD datasets constructed in the experiments? Is there an explanation for why SPARK achieved better performance on OOD cases even than other models did on non-OOD cases?

A4. Thank you for your comment.

• Experimental settings. For dataset setting, we propose that training and testing in the in-domain parameters is called w/o OOD experiments, while training in the in-domain parameters and testing in the out-domain parameters is called w/ OOD experiments. Here we present the in-domain and out-domain parameters for different benchmarks in the table below.

• Explanation for SPARK's better performance. Our SPARK is specifically designed for OOD problem. SPARK's physics-guided enhancement improves generalization by leveraging physical priors. This reduces sensitivity to distribution shifts. Additionally, the fourier-enhanced graph ODE module provides a robust mechanism for prediction, outperforming other baselines. To address your concern, we use three models specialized in OOD dynamical system modeling (LEADS[2], CODA[3], NUWA[4]), along with FNO, for comparison. The results shown below indicate that OOD-specific models outperform FNO in both OOD and non-OOD scenarios, with SPARK achieving the best performance. We will include these in our revised version.

---

[1] Van Den Oord A, et al. "Neural discrete representation learning." NeurIPS2017.

[2] Kirchmeyer M, et al. "Generalizing to new physical systems via context-informed dynamics model." ICML 2022.

[3] Yin Y, et al. "LEADS: Learning dynamical systems that generalize across environments." NeurIPS2021.

[4] Wang K, et al. "NuwaDynamics: Discovering and Updating in Causal Spatio-Temporal Modeling." ICLR2024.

Dear Reviewer 89BR,

Thank you for your valuable feedback and we will polish the paper based on your suggestions. We greatly appreciate your recognition of our extensive experiments and your support in raising the score!

Warm regards,

Authors

Dear Reviewer Rucb,

We sincerely appreciate the time you’ve dedicated to reviewing our paper, as well as your valuable insights and support. Your positive feedback is highly motivating for us. Below, we address your primary concern and offer further clarification.

---

Q1. It is unclear how much each strategy contributes to the success of the model, and some ablation studies would be useful.

A1. Thanks for your valuable feedback. To further demonstrate the contribution of each strategy, we conduct ablation experiments with five model variants. The experiments are conducted on Prometheus and Navier–Stokes datasets with OOD scenarios, and the results are shown below.

As observed, removing physical parameters or boundary conditions during pretraining leads to a performance decline, with an even greater drop when the memory bank is not used. This validates the effectiveness of physical compression when addressing OOD problems. We will include these in our revised manuscript.

---

Q2. Can the data bank be used for direct retrieval-augmentation?

A2. Thank you for your comment. We acknowledge that the memory bank in our SPARK framework can be used for direct retrieval augmentation. After pre-training, the memory bank’s parameters are frozen, allowing new samples to retrieve physics-rich embeddings in the bank for augmentation. For convenient understanding, we have illustrated this with diagram in Appendix H.9.