PDE-constrained Learning with Multi-time-stepping for Accelerated Fluid Simulation

Authors: Qi Wang, Yuan Mi, Wang Haoyun, Yang Liu, Yi Zhang, Ruizhi Chengze, Hongsheng Liu, Ji-Rong Wen, Hao Sun
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
Keywords: physics-informed learning, multiscale time stepping, spatiotemporal dynamics prediction

Abstract:
Solving partial differential equations (PDEs) by numerical methods meet computational cost challenge for getting the accurate solution since fine grids and small time steps are required. Machine learning can accelerate this process, but struggle with weak generalizability, interpretability, and data dependency, as well as suffer in long-term prediction. To this end, we propose a PDE-embedded network with multiscale time stepping (MultiPDENet), which fuses the scheme of numerical methods and machine learning, for accelerated simulation of fluid flows. In particular, we design a convolutional filter based on the structure of finite difference stencils with a small number of parameters to optimize, which estimates the equivalent form of spatial derivative on a coarse grid to minimize the equation's residual. A physics block with a 4th-order Runge-Kutta integrator at the fine time scale is established that embeds the structure of PDEs to guide the prediction. To alleviate the curse of temporal error accumulation in long-term prediction, we introduce a multiscale time integration approach, where a neural network is used to correct the prediction error at a coarse time scale. Experiments across various PDE systems, including the Navier-Stokes equations, demonstrate that MultiPDENet can accurately predict long-term spatiotemporal dynamics, even given small and incomplete training data, e.g., spatiotemporally down-sampled datasets. MultiPDENet achieves the state-of-the-art performance compared with other baseline models, with over 5$\times$ speedup compared to classical numerical methods.

Decision: reject
Reviewer scores: 5, 6, 6, 3, 8, 6

Reviews and discussion:
<thread>
creation date:
1730393078358

thread_type:
official_review

<thread_object>
creator: reviewer_yrnc

<content>
summary:
The authors introduce MultiPDENet. The goal is to accelerate the solution of time-dependent PDEs using deep learning. MultiPDENet has a learnable multiscale timestepper, inspired by Runge-Kutta updates. It also has some interesting inductive biases, such as symmetry constraints on layers that mimic the operation of finite difference stencils. The results show impressive gains over standard neural PDE surrogates from the literature in terms of long term rollout accuracy and stability.

strengths:
**Novelty**

There are several novel aspects of MultiPDENet
- The schematic diagram of how everything fits together is unique
- While learning a timestepper is not unique, I have not since it in the context of PDEs, only ODEs.
- While placing constraints on learnable derivative filters is not unique e.g. **Learning data driven discretizations for partial differential equations** Bar Sinai et al. (2018), the symmetry constraints of the MultiPDENet is unique

**Quality**

The results are really good! This is probably because of the sensible use of inductive bias in the learned solver structure. I think this is the highlight of the paper. In particular:
- The high correlation time plots in Figure 3 b/e/h are highly motivating
- The results in Table 2 show some massive performance improvements over baselines
- The results in Section 4.2 show large improvements over baselines even outside the ICs of the training range, demonstrating robustness to distribution shift
- I found the ablation in Table 3 very informative. This will also be important for anyone wishing to build on this work

**Significance**

For the field of ML4PDEs, any model that can push the envelope of performance forward is significant, and in this sense, MultiPDENet is significant.

weaknesses:
**Clarity**

I found it hard to follow the mathematical description of the model, since many of the terms are not laid out in the main text. I had to refer to the diagrams to understand the general gist of how everything fits together. Please see my questions below for further requests for clarification. I found I was able to understand, at a gut-feeling level, why main design choices were made, but because of the ambiguity in the exact model design I am not sure whether my reading is correct. The clarity is the part of the paper I would implore the authors to focus on most.

**Quality**

One subtle, but very important weakness of the paper is the very small number of training data used. The number of snapshots per equation are as follows:
- KdV: 3000 snapshots
- Burgers: 700 snapshots
- Gray-Scott: 540 snapshots
- Navier-Stokes: 24000 snapshots
For deep learning applications this is extremely small and could explain, to a large degree, just why the performance of MultiPDENet is just so good when compared with the other models, which exhibit much less inductive bias. As is well known, inductive bias helps in the low data regime, such as this. Indeed this could become the angle that this paper takes, but unless I missed it, there is no data ablation in the paper, which would have been a useful comparison

questions:
- Line 55: In what way are PiNNs data-drive? I think the claim that PiNNs are “data-driven”, at least in their vanilla form, is contentious. If you consider collocation points as data then this is somewhat true, but there are no ground-truth targets fed to PiNNs, which would render them data-free. 
- Line 103: One the citations reads “_Transformers Cao (2021)_”. I think this should read “_Cao (2021)_”.
- Equation 1: Please explain all the mathematical terms in full.
  - What is $\tilde{\textbf{x}}$?
  - Why is $\tau$ a dummy variable in the integral, but also found outside the integral?
  - Is $\Delta$ a difference operator or a Laplacian?
  - What is $\mu$
  - What do $i$ and $j$ correspond to? If timestep indices, how are they different from $k$?
  - Referring back to Figure 1.b), please explain what $k+\delta$ means? Is $\delta$ a fractional time index? What values does $\delta$ take?
- Equation 3:
  - What is $\lambda$?
  - What values can $m$ take on? Does this imply that Equation 3 is an $m$-step block? How does it correspond to Figure 1a)? Is Figure 1a) the Runge-Kutta style discretization of the integral?
  - In Equation 1 $\mathcal{P}$ accepts argument $\textbf{u}$, but in Equation 3 it also includes $\Delta \textbf{u}$ and $\lambda$.
  - What is the significance of square brackets?
- Figure 2: Please expand on the symmetry constraints for the central difference stencils. We see even symmetry in the vertical direction and odd/even symmetry in the horizontal direction. If you are working in 2 dimensions, what does the mixed derivative operator $\partial_{xy}$ looks like?
- Line 233: e -> We
- Line 269: What is $T$? Is it time? How is $T$ different from $t$?
- Figure 3: b)e)h).
  - Between each plot the colors of the curves per model change, please fix this.
  - In e) you write "Ours"; whereas, in b) and h) you write "MultiPDENet". Please choose one.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730602379571

thread_type:
official_review

<thread_object>
creator: reviewer_pgra

<content>
summary:
The paper introduces MultiPDENet, a novel neural network architecture designed to accelerate simulations of fluid dynamics by integrating machine learning with partial differential equation (PDE) constraints. The architecture uses a multi-scale time-stepping approach and a learnable convolutional filter that embeds physical equations directly into the network, enhancing both accuracy and computational efficiency. The numerical findings of MultiPDENet demonstrate the accuracy and generalizability across different initial conditions, Reynolds numbers, and external forces. The inclusion of correction blocks enhances its long-term stability and robustness. An ablation study validates the effectiveness of the model’s components, highlighting its potential as a powerful tool for fluid simulations and complex PDE-constrained problems. Overall, this paper is well-written and could be a significant algorithmic contribution. I am willing to recommend accept. However, there are still some comments and questions need to be clarified or further discussion,  see the weakness and questions.

strengths:
1. Originality: the paper introduces an innovative approach by integrating a physics-informed neural network (PINN) with a multi-scale time-stepping mechanism and a learnable filter, making it distinct from traditional numerical methods and existing data-driven models. This combination of deep learning and PDE constraints is a novel contribution that enhances the model's interpretability and applicability in fluid dynamics simulations.
2. Quality: this work is thorough, with comprehensive experiments demonstrating MultiPDENet’s performance across various PDE systems. The detailed ablation study provides strong evidence that each component of the model contributes meaningfully to its success. The results are rigorously compared against baseline methods, showcasing significant improvements in accuracy and computational efficiency.
3. Clarity: the paper is well-structured, providing clear explanations of the methodology, architecture, and experimental setup. The figures and tables are informative and support the textual content effectively. The presentation of the multi-scale time-stepping framework and the incorporation of physics into the network are clear.
4. Significance: the findings of the paper have potential implications in the field of computational fluid dynamics and PDE-constrained modeling.

weaknesses:
1. The model is limited to regular grids due to the use of convolutional operations. This restricts its applicability in cases involving irregular or adaptive grids.
2. The current study only addresses 1D and 2D problems, which limits its applicability to more complex 3D simulations. Examples of 3D fluid simulations are more convincing.
3. Although it is mentioned that the method works with sparse data, the paper could elaborate more on the data requirements for practical cases with very noisy or incomplete data.
4. The inclusion of multi-scale architectures, learnable filters, and correction blocks adds complexity to the model. This could impact the ease of implementation and adoption by other researchers or practitioners, especially when compared to more straightforward neural network models. The author should address this and provide open-source code for reproducibility.

questions:
1. On page 5 line 233, "e divide the neural network..." what is statement? 
2. On page 9 line 478, "This demonstrates our model..." this claim is questionable. The current test has Re = 4000. To demonstrate the capability of this approach for high Re turbulence, a 3D flow with Re of at least 10^4 would be necessary to substantiate this claim.
3. The training data collected are from uniform timestep (\Delta t). However, many fluid simulations have non-uniform timestep. How does the model react to these scenarios?
4. The paper mentions future work involving graph neural networks for irregular grids. How feasible is this adaptation, and what challenges might arise in integrating GNNs with the current architecture?
5. Are there foreseeable issues when extending the current approach to higher dimensions?
6. How sensitive is the model’s performance to hyperparameters, such as the choice of kernel size for the learnable filter or the learning rate for training?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730602979393

thread_type:
official_review

<thread_object>
creator: reviewer_nnlj

<content>
summary:
The authors propose a hybrid method for solving PDEs, combining elements of classical numerical methods with learnable parameters.  Elements of classical numerical methods include an RK4 time integrator, finite-difference derivatives, and a Poisson solver for pressure.  The learned correction block uses a Fourier Neural operator. The problem is tested on 1D and 2D PDEs with periodic boundary conditions.  The authors claim that using a learned correction at a coarse time scale improves the prediction error and provides "53x speedup" compared to classical numerical methods.

strengths:
This paper proposes an architecture to solve PDEs. Overall I cannot rate this paper strongly and am unable to offer a view of strengths because I find much of the presentation unclear.

weaknesses:
The approach here is similar to other hybrid classical+learned approaches.  The presentation of the work is rather unclear in my view.  It is not explained in detail what the micro (MiNN) and macro (MaNN) blocks do exactly and what the difference is between them.  The authors refer to "sparse grids" on occasion, and it is unclear what is meant because the authors also use the separate term "coarse grid".

The claim that this method offers a 53x speedup compared to traditional numerical methods is not backed up by evidence, nor details of how it is measured.  Overall, what is here is not sufficiently convincing or compelling.

questions:
Please provide details of how the speedup compared to traditional numerical methods is obtained.  Are the same resolutions being used for both?  Is the same type of device being used, or does this a comparison where one simulation is on CPU and another on GPU?  Is accuracy being checked at all for this comparison?  Please see https://www.nature.com/articles/s42256-024-00897-5.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730635031645

thread_type:
official_review

<thread_object>
creator: reviewer_6waj

<content>
summary:
This paper presents MultiPDENet, a physics-constrained neural network framework developed to accelerate the simulation of complex spatiotemporal dynamical systems, particularly for fluid dynamics applications. The paper addresses two key challenges: long-horizon prediction and generalization. To tackle these issues, MultiPDENet combines numerical schemes with deep learning through a multiscale time-stepping approach, which includes both fine-scale and coarse-scale predictions, effectively mitigating temporal error accumulation in long-term forecasting.

[-] Claims on Spatial Derivative Accuracy (Abstract, line 019): The authors state that MultiPDENet "enables accurate estimation of spatial derivatives on coarse grids." However, there is no theoretical analysis or empirical evidence provided in the main text or experiments to support this claim. Additional analysis or experimental validation is needed to substantiate this assertion. Actually, this claim appears many times in the paper.

[-] Oversimplification of Method Section (Section 3.2): The method description in Section 3.2 is oversimplified, with essential details, including the Poisson and correction blocks, relegated to the appendix. These blocks are central to the model's architecture and should be introduced in the main text to improve readability and clarity. Readers currently need to constantly refer to the appendix for critical information about the methodology.

[-] Computational and Memory Cost Analysis: Since the network architecture includes the computation of partial derivatives (especially second-order derivatives), this likely increases computational complexity and memory requirements. To provide a more comprehensive evaluation, the authors should compare the time and memory costs of MultiPDENet with those of other baseline models, during both training and testing. Additionally, reporting the parameter count for each method would help control variables, clarifying whether the performance gains result from the proposed architecture or from increased computational resources and model size.


[-] Framework Independence and Network Architecture Variability: Given that the proposed framework is architecture-agnostic, it would be beneficial for the authors to demonstrate that MultiPDENet performs well with different backbone architectures (such as FNO, UNet, or Transformers). This would further validate the framework's flexibility and effectiveness across architectures. As operator learning progresses rapidly, incorporating more recent models, such as transformer-based operators, as baselines would strengthen the paper's relevance and rigor.

strengths:
See above

weaknesses:
See above

questions:
See above
</content>
</thread_object>

</thread>
<thread>
creation date:
1730702742226

thread_type:
official_review

<thread_object>
creator: reviewer_xzfb

<content>
summary:
This paper proposes an interesting hybrid neural learner for fluid problems. The neural network builds the PDE in the pipeline, instead of using it as a loss function. A time-stepping scheme is applied explicitly instead of applying the NN directly in an autoregressive fashion.

strengths:
The method proposed is pretty straightforward to implement and uses the physics (PDE) in a good way, not just simply throwing the PDE residual into a loss function and calling it for a day. The numerics are competitive among end-to-end models.

weaknesses:
- The writings are pretty bad, there are many typos, e.g., "spatialemporal". The word choices and phrasing definitely read pretty weird at places, for example, the term "promotion" in Table 2. I felt that this paper could really use proof-reading from a native speaker.
- As someone from a background of many years of training in numerical analysis, I would not name the NN in consideration a "PDE-constrained" NN. The reason is that in traditional numerical methods, when talking about "constraint", it refers to the fact that the method of interest imposes the constraint ***exactly*** up to the machine eps. For example, the divergence-free constraint in NS.
- These micro and macro blocks $M_iNN$ and $M_aNN$ are not properly defined. After digging into the figures (S1 to be specific) and Section 3.2.4, I still do not know what their exact architectures are. For example, what hyperparams are used in the FNO and UNet? Why not use UNet as $M_iNN$ and FNO as $M_aNN$?
- $\mathcal{B}$ is not clearly defined either. Given the form of (S8), one would guess it is $\mathcal{F}$ plus the external forcing but (S6) states otherwise.
- The Poisson block is essentially a pressure projection scheme, which is nothing new and has been introduced back by Chorin in 1967 and by Shen later in 1980s. 
- The ablation lacks some in-depth studies, examples include: 
  - in theory, making the FD operator as learnable filters sounds pretty good, but do the learned parameters actually replicate the FD operator and have a truncation error of four-order accuracy?
  - instead of simply removing different components of the model in 4.6, I think it is better to compare the RK4 used with another time stepping scheme (with a less truncation error order) to test whether it is the RK4 adopted really works, or the pipeline is really robust with respect to any time-stepping schemes.
- Some missing references on making the FD stencil learnable as a convolution filter, e.g., a paper by Kossaczká-Ehrhardt-Günther and arXiv:2201.01854. There is also arXiv:2003.09573 on using NN as correction to a time stepping scheme.

questions:
I don't understand "FD method can yield inaccurate derivatives on sparse grids." What does "sparse grids" mean here?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730721318358

thread_type:
official_review

<thread_object>
creator: reviewer_925d

<content>
summary:
The authors propose to learn fluid simulation with multi-step predictions.

strengths:
- Using numerical integrator with neural networks is an interesting approach.
- The reported experimental results show improvement over the reported baselines.

weaknesses:
- The multi-step prediction naturally makes the network to be much deeper. For example, the computation will go through at least 4x layers with 4 micro step prediction, albeit the weight sharing. However, the benchmarked baselines are only single networks which are not as deep. It has been shown that using more stacks of shared-weight networks can greatly increase the performance [1, 2]. Given the big improvement in the reported metrics, the authors need to compare their methods with stronger baselines with more layers. 

- Another related point is regarding the numerical integration. The numerical integration is naturally similar to the simple residual network. The authors need to benchmark the proposed solution with a network with the same amount of network blocks but without using the RK4 integration, e.g., by stacking them or using the simple residual connections.

- Since the paper focuses on flow problem, it can benefit from testing larger flow benchmarks, such as PDEArena [3].

[1] Tran, Alasdair, et al. "Factorized fourier neural operators." ICLR 2023.

[2] Zhang, Xuan, et al. "SineNet: Learning Temporal Dynamics in Time-Dependent Partial Differential Equations." ICLR 2024.

[3] Gupta, Jayesh K., and Johannes Brandstetter. "Towards Multi-spatiotemporal-scale Generalized PDE Modeling." Transactions on Machine Learning Research.

questions:
- For Eq. 3, does it mean that inside each PDE block there will be multiple network evaluation to perform the RK4 integration? E.g., there will be multiple steps for each micro step.

- For Eq. 2, how are the derivative quantities also predicted by the network at the coarser grid? Where is the poisson block defined in Eq. 2?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732279585514

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General response

comment:
Dear Reviewers,

We deeply appreciate the insightful and constructive comments from you, which are helpful in improving our paper. We are pleased that all the reviewers recognized the *novelty* and *excellent generalizability* of our work. In particular, we thank the reviewers for recognizing the ***interesting hybrid approach*** (925D, XzFb), ***improved results*** (925D, PGra, YrNc), and ***novelty*** (yrNC, PGra) of our method.

In addition, we have thoroughly proofread our paper, corrected typos and grammatical errors, and reorganized the content to enhance its clarity. **Substantial revisions**, highlighted in red color, have been made in the updated version of the paper (please see ***the updated .pdf file***). A list of the main revisions include:

- Re-wrote and re-orgnized the majority of the paper (all sections).
- Added additional references to the related work section (Section 2, Page 3). 
- Provided a clearer description of the model architecture (Section 3.2.1, Page 4). 
- Refined the explanation of the Physics block (Section 3.2.2, Page 4). 
- Clarified the description of the NN block (Section 3.2.4, Page 5). 
- Improved the statements regarding the experiments (Section 4, Page 6). 
- Analyzed the scalability of the MiNN block and MaNN block (Appendix C.2, Page 18). 
- Clarified the inference cost of our approach (Appendix F.2, Page 22). 
- Added additional experimental results, e.g., ablation study, scalability tests, computational cost, etc. (Section 4.6, Page 10; Appendix Section C.2, Page 19; Appendix Section F.2, Pages 22-23).

Thank you once again for your thoughtful reviews. Please do feel free to let us know if you have any further questions. 

Best regards,

The Authors of the Paper
</content>
</thread_object>

</thread>
<thread>
creation date:
1732508502621

thread_type:
official_comment

<thread_object>
creator: area_chair_hjmg

<content>
title:
Reviewers' Response

comment:
Dear Reviewers,

As the author-reviewer discussion period is approaching its end, I would strongly encourage you to read the authors' responses and acknowledge them, while also checking if your questions/concerns have been appropriately addressed.

This is a crucial step, as it ensures that both reviewers and authors are on the same page, and it also helps us to put your recommendation in perspective.

Thank you again for your time and expertise.

Best,

AC
</content>
</thread_object>

</thread>
<thread>
creation date:
1733214798430

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Thanks to the reviewers

comment:
Dear Reviewers,

We would like to thank your engagement in the author-reviewer discussion period. Special thanks are given to Reviewers PGra, XzFb and yrNc for your positive feedback. Your constructive comments and suggestions are greatly helpful in improving the quality of our paper.

In the meantime, Reviewers 6Waj and NnLJ kept silent during the discussion period, despite our repeated reminders. However, we are confident that your concerns have been fully addressed through detailed clarifications and additional experimental results.

Overall, we appreciate your time and effort placed on reviewing our paper. Thank you very much!

Best regards,

The Authors
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper introduces MultiPDENet, a neural network architecture designed to accelerate fluid dynamic simulations by combining classical numerical methods, such as a multi-scale time-stepping scheme inspired by Runge-Kutta methods, finite-difference derivatives, and a Fourier Neural Operator for learned corrections. This approach embeds physical constraints from partial differential equations (PDEs) directly into the network. The model reportedly demonstrates significant speedups (53x) and improved long-term stability and accuracy compared to both traditional numerical solvers and other neural PDE surrogates. The method is validated on several one- and two-dimensional problems.

The original paper contained several inaccuracies, making it difficult to read. Regarding the content, some reviewers agreed that the paper was not clearly articulated. The authors introduced many different concepts, making it unclear what the individual enhancements actually are (see comment by reviewer 6Waj). Additionally, as mentioned by 6Waj, some parts of the paper seem to lack a logical thread, making the innovations difficult to assess. Using convolutional networks to learn discretizations of differential operators is not a new concept. In addition, the pressure correction step the authors used (as mentioned by reviewer XzFb) is far from new.

One of the paper's main claims, a "53x speedup" over classical methods, appears to have been revised down to 5x. However, further examination of the baselines, raised by many reviewers, suggests that they are still quite weak (see [1]). Solving the Navier-Stokes equations on a torus can be performed extremely fast with spectral methods using a vorticity formulation, as demonstrated in [2], which the authors cite. Extrapolating from Figure 7b to Figure 1 in [2], it appears that similar performance to the authors' 1024x1024 finite volume (FV) solver can be achieved with a 128x128 grid. Therefore, the comparison is misleading, as one needs 64x fewer degrees of freedom (with an extra factor of 2 from the vorticity formulation) and 8x fewer time steps (so we can "expect" the solver to be much faster at that resolution). Additionally, given that the authors also used a spectral Poisson solver for their implementation of the Navier-Stokes emulator, the comparison with an FV-based method seems unfair.

Regarding the claim of long-term stability, the authors primarily present trajectory-wise errors. However, they provide no metrics demonstrating that the trajectories follow the correct statistics for very long rollouts (see [3, 4] for examples of such metrics). As such it is not clear how to assess that claim, given that several new formulation have been develop to tackle the long-term statistics of chaotic/fluid systems.

Furthermore, the overall solver remains subject to the CFL condition, which limits the time steps. Thus, the solver can likely only remain competitive in the coarse-scale regime. This limitation should be clearly acknowledged in the paper's claims unless the authors can provide evidence to the contrary.

Given the numerous issues raised by the reviewers, which the updated version failed to fully address, I recommend rejection.

[1] McGreivy, Nick, and Ammar Hakim. "Weak baselines and reporting biases lead to overoptimism in machine learning for fluid-related partial differential equations." Nature Machine Intelligence (2024): 1-14.

[2] Dresdner, Gideon, et al. "Learning to correct spectral methods for simulating turbulent flows." arXiv preprint arXiv:2207.00556 (2022).

[3] Jiang, Ruoxi, et al. "Training neural operators to preserve invariant measures of chaotic attractors." Advances in Neural Information Processing Systems 36 (2024).

[4] Schiff, Yair, et al. "DySLIM: Dynamics Stable Learning by Invariant Measure for Chaotic Systems." Forty-first International Conference on Machine Learning (2024).