MultiPDENet: PDE-embedded Learning with Multi-time-stepping for Accelerated Flow Simulation

We appreciate your constructive comments. To enhance clarity, we have thoroughly proofread the manuscript and corrected all identified typographical errors. We believe these revisions will significantly improve the presentation. The updated version will be uploaded as soon as file submissions are enabled.

Weaknesses

W1. The benchmarks include few recent pure data-driven models.

Reply: Great remark! Our proposed MultiPDENet integrates PDE priors with multiscale time-stepping scheme, enabling efficient spatiotemporal simulations (e.g., turbulence) on coarse grids with limited data. While recent pure data-driven models (e.g., modern-UNet [1]) perform well in data-rich scenarios (e.g., requiring 128 trajectories), their performance deteriorates significantly in small data settings.

While we have selectd widely recognized data-driven baselines (e.g., FNO, modern-UNet, DeepONet), we have also considered hybrid physics-learning models (e.g., LI, TSM, PhyFNO, PeRCNN) for comparison. Following your suggestion, we added the recent model ONO [2] as an additional baseline. Its inferior performance under data scarcity (see Table A) further highlights the robustness of our framework.

Table A: Comparison of MultiPDENet and baselines for NSE.

W2. The selection of baselines for varying experimental cases.

Reply: Insightful comment! Our baseline selection carefully balances pure data-driven models (e.g., FNO, UNet, DeepONet) and physics-inspired counterparts (e.g., PhyFNO, LI, TSM, PeRCNN) to ensure fairness and relevance.

To ensure consistent comparison, we added PhyFNO as a baseline for the Burgers and GS equations, with results shown in Table B below. So far, we have maintained a unified setting of baseline comparsion (FNO, UNet, DeepONet, PhyFNO, PeRCNN) for KdV, Burgers, and GS equations. In the case of NSE, we slightly altered the baseline models by considering two other well-recognized models specifically tailored for NSE (aka, LI and TSM). In all cases, we have demonstrated that MultiPDENet surpasses the baselines.

We believe the current baseline setup ensures a fair and consistent comparison. Thanks for your great suggestion!

Table B: Comparison of MultiPDENet and PhyFNO for Burgers and GS.

Suggestions

S1. Typos of "Sovling PDE".

Reply: We have thoroughly proofread the paper and corrected all typos in the revised version.

Questions

Q1. Design motivation for the Correction Block and the effectiveness of state correction.

Reply: Great remark! The Correction Block is designed to mitigate the information loss introduced by resolution reduction before computing the equivalent derivatives, enabling the model to adapt to the coarse grid. Instead of explicitly recovering the high-resolution solution, the Correction Block implicitly corrects the coarse-grid state by adjusting the scaling of the equivalent derivative term. During training, this block focuses on minimizing the overall PDE residual rather than directly reconstructing fine-grid details. The state $\hat{\bar{\mathbf{u}}}{_m^k}$, obtained via the Correction Block, represents a neural-corrected version of the coarse solution. However, this correction should not be interpreted as an explicit recovery of the fine-grid solution. Rather, it ensures that the equivalent derivative term computed via the Symmetric Filter is optimally adjusted to minimize the overall PDE residual. Our training objective is solely aimed at making the results more closely approximate the ground truth solution, as demonstrated in the ablation study (Model-E in Table 3 on Page 8). Therefore, there are no additional training objectives specifically designed for correction beyond this implicit adjustment mechanism.

Refs:

[1] Gupta et al. Towards Multi-spatiotemporal-scale Generalized PDE Modeling. TMLR, 2023.

[2] Xiao, et al. Improved Operator Learning by Orthogonal Attention. ICML, 2024.

Remark: Once again, we sincerely appreciate your constructive comments. Looking forward to your feedback!

Thanks for your constructive comments and suggestions! We have carefully addressed them, and the following responses have been incorporated into the revised paper.

Questions

Q1. Essential References Not Discussed.

Re: We appreciate your comment. We have noted that our initial submission omitted several important black-box methods in ICLR 2022, such as MP-PDE [1], GMR-Transformer-GMUS [2], and SiT [3]. We will include these references in the related work section in the revised manuscript.

Q2. Is the white-box modules and black-blox moduls applied on the coarse mesh for NSE?

Re: Great remark! To accelerate the prediction of spatiotemporal dynamics on coarse grids, we integrated the numerical scheme with neural networks through a multiscale time-stepping strategy. Consequently, the entire network, encompassing both white-box and black-box modules, is designed to operate on a unfied coarse mesh.

Q3. Does the model need to satisfy the CFL condition?

Re: Insightful question! According to the CFL condition, the maximum allowable time step is given by: $\delta t _{\max}=\mathrm{CFL} _{\max} \cdot \frac{\delta x}{\left| u \right| _{\max}}$, where the standard choice for numerical simulations is $\mathrm{CFL} _{\max}=0.5$. Based on this, we obtained $\delta t _{\max} = 0.007$ for the coarse grid with $\delta x = 2\pi/64$, which indeed means that our model satifies the CFL condition in our original test setting.

However, to further investigate the extent to which MultiPDENet can go beyond the CFL condition, we conducted experiments with increased time intervals, namely, $\delta t_{\max} = 0.028$ (four times the required smallest time step). The resulting model is called MultiPDENet-L. With this setup, we achieved an extra 4$\times$ speedup (e.g., inference time of 6 sec vs. the original 26 sec) while maintaining a similar model accuracy, aligning with the speed of other baseline models as shown in Table A. Note that MultiPDENet-L was trained based on a rollout strategy over 8 macro-steps.

In summary, MultiPDENet does not need to adhere to the CFL condition, demonstrating its capability to operate effectively beyond conventional stability constraints.

Table A: Comparison of MultiPDENet and baselines for NSE

Refs:

[1] Brandstetter et al. Message passing neural PDE solvers. ICLR, 2022.

[2] Han et al. Predicting Physics in Mesh-reduced Space with Temporal Attention. ICLR, 2022.

[3] Shao et al. SiT: Simulation transformer for particle-based physics simulation. ICLR, 2022.

Remark: Once again, we sincerely appreciate your constructive comments. Please feel free to let us know if you have any further questions. Looking forward to your feedback!

Thanks for your constructive comments! We have addressed them thoroughly and added new figures/tables (see the rebuttal.pdf via https://anonymous.4open.science/r/Rebuttal-5D8B/rebuttal.pdf). These results will be added to the revised paper.

Weaknesses

W1. Uniform grids.

Re: Thanks for your thoughtful comments on irregular geometries and complex boundary conditions (BCs). Our study focuses on fluid simulations on regular grids with periodic BCs, a common constraint shared with methods like FNO, LI, and TSM. This indicates that the limitation is not specific to MultiPDENet but represents a general challenge in the field. Similar to Geo-FNO [1], a geometry-aware netwok (e.g., a geometry encoder-decoder) could be established on top of MultiPDENet (latent structured geometry learner) to handle general geometries. Thank you for pointing out the important direction for our future work!

To demonstrate our model's generalization to complex BCs, we conducted experiments on Burgers equation with Dirichlet and Neumann BCs, while keeping other settings consistent with the original data generation setup. We tested our previously trained model for inference on 10 trajectories with complex BCs through BC encoding (Table R1 in rebuttal.pdf). Figure R1 in rebuttal.pdf shows predicted snapshots, confirming the model's generalizability over complex BCs.

W2. Lacking 3D cases.

Re: Great remark! We followed your comment and tested our model on the 3D Gray-Scott (GS) equation. We generated 5 datasets (1 for training and 4 for testing).

Snapshots of trajectory evolution from 0 to 600 s for MultiPDENet and baselines are shown in Figure R2(a) in rebuttal.pdf. Figure R2(b) in rebuttal.pdf shows MultiPDENet maintains a Pearson correlation coefficient $>$0.8 throughout the evolution. The error distribution in Figure R2(c) in rebuttal.pdf highlights our model’s superior performance. Table R2 in rebuttal.pdf summarizes our model’s performance suggesting strong potential to solve 3D problems.

Questions

Q1. Model inference speed.

Re: Thanks for raising this important point! It's crucial to clarify that the reported orders of magnitude speedups for popular models like FNO, LI and DeepONet (10$^3\times$, 80$\times$, 24$\times$) are often not evaluated against numerical methods with comparable accuracy. McGreivy et al. [2] revealed this by reimplementing these models and comparing them under consistent precision, resulting in only 7$\times$ speedup for FNO, while LI and DeepONet exhibited slower performance. Their evaluation effectively highlighted the prior discrepancies in speedup reporting.

Consistent with [2], we demonstrate a 5–7$\times$ speedup under the condition of sonsistent accuracy, althouh implementation in JAX may further yield extra speedup gains (e.g., $>5\times$) [3]. We must admit that the rollout strategy used in our model inherently limits its speed, while reducing the network size provides only marginal gains. Nonetheless, MultiPDENet generalizes well across diverse initial conditions, varying $Re$, external forces, and larger domains.

Q2. Does MultiPDENet still outperform other models at the same inference speed?

Re: Great question! As mentioned above, solely reducing network parameters does not substantially accelerate the inference due to the employed rollout strategy. However, leveraging a larger time step $\delta t$ proves effective. MultiPDENet's multiscale time-stepping design allows it to circumvent the CFL condition, ensuring accuracy and stability even with increased $\delta t$. Specifically, with $4\delta t$ (MultiPDENet-L), we achieved an extra 4$\times$ speedup (e.g., inference time of 6s) while maintaining a similar model accuracy, aligning with the speed of other models as shown in Table A. Note that MultiPDENet-L was trained based on a rollout strategy over 8 macro-steps.

These results demonstrate our model’s multiscale time stepping scheme consistently outperforms all baselines across all metrics, even at similar inference speeds. Hope this clarifies your concern.

Table A: Comparison of MultiPDENet and baselines for NSE

Refs:

[1] Li et al. Fourier neural operator with learned deformations for pdes on general geometries. JLMR, 2023

[2] McGreivy et al. Weak baselines and reporting biases lead to overoptimism in machine learning for fluid-related partial differential equations. Nature Machine Intelligence, 2024.

[3] Takamoto et al. PDEbench: An extensive benchmark for scientific machine learning. NeurIPS, 2022.

Remark: Please let us know if you have other questions. Looking forward to your feedback!