We are truly grateful for the time you have taken to review our paper, your insightful comments and support. Your positive feedback is incredibly encouraging for us! In the following response, we would like to address your major concern and provide additional clarification.

Q1. The contribution of the proposed method in dealing with the OOD problem needs to be further clarified since the advantages of PURE over the previous efforts, such as Refs. [7, 67, 14, 72], etc., to address the OOD problem are not listed.

A1. Thanks for your comment. Compared with previous OOD methods, our contribution can be listed as follows:

• Underexplored Scenarios. Our work studies an underexplored yet practical problem of OOD generalization in fluid dynamic modeling while previous OOD methods usually study the OOD problem in image classification scenarios.
• Innovative Methodlogy. Our work not only learns time-evolving prompts using context mining and a graph ODE but also decouples prompt embeddings and observation embeddings via mutual information for invariant learning to address the OOD problem.
• Theoretical Analysis. We provide a comprehensive theoretical analysis to support our designs, which makes our framework more solid.
• Superior Performance. Comprehensive experiments validate the effectiveness of our method in different challenging settings. In particular, it improves performance by an average of 25.90% in OOD scenarios.

Q2. The writing of the paper needs to be improved. Some of the symbols in the method description section are not defined, e.g., what do P and N in Equation 9 refer to?

A2. Thank you for your comment. $P$ denotes the set collecting all the positive pairs (i.e., from the same trajectory) of observation embeddings and prompt embeddings while $N$ denotes the set collecting all the possible pairs of observation embeddings and prompt embeddings in the dataset. We will include it in our revised version.

Q3. The experiment is not comprehensive enough. (a) The reasons for selecting baselines are not explained. Data augmentation [66, 7], invariant feature learning [39, 69, 38], adversarial training [67, 7], and domain adaptation [32, 14] are mentioned in the paper in related work for solving the OOD problem, but they are not be compared as baselines in the experiment.

A3. Thank you for your comment. We have added three more baselines LEADS [1], CODA [2] and NUWA [3] for performance comparison. The results show that our method is superior to these baselines, which validates the effectiveness of our method when addressing OOD problems. In addition, the other references focus on the image classification problem, which cannot be adopted to solve our fluid dynamics modeling problem. We will include it in our revised version.

Q4. (b) The experiments in this paper do not state whether noisy data are considered.

A4. Thank you for your comment. We have added experiments with noisy data to evaluate the robustness of our method. The results are shown below, which validate that our method is more robust to interference from noisy data. We will include it in our revised version.

Q5. (c) The authors just give a brief description of the results without analyzing the reasons behind the high performance.

A5. Thank you for your comment. The potential reason for our high performance is that (1) our model enhances model invariance across different distributions through decoupling prompt embeddings and observation embeddings via mutual information which results in high generalizationability to different environments; (2) our model utilizes multi-view context mining and graph ODE to extract prompt embeddings, which capture environment information accurately. We will include it in our revised version.

Reference

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

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

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

We will also add your suggestion about future works to our revised version. Thanks again for appreciating our work and for your constructive suggestions. Please let us know if you have further questions.

Dear Reviewer,

We summarized our rebuttal content at your convenience as follows:

• We have included a detailed explanation of our contribution.
• We have included more competing baselines to demonstrate the superiority of our approach.
• We have included more real-world settings such as noisy scenarios.
• We have included more analysis about our superiority.

We will also add your suggestion about future works to our revised version. Thanks again for appreciating our work and for your constructive suggestions. Please let us know if you have further questions.

Best,

the Authors

We are truly grateful for the time you have taken to review our paper and your insightful review. Here we address your comments in the following.

Q1. My major concern with the paper is that the OOD challenge in dynamics modeling is not well-formulated. The paper describes the OOD scenario verbally as "different dynamical systems could involve different parameters in underlying rules" and "during long-term auto-regressive forecasting, the input data distribution could vary hugely during temporal evolution," which is straightforward and easy to understand. However, the mathematical formulation of these scenarios is absent. This formulation should be the foundational basis of the topic, as we need to clearly define the problem before addressing it.

A1. Thank you for your comment. In dynamical systems, the OOD problem studies the prediction performance of models under parameter distributions or environments not seen during training. In formulation, the evolution of dynamical systems is defined by $\frac{du}{dt} = F(u, \xi)$, where $u$ is the observation, and $\xi$ is the system parameter. If these parameters come from a distribution $\xi \sim P(\xi)$, the state trajectory comes from the distribution $u^{1:T_0} \sim P(u^{1:T_0} | \xi)$. Assume we learn a state mapping $f$ from time $u^{1:T_0}$ to $u^{T_0+1:T_0+T}$, i.e., $u^{T_0+1:T_0+T}= f(u^{1:T_0})$. We could have different distributions across training and test datasets, i.e., $P_{\text{train}}(\xi)\neq P_{\text{test}}(\xi)$, which results in $P_{\text{train}}(u^{1:T_0})\neq P_{\text{test}}(u^{1:T_0})$. Moreover, when conducting rollout prediction, we need to feed the output back to the model, i.e., $u^{T_{start}:T_{start}+T-1}= f(u^{T_{start}-T_0:T_{start}-1})$ with $P(u^{1:T_0} | \xi)\neq P(u^{T_{start}-T_0:T_{start}-1}|\xi, T_{start})$, which demonstrates temporal distribution shift. We will include our explanation in the revised version.

Q2. Given the lack of mathematical formulation of the challenge, I find myself lost in the proposed approach section, unsure of the necessity for specific components. While I understand the function of each component, I cannot see why it is needed or which gaps it aims to bridge in the absent mathematical framework.

A2: Thank you for your comment. From our definition, the mapping $u_{output}=f(u_{input})$ could suffer from a serious distribution shift result from different $\xi$ and $T_{start}$, i.e., $P(u_{input}|\xi, T_{start})$. To reduce the impact of distribution shift, we aim to learn invariant observation embeddings $\mu^t$ to environments, i.e., $\xi$ and $T_{start}$, and utilize prompt embeddings $z^t$ to indicate the current environment. The basic idea of our method is to ensure the invariance of observation embeddings for better generalization, i.e., $z^t\perp \mu^t$. Then, the prompt embeddings would be combined to generate the future predictions, i.e., $u_{output} = \phi ([\mu^t,z^t])$. In our framework, a basic model is adopted to generate observation, i.e., $\mu^t=Basic Model(u_{input})$. We adopt context mining and graph ODE to learn time-varying prompt embeddings., i.e., $z^0= Context Mining (u_{input})$, $z^t = GraphODE (z^0,t)$, which can explore the temporal evolution of environments. We will include our explanation in the revised version.

Q3 Why is the proposed method termed "prompt"? Is there a connection to prompt tuning in large language models?

A3. Thank you for your comment. "Prompt" means a supplementary hint to indicate the context, which is incorporated to the input (observation embedding) for better generalization. Generally, it shares the basic idea with prompt tuning, which also aims to incorporate optimal tokens into the input sequence. However, prompt tuning usually utilizes prompts to indicate different tasks in NLP scenarios, our prompt refers to the current environment in fluid dynamical modeling, which determines the future evolution with better generalization. We will include it in our revised version.

Q4. How do you quantify the distribution shift in dynamics modeling? Can you rank the 'difficulty level' of OOD generalization in your experiments and analyze in which scenarios your method stands out and why?

A4. Thank you for your comment. We have added experiments to demonstrate the performance with varying difficulty levels. In particular, we measure the difficulty levels based on the distance between $P_{\text{train}}(\xi)$ and $P_{\text{test}}(\xi)$ and generate three levels on the Prometheus dataset. The compared results are shown below. We can observe that although all the model performs worse in hard scenarios, our method consistently outperforms these baselines. The potential reason is that (1) our model enhances model invariance across different distributions through decoupling prompt embeddings and observation embeddings via mutual information which results in high generalizationability to different environments; (2) our model utilizes multi-view context mining and graph ODE to extract prompt embeddings, which capture environment information accurately. We will include it in our revised version.

Reference

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

In light of these responses, we hope we have addressed your concerns, and hope you will consider raising your score. If there are any additional notable points of concern that we have not yet addressed, please do not hesitate to share them, and we will promptly attend to those points.

Thanks for your feedback, and we have finished the revision based on your helpful suggestions.

Q1. The new information on OOD in dynamic systems is too basic and should be included in the problem setup at the start of the paper.

Thanks for your comment. We have revised the paper based on your suggestions. The revised draft of Sec. 2 (Problem Definition) and Sec. 3.1 (Motivation and Framework Overview) is shown as below:

Sec. 2 Problem Setup

Given a fluid dynamical system, we have $N$ sensors within the domain $\Omega$, with their locations denoted as $x\_1, \cdots, x\_N$, where $x_i \in \mathbb{R}^{d_l}$. The observations at time step $t$ are represented as $s_1^t, \cdots, s_N^t$, where $s_i^t \in \mathbb{R}^{d_o}$ and $d_o$ indicates the number of observation channels. Dynamical systems are governed by underlying system rules, such as PDEs with coefficient $\xi$. Variations in system parameters may lead to different environments, potentially resulting in distribution shifts. In our study, we are provided with historical observation sequences $\{s_i^{1: T_0}\}\_{i=1}^N$ and physical parameters $\xi$ (e.g., coefficients in the PDEs). Our goal is to predict the future observations of each sensor $s_i^{T_0+1: T_0+T}$. In dynamical systems, the out-of-distribution problem examines model performance when predicting under unseen parameter distributions or environments. Let $u^t=[s_1^t,\cdots, s_N^t]$, these systems evolve according to $\frac{d{u}}{dt} = F({u}, {\xi})$, where ${u}$ represents the observations and ${\xi}$ denotes the system parameters. When ${\xi} \sim P({\xi})$, the state trajectory ${u}^{1:T_0}$ follows the distribution $P({u}^{1:T_0} | {\xi})$. Assume we learn a learned mapping function $f$ from ${u}^{1:T_0}$ to ${u}^{T_0+1:T_0+T}$, i.e., ${u}^{T_0+1:T_0+T} = f({u}^{1:T_0})$ and there could be different distributions across training and test datasets, i.e., $P_{\text{train}}({\xi})\neq P_{\text{test}}({\xi})$, which results in $P\_{{train }}\left({u}^{1: T_0}\right) \neq {P}\_{{test }}\left({u}^{1: T_0}\right)$. Moreover, when conducting rollout prediction, we need to feed the output back to the model, i.e., ${u}^{T_{start}:T_{start}+T-1} = f({u}^{T_{start}-T_0:T_{start}-1})$, with $P({u}^{1:T_0} | {\xi})\neq P({u}^{T_{start}-T_0:T_{start}-1}|{\xi}, T_{start})$, which demonstrates temporal distribution shift.

Sec. 3 The Proposed PURE

Sec. 3.1 Motivation and Framework Overview

This paper addresses the challenge of out-of-distribution fluid system modeling, which is complicated by parameter-based and temporal distribution shifts. Specifically, our function $f(\cdot)$ can suffer from a serious distribution shift result from different ${\xi}$ and $T\_{start}$, i.e., $P({u}\_{input}|{\xi},T\_{start})$. To reduce the impact of distribution shift, we aim to learn invariant observation embeddings ${\mu}^t$ to different environments, i.e., ${\xi}$ and $T_{start}$ for better generalization and utilize prompt embeddings ${z}^t$ to indicate the current environment for final prediction. In formulation, we have:

The first term ensures the invariance of observation embeddings by decoupling observation embeddings and prompt embeddings. The second term aims to combine both two embeddings to generate the future predictions. Therefore, we propose a novel approach named PURE as:

where a basic model is adopted to generate observation, and we adopt context mining and graph ODE to learn time-varying prompt embeddings. Given a basic forecasting model (Eqn. 2), our PURE contains three key modules: (1) Multi-view Context Exploration, which explores spatio-temporal data using both the attention mechanism and the frequency domain to initialize prompt embeddings (Eqn. 3). (2) Time-evolving Prompt Learning, which incorporates the interpolation of observation sequences into a graph ODE to learn the evolution of prompt embeddings (Eqn.4). (3) Model Adaptation with Prompt Embeddings, which leverages the time-evolving prompts to mitigate the temporal distribution shifts in fluid dynamics models (Eqn.1). More details are in Figure 1.

Q2. It's essential to explain how the proposal's components address dynamic OOD, focusing on parameter invariance between training and testing.

Thanks for your comment. We have included our motivation in Sec. 3.1.

Q3. Without revision, I can't connect the problem to the proposal from the rebuttal and submission exchanges.

Thanks for your comment. We have revised the manuscript and included all the related content here for your convenience.

Thank you again for your feedback! Please let us know if you have further questions.

We are truly grateful for the time you have taken to review our paper, your insightful comments and support. Your positive feedback is incredibly encouraging for us! In the following response, we would like to address your major concern and provide additional clarification.

Q1. Some results may be incorrectly labeled as suboptimal in table, and there are errors in the use of some symbols. There might be misuses of symbols in the paper, such as, Change "xi" to "si" in line 65, Change "xq" to "xq" in line 96, "Zero-shot Experiments" and "Generalization Experiments" in line 215 should be treated equally and placed on separate lines. Is there an issue with the second-best data in Table 2? For example, in the column w/OOD of SPHERICAL-SWE, the second-best should be DGPDE 0.0028. The corresponding improvement results also need to be modified.

A1. Thank you for pointing this out. We will correct all the typos carefully in the revised version.

Q2. The explanation of the experimental results is not detailed enough, making some experiments difficult to understand. What does the clustering in Figure 5 represent? Could you provide a detailed explanation?

A2. Thank you for your comment. We demonstrate the t-SNE visualization of the predicted trajectories for compared methods. From the results, we can observe that our method outperforms baselines, which validate the superiority of our method. Here, clustering refers to the technique of reducing the dimensions in t-SNE. We will include more details in our revised version.

Q3. The proposed method is aimed at OOD (Out-Of-Distribution), but the experiments lack comparison and discussion with methods specifically targeting OOD, such as [1] and [2]. Reference: [1] Kirchmeyer, Matthieu, et al. "Generalizing to new physical systems via context-informed dynamics model." International Conference on Machine Learning. PMLR, 2022. [2] Yin, Yuan, et al. "LEADS: Learning dynamical systems that generalize across environments." Advances in Neural Information Processing Systems 34 (2021): 7561-7573.

A3. Thank you for your comment. We have added two baselines LEADS [1] and CODA [2] for performance comparison. The results below show that our method is superior to these baselines, which validates the effectiveness of our method when addressing OOD problems. We will include it in our revised version.

Q4. The paper utilizes mutual information to decouple different prompt embeddings and observation embeddings, reducing the sensitivity of observation embeddings to different distributions. However, I don't quite understand the purpose of decoupling. Observational embeddings are related to the environment, and prompt embeddings are related to the environment as well; they are inherently correlated. Please provide further explanation.

A4. Thank you for your comment. Our decoupling aims to build invariance of our observation embeddings to environments for better generalization. Here, both observation embeddings and prompt embeddings are combined for the final prediction, and prompt embeddings are utilized to provide the environment information. In other words, observational embeddings are not correlated with the environment for better generalization. We will include it in our revised version.

Q5. This method does not apply to real-world scenarios, such as rigid dynamics modeling and traffic flow forecasting.

A5. Thank you for your comments. We have added four baselines, i.e., EGNN [1], SGNN [2], SimVP [3], PastNet [4] for performance comparison on the RigidBall and TaxiBJ datasets. The table shows our method outperforms the baseline models in rigid dynamics and traffic flow forecasting. Please see the PDF for visual results. We will include it in our revised version.、

Rigid dynamics (MSE)

Traffic flow

Reference

[1] Satorras, Vıctor Garcia, Emiel Hoogeboom, and Max Welling. "E (n) equivariant graph neural networks." ICML2021

[2] Han, Jiaqi, et al. "Learning physical dynamics with subequivariant graph neural networks." NeurIPS2022

[3] Gao, Zhangyang, et al. "Simvp: Simpler yet better video prediction." CVPR. 2022.

[4] Wu, Hao, et al. "Pastnet: Introducing physical inductive biases for spatio-temporal video prediction." MM2024

We will also add your suggestion about future works to our revised version. Thanks again for appreciating our work and for your constructive suggestions. Please let us know if you have further questions.

Dear Reviewer,

We summarized our rebuttal content at your convenience as follows:

• We have included more competing baselines to demonstrate the superiority of our approach.
• We have included more real-world settings such as rigid dynamics modeling and traffic flow forecasting for performance comparison.
• We have explained the purpose of decoupling and visualization.
• We will proofread our paper to clear every typo in our final version.

We will also add your suggestion about future works to our revised version. Thanks again for appreciating our work and for your constructive suggestions. Please let us know if you have further questions.

Best,

the Authors

We are truly grateful for the time you have taken to review our paper, your insightful comments and support. Your positive feedback is incredibly encouraging for us! In the following response, we would like to address your major concern and provide additional clarification.

Q1. As the method is based on attention, the proposed approach may have limited scalability and take long computation time. Is there a comparison on these with the previous works?

A1. Thank you for your comment. We have added a comparison of computational costs below. From the results, we can observe that our method has a competitive computation cost with huge performance increasement. Without OOD, it improves performance by an average of 4.44%, and with OOD, it improves performance by an average of 25.90%. We will include it in our revised version.

We will also add your suggestion about future works into our revised version. Thanks again for appreciating our work and for your constructive suggestions. Please let us know if you have further questions.

Thank you for your feedback and support! We will add the rebuttal contents to the main paper in the final version following your valuable suggestions.