PURE: Prompt Evolution with Graph ODE for Out-of-distribution Fluid Dynamics Modeling

Authors: hao  wu
changhu  wang
fan  xu
jinbao  xue
chong  chen
xian-sheng  hua
xiao  luo
rnn  rnn
rnn  mhsa
mhsa  mhsa
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 15:08:39

Abstract: This work studies the problem of out-of-distribution fluid dynamics modeling. Previous works usually design effective neural operators to learn from mesh-based data structures. However, in real-world applications, they would suffer from distribution shifts from the variance of system parameters and 
temporal evolution of the dynamical system. In this paper, we propose a novel approach named \underline{P}rompt Evol\underline{u}tion with G\underline{r}aph OD\underline{E} (\method{}) for out-of-distribution fluid dynamics modeling. The core of our \method{} is to learn time-evolving prompts using a graph ODE to adapt spatio-temporal forecasting models to different scenarios. In particular, our \method{} first learns from historical observations and system parameters in the frequency domain to explore multi-view context information, which could effectively initialize prompt embeddings. More importantly, we incorporate the interpolation of observation sequences into a graph ODE, which can capture the temporal evolution of prompt embeddings for model adaptation. These time-evolving prompt embeddings are then incorporated into basic forecasting models to overcome temporal distribution shifts. We also minimize the mutual information between prompt embeddings and observation embeddings to enhance the robustness of our model to different distributions. Extensive experiments on various benchmark datasets validate the superiority of the proposed \method{} in comparison to various baselines.

Main paper: 
 introduction: 
 Fluid dynamics [44, 89] is a critical area in the field of mechanics and computational fluid dynamics has emerged as a powerful tool to understanding fluid flow [32, 58, 67, 45] . Recently, various machine learning approaches have been widely adopted to solve the problem in a data-driven manner [59, 46, 66, 65, 18, 7, 81] , which can achieve high efficiency in comparison to previous traditional numerical solvers. Moreover, they enjoy strong applicability when the underlying rules are not explicit, such as real-world weather forecasting [4] and disease transmission [68] .In literature, existing data-driven fluid dynamics modeling approaches can be roughly divided into grid-based approaches [12, 17] and geometry-based approaches [59, 66, 65, 20] . Grid-based approaches construct regular meshes and then utilize neural operators to explore spatio-temporal relationships. In contrast, geometry-based approaches focus on irregular point clouds and then utilize graph neural networks (GNNs) [30, 70] to learn from the interaction between mesh points.Despite their great success, existing approaches [88, 25] generally assume that training and test data share the same data distribution [59, 66, 65, 47] , which could be not the case in real-world applications. In particular, there are two typical types of distribution shifts in dynamical systems, i.e., parameter-based shifts, and temporal distribution shifts. Firstly, different dynamical systems could involve different parameters in underlying rules, such as coefficients in PDEs and pressures in fluid systems [3, 65] . Secondly, during long-term auto-regressive forecasting, the input data distribution could vary hugely during the temporal evolution [91] . As in previous works [22, 33] , machine learning approaches usually suffer from huge performance degradation when it comes to distribution shifts. Therefore, in this paper, we focus on the problem of out-of-distribution fluid dynamics modeling to enhance the performance under potential distribution shifts.In this paper, we propose a new approach named Prompt Evolution with Graph ODE (PURE) for out-of-distribution fluid dynamics modeling. The high-level idea of our proposed PURE is to adapt well-trained forecasting approaches to different out-of-distribution scenarios by learning timeevolving prompts [51, 84, 9] . To begin, we extract multi-view context signals from both historical observations and system parameters in the frequency domain using the attention mechanism, which can effectively initialize prompt embeddings under parameter-based shifts. More importantly, to capture temporal distribution shifts, we combine the interpolation of observation sequences into a graph ODE framework, which can utilize the interaction between prompt embeddings and observation embeddings for high-quality time-evolving prompt embeddings. Then, we concatenate our prompt embeddings and observation embeddings for model adaptation and enhance the robustness of our PURE to distribution variance by minimizing their mutual information using adversarial learning. Extensive experiments on a range of fluid dynamics datasets validate the superiority of the proposed PURE in comparison to various state-of-the-art approaches.In summary, the contribution of our paper can be summarized as follows: (1) Problem Connection. We are the first to connect prompt learning with dynamical system modeling to solve the issue of out-of-distribution shifts. (2) Novel Methodology. Our PURE first learns from historical observations and system parameters to initialize prompt embeddings and then adopts a graph ODE with the interpolation of observation sequences to capture their continuous evolution for model adaptation under out-of-distribution shifts. (3) Superior Performance. Comprehensive experiments validate the effectiveness of our PURE in different challenging settings. 
problem setup: 
 Given a fluid dynamical system, we have N sensors within the domain Ω, with their locations denoted as x 1 , • • • , x N , where x i ∈ R d l . The observations at time step t are represented as s t 1 , • • • , s t N , where s t i ∈ R do and d o indicates the number of observation channels. Dynamical systems are governed by underlying system rules, such as PDEs with coefficient ξ. Variations in system parameters may lead to different environments, potentially resulting in distribution shifts [54, 80, 6, 27] . In our study, we are provided with historical observation sequences {s 1:T0 i } N i=1 and physical parameters ξ (e.g., coefficients in the PDEs). Our goal is to predict the future observations of each sensor s T0+1:T0+T i . In dynamical systems, the out-of-distribution problem examines model performance when predicting under unseen parameter distributions or environments. Letu t = [s t 1 , • • • , s t N ], these systems evolve according to du dt = F (u, ξ), where u represents the observations and ξ denotes the system parameters. When ξ ∼ P (ξ), the state trajectory u 1:T0 follows the distribution P (u 1:T0 |ξ). Assume we learn a learned mapping function f from u 1:T0 to u T0+1:T0+T , i.e., u T0+1:T0+T = f (u 1:T0 ) and there could be different distributions across training and test datasets, i.e., P train (ξ) ̸ = P test (ξ), which results in P train u 1:T0 ̸ = P test u 1:T0 . Moreover, when conducting rollout prediction, we are required to feed the output back to the model, i.e., u Tstart:Tstart+T -1 = f (u Tstart-T0:Tstart-1 ), with P (u 1:T0 |ξ) ̸ = P (u Tstart-T0:Tstart-1 |ξ, T start ).3 The Proposed PURE 
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 (•) can suffer from a serious distribution shift result from different ξ and T start , i.e., P (u input |ξ, T start ). To reduce the impact of distribution shift, we aim to learn invariant observation embeddings µ t to different environments, i.e., ξ and T start for better generalization and utilize prompt embeddings z t to indicate the current environment for final prediction. In formulation, we have:EQUATIONThe 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:EQUATION)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.2). ( 2) Time-evolving Prompt Learning, which incorporates the interpolation of observation sequences into a graph ODE to learn the evolution of prompt embeddings (Eqn. 2). (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 .The main idea of our PURE is to utilize prompt learning to solve the issue of out-of-distribution shifts [54, 80, 6, 27] in fluid dynamical systems. Prompt learning [51, 84, 9] is an effective manner to adapt language models to various downstream tasks. In our scenarios, we aim to learn from both historical spatio-temporal information and system parameters to initialize our prompt embeddings, which can effectively solve the parameter-based distribution shifts. Here, we first follow the attention mechanism [69, 86, 10, 50] to reconstruct the field value and then adopt the Fourier neural operator [45] to integrate multi-view context information.In particular, given each location, we map each location x i and each initial observation s 0 i into a position embedding p i and an observation embedding q i using two feed-forwarding networks (FFNs) ϕ P E (•) and ϕ OE (•), and then aggregate p i and q i using the Hadamard product followed by stacking L self-attention blocks for representation learning. In formulation,EQUATIONwhere ⊙ denotes the Hadamard product, E 0 is constructed by stacking {e i } N i=1 , and ϕ SA,(l) is the self-attention block at the layer l. Afterward, we adopt the attention mechanism [69] to retrieve the representations for each query position x q as:EQUATIONwhere W * is a learnable weight matrix for feature transformation and d is the hidden dimension. By retrieving the representations at each regular grid, we generate the 3D representation tensor U . Each tensor would be concatenated with the parameter embedding u p = ϕ P A (ξ) for multi-view information integration, which results in the final tensor Ũ . Then, we utilize the frequency domain for representation enhancement to generate a prompt tensor H. Here, we first transfer the tensor into the frequency domain using a Fast Fourier Transformer (FFT) operator [45] and then adopt an FFN for feature transformation. Lastly, an inverse Fast Fourier Transformer (iFFT) operator is adopted to convert the features back to the spatial domain. Formally,EQUATIONwhere FFT(•) and iFFT(•) denote the FFT and iFFT operators, respectively. Since the input of spatio-temporal models would be irregular, we flatten the tensor H, and retrieve prompts for each sensor from the prompt tensor using:EQUATIONwhere f latten(•) is a flattening operator to transform 3D tensors to 2D matrices. Through the frame reconstruction, we can extract important spatio-temporal signals from the frequency domain, which is effective in initializing the prompt embedding for each sensor.To capture temporal distribution shifts within one system, static prompt embeddings [51] from context exploration are far from satisfactory. Our solution is to obtain continuous time-evolving prompts at any timestamp. To achieve this, we view the output of Eqn. 6 as the initial prompt embeddings and then incorporate the attention mechanism into a continuous graph ODE, which combines the interpolations of observations with the graph structure to learn the evolution of prompt embeddings.In particular, given the initial prompt embeddings, we introduce two functions ψ a (•) and ψ r (•) for relation mining and feature aggregation. ψ r (•) calculates the interaction between the centroid node and each of its neighboring nodes and ψ a (•) aggregates all the neighborhood interactions to determine the evolution. Therefore, a graph ODE can be formulated by the following formulation:EQUATIONwhere S t (i) collects the sensors from the neighbours of i at timestamp t. However, Eqn. 7 neglects observations themselves during evolution, which are directly related to temporal distribution shifts in dynamical systems. Thus, it could generate suboptimal prompt embeddings. To tackle the issue, we conduct the interpolations of observation sequence s 1:T0 i , which results in s t i at any timestamp. Then, we incorporate them into our graph ODE using the attention mechanism by rewriting Eqn. 7 into:EQUATIONwhere W Q and W K are two matrices to generate the query and key, respectively. Here, we utilize the prompt embeddings and interpolated observations to serve as the query and the key. In this way, we effectively model their interaction to adjust the derivative in the graph ODE, which can help generate proper prompt embeddings for our model adaptation under temporal distribution shifts.Finally, we incorporate our prompt embedding into our basic spatio-temporal forecasting model, and then introduce the optimization objective for the end-to-end training.Basic Forecasting Model. Our time-evolving prompt embeddings can be easily incorporated into any spatio-temporal forecasting model. To make the best of our efficacy, we utilize a simple yet powerful basic model as our default model and also explore the performance of our PURE on more existing forecasting models. The input of our model is the observations of sensors from between the interval [1, T 0 ] and output the predictions in [T 0 + 1, T 0 + T ], i.e., S 1:T0 → S T0,T0+T where S * is stacked by s * i . In particular, our basic model first generates the embeddings of different observations, and then reconstructs the irregular observations into frames on grids using the reconstruction modules in Sec. 3.2. More importantly, we introduce two parallel modules, i.e., a Fourier neural operator [45] and a ViT-based convolution [19] and extract complementary feature maps [81] , which would be fused to generate the predicted frames in the future. More details of our basis forecasting model can be found in Appendix B. Additionally, we also use other basic models.Adaptation and Optimization. Note that we generate observation embeddings in our basic module, i.e., µ t i = ϕ enc (s t i ). To adapt our model under distribution shifts, we concatenate the observation embeddings and prompt embeddings into updated embeddings μt i as follows:EQUATIONwhich will be fed into the subsequent modules in the basic model. To optimize the whole framework, we first minimize the mean squared error (MSE) between the predicted observation and the ground truth as follows:EQUATIONwhere Ŝt denotes our predicted observation for every node and S t denotes the ground truth observations. Moreover, to enhance the invariance of our model to different scenarios, we turn to invariant learning [72, 42, 77] to decouple various prompt embeddings and observation embeddings, which promotes the observation embeddings to be less sensible to different distributions. To achieve this, we minimize the mutual information between observation embeddings and observation embeddings, i.e., I(µ t i ; z t i ). In our work, we adopt a Jensen-Shannon mutual information estimator [40, 53] T γ (•, •) where γ denotes the parameters to estimate their mutual information. Then, we collect all the corresponding pairs of (µ t i , z t i ) using P and all the possible pairs of (µ t i , z t j ) using N . The adversarial learning objective can be written as:EQUATIONin which sp(x) = log(1 + e x ) represents the softplus function. In summary, the overall objective can be written as:EQUATIONwhere λ is a coefficient to balance two loss objectives. The algorithm is summarized in Appendix D.In this part, we provide a theoretical analysis to demonstrate how PURE works. Our focus is primarily on theoretically showing the necessity of incorporating the observations themselves during evolution. For simplicity of analysis, we assume that Eqn. 7 can be rewritten as:EQUATIONwhere #(•) calucates the size of the set. For the sake of simplicity in the proof, we assume that z t i is one-dimensional and consider only the ODE above for i (not the entire system of ODEs). Then, Eqn. 13 can be rewritten as:EQUATIONwhere b(t) is a function. To characterize temporal distribution shifts, we assume that a portion of the corresponding true z t j has a constant shift. With the potential environmental change, Eqn. 13 can be rewritten as:EQUATIONwhere |b(t) -b ′ (t)| ≥ c 0 , suggesting the constant shift c 0 .Theorem 3.1. Given the following ODEs in R,EQUATIONwhereEQUATION)The proof of Theorem 3.1 can be found in Appendix A. Theorem 3.1 suggests that even with the simplest one-dimensional linear ODE, significant differences in the solutions will arise if temporal distribution shifts are neglected. Next, we will focus on how Eqn. 8 addresses this issue. In this case, we assume thatEQUATIONThen, Eqn. 7 can be written as:EQUATIONEQUATIONSimilarly, Eqn. 8 can be written as:EQUATIONwhereEQUATIONFor simplicity of notation, we omit the superscript i.Write F (z, t) = ψ α M 1 z + b(t) and G(z, t) = ψ α M 1 z t + b ′ (t). Then, we have the following theorem with the proof in Appendix A.Theorem 3.2. Assume that the attention mechanism satisfies that |b ′ (t) -b(t)| ≤ ϵ, for all t > 0, and the function ϕ α is L-Lipschitz. Given the following ODEs in R,EQUATIONthere exists two constants c 2 and c 3 such thatEQUATION)Thoerem 3.2 shows that as long as the attention mechanism is sufficiently good, we can approximate the true ODE with arbitrary precision using Eqn. 8, even in the presence of environmental change. Benchmarks. We study Benchmarks from three domains, as shown in Table 5 . ▷ Computational Fluid Dynamics. We use Prometheus [80] and follow the original setup for environment segmentation. ▷ Real-world Data. We employ the ERA5 [23] , using different combinations of variables as the environment. In detail, we use ERA5 data with variables such as surface pressure (Sp), sea surface temperature (SST), sea surface height (SSH), and two-meter temperature (T2m) to predict temperature.▷ Partial Differential Equations. The 2D Navier-Stokes equations [45] describe fluid motion, with the primary variable being the viscosity coefficient ν, which quantifies internal friction in the fluid, simulating vorticity values under ten different viscosity coefficients. The spherical shallow water equations [14] simulate large-scale atmospheric and oceanic fluid motion on Earth's surface, also with viscosity coefficient ν as the main variable, involving tangential vorticity (w) and fluid thickness (h) on a spherical surface. The 3D reaction-diffusion equations describe the diffusion and reaction of chemicals in space [62] , with the primary variable being the diffusion coefficient D, representing the rate of chemical diffusion in space, including u, v velocity components. More details see Appendix E.Baselines. We select representative models from three domains as baselines. ▷ Visual Backbone Networks. We include ResNet [21] , U-Net [64] , Vision Transformer(ViT) [10] , and Swin Transformer(SWINT) [49] . ▷ Neural Operator Architectures. We cover FNO [45] , UNO [1] , CNO [63] , and NMO [82] . ▷ Graph-ODE Architectures. We feature CG-ODE [26] , and DGODE [80] .Tasks. We evaluate model performance for various prediction tasks through the following scenarios and use MSE as metrics. The specific tasks are as follows:▷ Generalization Experiments:  
generalization experiment results: 
 In this section, we focus on the issue of generalization. Based on our experimental findings, we make the following observations. Out-of-Distribution Generalization. The results as shown in Table 1 . On the Prometheus dataset, PURE outperforms all benchmark models with an MSE of 0.0323 in-distribution and 0.0328 OOD. It improves over the second-best model, DGODE (MSE 0.0344 in-distribution, 0.0359 OOD), by 6.10% and 8.63%, respectively. On the Navier-Stokes dataset, PURE achieves the best performance with an MSE of 0.0752 in-distribution and 0.0763 OOD, improving by 6.58% and 26.07% over the second-best model. On the Spherical-SWE dataset, PURE has an MSE of 0.0022 both in-distribution and OOD, which is 41.46% better than the secondbest model. Additionally, in Table 2 , the performance of various benchmark models significantly improves when using PURE, in summary, the PURE framework performs excellently in handling OOD fluid dynamics modeling.Spatial & Temporal Generalization. Visualization and Analysis. Figure 3 compares the performance of different methods in fluid dynamics modeling, including the Prometheus dataset, Navier-Stokes equations, and the 3D Reaction-Diffusion Equation. In the Prometheus dataset, using PURE significantly reduces DGODE's prediction error, especially in complex dynamic regions. For the Navier-Stokes and Spherical Shallow Water equations, FNO and NMO models combined with PURE excel in capturing complex flow features. In the 3D Reaction-Diffusion Equation, DGODE with PURE significantly reduces prediction errors. Overall, PURE greatly enhances the prediction accuracy of models in fluid dynamics, allowing for better capture of complex dynamic evolution.As shown in Figure 4 , the PURE framework performs excellently in zero-shot super-resolution and environmental generalization experiments. In the Prometheus benchmark, the FNO model using In this section, we evaluate the effectiveness of the PURE method and the importance of its components through qualitative analysis and ablation studies.Qualitative Analysis. The Figure 5  
conclusion: 
 In this paper, we study a practical problem of out-of-distribution fluid dynamics modeling and propose a novel approach named PURE for this problem. The high-level idea of our PURE is to learn time-evolving prompts using graph ODEs, which can effectively adapt spatio-temporal forecasting models to different scenarios. Our PURE first initializes prompt embeddings by exploring multi-view context information from spatio-temporal data and system parameters. Then, PURE incorporates the interpolation of observation sequences into the graph ODE, which helps capture the temporal evolution of prompt embeddings to mitigate temporal distribution shifts. In future works, we will extend our PURE to more real-world scenarios such as rigid dynamics modeling and traffic flow forecasting.A Proofs of Theorem 3.1 and Theorem 3.2Proof of Theorem 3.1. We first introduce a lemma as follows.Lemma A.1. Suppose that ẋ = M 1 x + b(t), x(0) = x 0 . Then, the solution can be written asEQUATIONThe proof of Lemma A.1 is from Method of Variation of Parameters [37] . With Lemma A.1, we can prove Theorem 3.1.By Lemma A.1, we haveEQUATIONandEQUATIONIt follows thatEQUATIONBy the Mean Value Theorem for Integrals, we haveEQUATIONThus, we obtain that |x(t) -y(t)| ≥ c 0 (e M1t -1)/M 1 ,which we complete the proof.Proof of Theorem 3.2. To prove the theorem, we need the following lemma, which can be found in [5] .Lemma A.2. Given the following ODEs in R n ,EQUATIONassume that F is globally Lipschitz continuous and "close to G. In other words, there exist L ≥ 0 and ϵ ≥ 0 such thatEQUATIONand assume that ∥Φ(t, s)∥ i ≤ ce η(t-s) , for all 0 ≤ s ≤ t < T, (33) in which ∥ • ∥ i denotes the induced matrix norm associated with vector norm ∥ • ∥, Φ(t, s) denotes the transition matrix for A(t), and c ≥ 1. Then for all t ∈ [0, T ), if η + cL ̸ = 0EQUATIONIn this case, we set c = 1, η = 0. Then, Theorem 3.2 is clear from the Lemma A.2. 
experimental setting/details: 
 Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results? Answer: [Yes] Justification: we have introduced experiment settings and details, and even add more details in the appendix. Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. 
experiment statistical significance: 
 Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments? Answer:[No] Justification: We do not report specific error bars for the following reasons: (a) In fluid dynamics modeling, using a fixed random seed shows consistent performance, and the performance difference with different seeds is small [78, 45] . (b) Related work [45, 80, 55] in this field does not report error bars. (c) To ensure fairness, we fix all random seeds and conduct experiments on the same machine. Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors). • It should be clear whether the error bar is the standard deviation or the standard error of the mean. • It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified. • For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates). • If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text. 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments? Answer: [Yes] Justification: we have discussed the information about computer resources. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: Our base forecasting model combines two parallel modules [81] : the Fourier Neural Operator (FNO) [45] and a Vision Transformer (ViT)-based convolution [10] . The FNO processes the input observation embeddings in the frequency domain with Fast Fourier Transform (FFT) and inverse FFT (iFFT), capturing frequency features. The ViT module uses a multi-head attention mechanism to process the spatial features of the input data. The PURE framework generates time-evolving prompt embeddings using a Graph ODE, capturing dynamic changes in spatio-temporal features. These embeddings, along with the features from the FNO and ViT modules, are integrated using skip connections and a Multi-Layer Perceptron (MLP) to produce the final predictions. The model optimizes by reducing the mean squared error (MSE) between predictions and the ground truth, and by enhancing robustness through reducing mutual information between prompt and observation embeddings. This approach ensures high accuracy in out-of-distribution fluid dynamics forecasting.Dynamical System Modeling. The field combining machine learning with dynamical systems aims to use machine learning methods to model, predict, and control the behavior of dynamical systems [45, 80, 59, 55, 2, 54, 81] . Key techniques include using neural networks to extract patterns from spatiotemporal data, such as Convolutional Neural Networks (CNN) [63, 61] , Graph Neural Networks (GNN) [59, 43, 38] , and Transformer models [79, 4, 36] . Additionally, Physics-Informed Neural Networks (PINN) [34, 60] embed physical laws into neural networks to enhance the model's physical consistency. These methods apply to both short-term and long-term predictions of dynamical systems and optimize control strategies in areas like robotics and autonomous driving [11, 57] . To address the challenge of out-of-distribution data, researchers develop new datasets and benchmarks to evaluate model performance under different data distributions [80] . This field finds wide applications in aerospace, biomedical, and meteorological domains [93, 39] . In this work, we propose a framework named PURE, which uses prompt learning and graph neural ODE to address complex distribution shifts in fluid dynamics due to parameter and temporal changes.Out-of-distribution Generalization Out-of-distribution (OOD) [71, 73, 83, 85, 27] generalization means a model performs well on new, unseen data. The core goal in this field is to improve model performance when training and test data come from different distributions. Models that excel in OOD scenarios should be robust and adaptable. Researchers have proposed several methods, such as data augmentation [74, 8] , invariant feature learning [42, 77, 41, 72] , adversarial training [75, 8] , and domain adaptation [35, 15] . These methods are widely used in areas like autonomous vehicles, medical diagnosis, financial forecasting, and dynamical systems modeling [80, 13, 53] . In this work, we propose PURE, which uses prompt learning and graph neural ODE to adapt spatio-temporal forecasting models to address distribution shifts in fluid dynamics.Prompt Learning. Prompt learning [48, 16, 24] has recently gained significant attention as a strategy for adapting pre-trained models to various downstream tasks by leveraging the power of promptbased fine-tuning [9, 29, 94, 76, 52] . In the domain of large language models, prompt learning aims to incorporate optimal tokens into the input sequence, which can effectively improve performance without extensive retraining [28, 92, 90] . In the context of fluid dynamics modeling, prompt learning means a supplementary hint to indicate the context, which is incorporated into the input (observation embedding) for better generalization. Although it shares a similar meaning as prompt tuning in language models, our prompt refers to the current environment, which determines the future evolution with better generalization.The whole learning algorithm of PURE is summarized in Algorithm 1.We evaluate our proposed PURE on five physical benchmarks. Map location x i and initial observation s 0 i to embeddings p i and q i using FFNs ϕ P E (•) and ϕ OE (•)Aggregate embeddings: e i = p i ⊙ q i 5: end for 6: Stack initial embeddings E 0 = {e i } N i=1 and apply self-attention blocks ϕ SA,(l) 7: Retrieve representations u q for each query position x q using attention mechanism 8: Generate 3D representation tensor U and integrate with parameter embedding u p = ϕ P A (ξ) to obtain Ũ 9: Enhance representation in frequency domain: H = iFFT(FFN(FFT( Ũ ))) Navier-Stokes equations [45] depict the motion of a viscous, incompressible fluid. The equations are as follows in vorticity form on the unit torus:EQUATIONWe solve these equations using the stream function formulation and a pseudo-spectral approach.In particular, we first solve the Poisson equation in order to identify the velocity field. Afterward, we differentiate vorticity, compute the nonlinear terms, and apply de-aliasing. We use the Crank-Nicolson scheme for time-stepping, recording the solution at time intervals of t = 1 on a 256×256 grid followed by downsampling. For the Bayesian inverse problem, the timestep during data generation is 1e-4, and in MCMC, it is 2e-2. We simulate the Navier-Stokes equations with varying viscosity coefficients, adjusting ν to study its impact on fluid flow and vorticity distribution.Spherical Shallow Water Equations (Spherical-SWE) [14] describe the large-scale atmospheric and oceanic fluid motion on Earth's surface. The equations are:EQUATIONHere, h represents fluid thickness, u is the fluid velocity vector, f represents the Coriolis parameter, k represents the unit vertical vector, g represents the gravitational acceleration, ν represents the viscosity coefficient, and ∆ represents the Laplacian operator. The first equation (continuity equation) represents mass conservation, describing changes in fluid thickness. The second equation (momentum equation) represents momentum conservation, including advection, Coriolis force, pressure gradient force, and viscous diffusion. We simulate the Spherical Shallow Water Equations with different viscosity coefficients ν, adjusting ν to study its effects on fluid motion.3D Reaction-Diffusion Equations [62] describe the diffusion and reaction of chemical substances in space. The general form of these equations is:EQUATIONHere, u and v represent the concentrations of the chemical substances, D u and D v represent the diffusion coefficients for u and v, respectively, ∆ is the Laplacian operator, and R u (u, v) and R v (u, v) denote the reaction terms that represent the reaction rates between u and v. Diffusion terms, i.e., ∆u and ∆v denote the diffusion process of the chemicals in space. The diffusion coefficients D u and D v determine the rate of diffusion. Reaction terms, i.e., R u (u, v) and R v (u, v) describe the reaction rates of the chemical substances. These terms depend on the concentrations of u and v, and can include linear reactions, nonlinear reactions, and complex dynamic processes. We simulate the 3D reaction-diffusion equations with different diffusion coefficients D u and D v to study the effects of diffusion rates on the distribution and reaction rates of the chemical substances.ERA5 [23] is a global atmospheric reanalysis dataset produced by ECMWF, providing weather data from 1979 to the present with high spatial (31 km) and temporal (hourly) resolution. It includes variables like surface pressure, sea surface temperature, sea surface height, and two-meter temperature. ERA5 data supports applications in weather forecasting, climate research, environmental monitoring, energy management, and agriculture. Accessible via the Copernicus Climate Data Store, ERA5 is crucial for analyzing and predicting meteorological and climate phenomena. ν = {1e -1 , 1e -2 , ..., 1e -9 , 1e -10 } ν = {1e -11 , 1e -12 } SPHERICAL SHALLOW WATER EQUATION ν = {1e -1 , 1e -2 , ..., 1e -9 , 1e -10 } νt = {1e -11 , 1e -12 } 3D REACTION-DIFFUSION EQUATIONS D = {2.1 × 10 -5 , 1.6 × 10 -5 , 6.1 × 10 -5 } D = {2.03 × 10 -9 , 1.96 × 10 -9 } ERA5 V = {Sp, SST, SSH, T 2m} V = {SSR, SSS}The compared approaches involved in this study is as follows:• U-Net [64] is a convolutional neural network initially used for biomedical image segmentation. It has a symmetric U-shaped structure and uses skip connections to link the encoder and decoder, enabling efficient feature fusion.• ResNet [21] introduces residual blocks to solve the degradation problem in deep networks. It allows the network to be deeper and easier to train by using skip connections to directly pass information.• ViT [10] applies the Transformer model to image recognition. It divides the image sample into patches and uses self-attention mechanisms to process these patches, balancing computational efficiency and performance.• SwinT [49] introduces a sliding window mechanism for effective local and global feature extraction. It is suitable for various computer vision tasks.• FNO [45] uses Fourier transforms for global feature extraction, suitable for processing continuous field data and efficiently solving PDEs.• UNO [1] combines the U-Net architecture with optimization methods to enhance feature extraction and fusion capabilities, improving model performance.• CNO [63] combines convolution operations with operator learning, focusing on high-dimensional continuous data and modeling complex dynamic systems.• NMO [82] enhances the modeling capability for multi-scale dynamic systems by combining neural networks with manifold learning algorithms.• CGODE [26] is a neural ODE model that aims to capture the dynamics of both nodes and edges jointly.• DGODE [80] addresses the challenge of out-of-distribution (OOD) generalization in fluid dynamics modeling by learning disentangled representations using a temporal GNN and a frequency network. It minimizes mutual information between node and environment representations to mitigate distribution shifts and employs a coupled graph ODE framework for robust modeling.Mean Squared Error (MSE). Mean Squared Error measures the gap between the predicted and ground truth. The formula is:EQUATIONin which y i denotes the actual value, ŷi denotes the predicted value, and n denotes the number of data points.Relative L2 Error. Relative L2 Error evaluates the relative accuracy of the model's predictions. The formula is:EQUATIONwhere ∥ • ∥ 2 denotes the L2 norm, y is the vector of actual values, and ŷ is predicted values.▷ Sparse Reconstruction Experiments. The experimental setup includes sparse reconstruction experiments on the Prometheus and ERA5 datasets. For the Prometheus dataset, the sparsity rates are set to 25%, 50%, and 75%. For the ERA5 dataset, the sparsity rates are set to 5% and 25%. Each set of experimental results includes the original data, sparse input data, results from our method, and results from MMGNet [56] . The experiments evaluate the performance of each method by comparing the reconstruction results at different sparsity rates.Figure 6 shows two sets of sparse reconstruction experiment results from the Prometheus and ERA5 datasets. The sparsity rates for Prometheus are 25%, 50%, and 75%, while for ERA5, they are 5% and 25%. Each set displays the Ground Truth, Sparse Input, our reconstruction method (Ours), and MMGNet's results in order. As the sparsity rate increases, the quality of the reconstruction decreases. At low sparsity rates, both our method and MMGNet effectively restore image details. At high sparsity rates, our method performs better in preserving details and recovering overall structure.▷ More Ablation Experiments. ▷ Performance with respect to Different Difficulty Levels. Here, we demonstrate the performance of our PURE with varying difficulty levels. In particular, we measure the difficulty levels based on the distance between P train (ξ) and P test (ξ) and generate three levels on the Prometheus dataset.The compared results are shown in Table 7 . From the results, we can observe that all the model performs worse in hard scenarios and our method has consistently outperformed 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. ▷ Robustness to Noisy Data. We have also experimented with noisy data to evaluate the robustness of our method. The results are shown in Table 8 . From the results, the performance of both ResNet and NMO models degrades significantly when noise is introduced. However, the integration of our PURE with these models substantially mitigates the impact of noise, leading to much lower MSE values compared to their baselines.▷ Expanded Evaluation of OOD Generalization in Dynamical Systems. To further validate the effectiveness of our PURE, we include three additional models specifically designed for outof-distribution generalization in dynamical systems: LEADS [87] , CODA [31] , and NUWA [74] . Table 9 shows the performance of each method across different datasets in both in-distribution (ID) and out-of-distribution (OOD) scenarios. The results indicate that our proposed method outperforms these baselines on all datasets, especially in OOD scenarios. Although the PURE method shows superiority on multiple benchmark datasets, it has some limitations. First, we assume that training and test data are independent and identically distributed (IID). This may be not true in some extreme physical scenarios due to environmental changes causing significant data distribution shifts. Second, while the PURE method performs well in addressing distribution shifts, it may still face challenges when dealing with high-dimensional and complex fluid dynamics systems. Additionally, the PURE method has high computational complexity, requiring more computational resources and time in practical applications. Future research can focus on optimizing the algorithm to improve computational efficiency and extending it to more real-world scenarios such as rigid dynamics modeling and traffic flow forecasting.The PURE method significantly impacts out-of-distribution fluid dynamics modeling. It adapts to different scenarios, improving the model's performance in handling distribution changes. This is helpful for climate prediction, epidemic spread, aerospace, and biomedical fields. In future works, we will extend our PURE to more real-world scenarios such as rigid dynamics modeling and traffic flow forecasting.• The answer NA means that the paper does not include experiments.• If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not. • If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable. • Depending on the contribution, reproducibility can be accomplished in various ways.For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed. • While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. , with an open-source dataset or instructions for how to construct the dataset). (d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility.In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 


Decision: This paper proposes a new approach for out-of-distribution fluid dynamics modeling, PURE. It can learn time-evolving prompt embedding, using a graph ODE, for model adaptation, to overcome temporal distribution. Experiments show its effectiveness.

Firstly, the paper studies an important topic in fluid dynamics modeling. Secondly, the idea is new, and the techniques are sound. Finally, the paper provides extensive empirical results to support its claims. However, the presentation is not clear enough. For example, the studied problem is not well-formulated. And there are quite a few typos.

The authors should further improve the paper carefully!

Reviews: review: paper_summary: paper_summary: The paper proposes a graph ODE-based approach for OOD fluid dynamics modeling. PURE aims to learn time-evolving prompts via graph ODE for adaptation of spatio-temporal forecasting models on OOD scenarios. To address temporal distribution shifts, the interpolation of obersvation sequences are combined into graph ODE framework to learn evolution of prompt embeddings.
strength_weakness: strength_weakness: As I am not an expert in this field, I am unable to find major concerns or weakness of the approach.
- 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?
questions: questions: Please address the questions in the Weaknesses.
limitations: limitations: The limitations are explained in Appendix I.
paper_summary: paper_summary: This paper pioneers the connection of prompt learning with dynamical system modeling to address the challenge of out-of-distribution shifts. The proposed PURE method initializes prompt embeddings by learning from historical observations and system parameters.
strength_weakness: strength_weakness: 1. Some results may be incorrectly labeled as suboptimal in table, and there are errors in the use of some symbols.
2. The explanation of the experimental results is not detailed enough, making some experiments difficult to understand.
3. 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.
questions: questions: Q1. 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.
Q2. 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.
Q3. What does the clustering in Figure 5 represent? Could you provide a detailed explanation?
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.
limitations: limitations: This method does not apply to real-world scenarios, such as rigid dynamics modeling and traffic flow forecasting.
paper_summary: paper_summary: - The paper aims to improve the out-of-distribution (OOD) generalization of fluid dynamics modeling.

- Two types of OOD scenarios are targeted: OOD across different systems and OOD within the same system across different timestamps.

- The paper proposes a framework named PURE, composed of modules including:
    - Multi-view Context Exploration, which explores spatio-temporal data using both the attention mechanism and the frequency domain;
    - Time-evolving Prompt Learning, which incorporates the interpolation of observation sequences;
    - Model Adaptation with Prompt Embeddings, which leverages time-evolving prompts to mitigate temporal distribution shifts.

- Extensive experiments on a range of fluid dynamics datasets support the claim.
strength_weakness: strength_weakness: - 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.

- 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.

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

- 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?
questions: questions: NA
limitations: limitations: NA
paper_summary: paper_summary: The paper presents a new approach called Prompt Evolution with Graph ODE (PURE) for non-distributed fluid dynamics modeling. PURE first learns from historical observations and system parameters in the frequency domain to explore multi-view contextual information, which can efficiently initialize the cue embedding. Interpolations of the observation sequences are then merged into the graph ODE so that the time evolution of the model-adaptive cue embeddings can be captured. These time-evolving cue embeddings are then incorporated into the underlying predictive model to overcome spatio-temporal distributional variations. In addition the paper minimizes the mutual information between the cue embeddings and the observation embeddings to enhance the robustness of the model to different distributions. Finally, extensive experiments conducted on various kinds of benchmark datasets validate the superiority of the proposed PURE compared to various baselines.
strength_weakness: strength_weakness: 1. 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.

2. 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?

3. 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. (b) The experiments in this paper do not state whether noisy data are considered. (c) The authors just give a brief description of the results without analyzing the reasons behind the high performance.
questions: questions: 1.  What are the advantages of PURE over previous efforts to address OOD? 

2. Some of the symbols in the method description section are not defined, e.g., what do P and N in Equation 9 refer to?
limitations: limitations: The authors list limitations in the appendix but do not mention them in the main text. It is recommended that the author make a description in the main text.

Comments: title: Summary of our rebuttal
Summary of our rebuttal
Summary of Rebuttal
Further clarification
Thank you for your feedback and support!
Thanks for your feedback!
Thank you for your responses
comment: 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
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
Dear Reviewers and Area Chairs,

We would like to express our sincere gratitude for your great efforts, insightful comments, support, and the constructive suggestions you have provided once again! Through our discussions and the reviewers' responses, it appears that we have effectively addressed the major concerns raised by reviewers. This outcome has greatly benefited us, and we would like to thank all of you for your valuable support!

The reviewers held many positive comments on our paper. For example, **three reviewers** (Reviewer Sh2w, Reviewer MBSt, and Reviewer MFD5) commented that our work is **novel or new**. Reviewer MFD5 and Reviewer MBSt commented that the paper is **technically sound**. Reviewer MBSt acknowledged that our paper is **easy to follow** and shows **comprehensive experimental** comparisons. Reviewer vRkX believes that our work is **well-motivated**. 

The reviewers also raised insightful and constructive concerns. We made every effort to address all the concerns by providing detailed clarification and requested results. Here is the summary of the major revisions:

- We have **clarified our problem definition and motivation**, and provided the corresponding **revised draft** for reviewers' convenience. 
- We have added **efficiency analysis** to make the paper more complete.
- We have included **more competing baselines** including LEADS, CODA and NUWA to demonstrate the superiority of our approach.
- We have demonstrated **more performance comparisons** in different real-world scenarios to validate the effectiveness of the proposed method in the real world.

We firmly believe that our framework for out-of-distribution fluid dynamics modeling plays a significant role in advancing the community. And we are committed to making our complete code and training details publicly available. All the rebuttal contents will be properly included in the final version, following your valuable suggestions. Moreover, we are eager to engage in further discussions with you to enhance our understanding of the domain and further improve the quality of the paper.

Once again, thank you for your time and effort in reviewing our work. We greatly appreciate your assistance in improving our manuscript!

Best regards,

the Authors
Dear Reviewer,

As the deadline for the author-reviewer discussion phase is approaching, we would like to check if you have any other remaining concerns about our paper. We greatly appreciate your feedback and have worked diligently to address your comments. Thanks to your suggestions, we have revised the draft. **For your convenience, we show all the revised part of our draft as follows** and we believe that it may not be necessary to review the rebuttal text and the submission back-and-forth at this time. 

> **Major concern about OOD challenge formulation and component necessity:**


We have revised the manuscript and included all the related content [Section 2 (Problem Setup) and Section 3.1 (Motivation and Framework Overview)] in our last response for your convenience.

> **Clarification on the term "prompt":** 

We have revised the manuscript and included all the related content [Appendix C] below for your convenience.

**Prompt Learning.** Prompt learning [46, 15, 23] has recently gained significant attention as a technique for adapting pre-trained models to various downstream tasks by leveraging the power of prompt-based fine-tuning [8, 27, 85, 69, 49]. In the domain of large language models, prompt learning aims to incorporate optimal tokens into the input sequence, which can effectively improve performance without extensive retraining [26, 83, 81]. In the context of fluid dynamics modeling, prompt learning refers to a supplementary hint to indicate the context, which is incorporated into the input (observation embedding) for better generalization. Although it shares a similar meaning as prompt tuning in language models, our prompt refers to the current environment, which determines the future evolution with better generalization.

**[46]** Yajing Liu, Yuning Lu, Hao Liu, Yaozu An, Zhuoran Xu, Zhuokun Yao, Baofeng Zhang, Zhiwei Xiong, and Chenguang Gui. Hierarchical prompt learning for multi-task learning. CVPR2023

**[15]** Yuxian Gu, Xu Han, Zhiyuan Liu, and Minlie Huang. PPT: Pre-trained prompt tuning for few-shot learning. arXiv2021

**[23]** Tony Huang, Jack Chu, and Fangyun Wei. Unsupervised prompt learning for vision-language models. arXiv 2022

**[8]** Ning Ding, Shengding Hu, Weilin Zhao, Yulin Chen, Zhiyuan Liu, Hai-Tao Zheng, and Maosong Sun. OpenPrompt: An open-source framework for prompt-learning. arXiv 2021.

**[27]** Muhammad Uzair Khattak, Hanoona Rasheed, Muhammad Maaz, Salman Khan, and Fahad Shahbaz Khan. Maple: Multi-modal prompt learning. CVPR2023

**[85]** Kaiyang Zhouet al. Conditional prompt learning for vision-language models. CVPR2022

**[69]** Zifeng Wang et al.. Learning to prompt for continual learning.CVPR2022

**[49]** Yuning Lu et al. Prompt distribution learning. CVPR2022

**[26]** Woojeong Jin et al. A good prompt is worth millions of parameters: Low-resource prompt-based learning for vision-language models. arXiv 2021.

**[83]** Zizhuo Zhang et al. Prompt learning for news recommendation. SIGIR 2023.

**[81]** Yaohua Zha et al. Instance-aware dynamic prompt tuning for pre-trained point cloud models. CVPR2023

> **Concern about Quantifying distribution shifts:**

We have revised the manuscript and included all the related content [Appendix H] here for your convenience.

**Performance with respect to Different Difficulty Levels.** Here, we demonstrate the performance of our PURE 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 in Table 7. From the results, we can observe that all the model performs worse in hard scenarios and our method has consistently outperformed 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.

Table 7. Performance comparison across varying levels of OOD generalization difficulty. The values represent the Mean Squared Error (MSE) for each method.
| Method      | U-NET  | ResNet | VIT    | Swin-T | FNO    | CGODE  | PURE(Ours) |
|-------------|--------|--------|--------|--------|--------|--------|------------|
| **Easy**    | 0.0945 | 0.0682 | 0.0654 | 0.0676 | 0.0452 | 0.0772 | *0.0325*    |
| **Mid**     | 0.1063 | 0.0922 | 0.0902 | 0.0912 | 0.0544 | 0.0863 |*0.0341*    |
| **Hard**    | 0.1432 | 0.1234 | 0.1076 | 0.1123 | 0.0623 | 0.0921 |  *0.0354*     |

We hope that these revisions address your concerns. Please let us know if there are any further questions or concerns.

Sincerely,  
The Authors
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.
Thank you for the additional experiments on computational costs.
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:

$$
{z}^t \perp {\mu}^t, u_{\text {output }}=\phi\left(\left[{\mu}^t,  {z}^t\right]\right) . \quad (1)
$$

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:
$$
{\mu}^t=\operatorname{BasicModel}\left({u}_{\text {input }}\right),  \quad (2)
$$

$$
z^0=\operatorname{ContextMining}\left({u}_{\text {input }}\right),\quad (3)
$$

$$
{z}^t=\operatorname{GraphODE}\left(z^0, t\right),\quad (4)
$$

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.
Thank you for your responses. I appreciate the authors provide additional explanation to what I asked. I would like to maintain my rating (4) for the reason that
- I think the manuscript needs significant update to reflect the new information.
- The new information including the problem formulation (OOD in dynamic systems) is too fundamental rather trivial, which should appear at the beginning of the paper (problem setup).
- It is essential in motivating the proposal, including the design of different components, on how they can address the dynamic OOD (on what parameter invariance across training-test).
- Without revision, I fail to connect the problem to the proposal by looking the rebuttal text and the submission forth-and-back.