Paper: GNRK: Graph Neural Runge-Kutta method for solving partial differential equations (ICLR-2024)

Reviewer comment (Novelty):
I struggle to see a motivation for using a neural network model to replace the evaluation of f in Eqn. (1) in the RK framework. The RK scheme described by Eqn. (3) is a standard explicit time integration scheme, and a numerical method (finite difference, finite elements, etc.) for evaluating f in Eqn. (1) can be highly efficient using trivial parallelization on GPUs. Numerical methods also have great generalizability as they can easily handle different initial conditions, physical parameters, spatial discretization, etc. in a principled way. After reading the paper, it remains unclear to me what the advantages of the proposed method over numerical simulation are.

Author response:
$\Rightarrow$ The advantage of replacing $f$ with a neural network in the GNRK lies in scenarios where the exact form of the governing equation is unknown.
In cases where the equation $f$ is not precisely defined, numerical methods may face challenges, while the GNRK can be trained using data alone, without prior knowledge of the governing equation. This flexibility allows the proposed method to adapt to situations where traditional numerical simulations struggle due to the lack of a well-defined analytical form for the function $f$.

Reviewer comment (Experiments):
Following my comment above, I feel the method should be compared with a high-quality numerical simulator (finite difference/element/volume, etc.) and report their runtime difference. I skimmed over the Burges equation part in the code, and it looks like the training data were generated using finite differences implemented in numpy. Was it faster or slower than the network inference time in the proposed method?

Author response:
$\Rightarrow$ In response, we have included the wall clock and GPU memory usage for the GNRK and other baseline models in Table. 1 for a more comprehensive comparison.
Regarding the implementation of numerical simulations, we acknowledge the use of a naive finite difference method with numpy, without GPU acceleration or mroe optimized methods due to the time constraints.
While the GNRK may exhibit slower wall clock times compared to typical numerical solvers, it offers a unique advantage in situations where little information about the governing equation is available, as highlighted in the previous response.

Reviewer comment (Experiments):
I feel the comparison with other neural network baselines is not well-balanced in that it exaggerates their weaknesses but ignores their strengths in other aspects. Taking PINN as an example, PINN has been tested against much more complicated PDEs in fluid and solid mechanics (e.g., Navier-Stokes) which explicit time integrations like RK4 struggle to solve, and I don’t see an easy way for extending the proposed method to these PDEs given that it is coupled with RK time integration. Another thing is that PINN can be trained with partial observations of the spatiotemporal states sampled irregularly in time and space, so it seems to impose fewer requirements on the training data than the proposed method. Of course, it is unnecessary to beat all baselines in all aspects before publishing a research paper, but I think a more balanced view and experiment design would make it more convincing.

Author response:
$\Rightarrow$ As you pointed out, for chaotic systems such as Navier-Stokes, error can accumulate when using recursive time integration.
In the case of the PINN, where no time integration is involved, it can predict arbitrary time without accumulating error.
However, applying the PINN to such complex PDEs, especially those with coupled dynamics like 2D Burgers' equation, presents challenges.
Training a neural network to predict the solution directly, as PINN does, becomes intricate, especially as the spatial domain dimension increases.
The exponential increment in required training points for PINN with higher dimensions, along with challenges in balancing data points at boundaries and non-boundary regions, makes it less straightforward.
Moreover, it is not applicable to Dataset I (different initial conditions), and Dataset II (different PDE coefficients)

Reviewer comment (Writing):
The title of the paper includes "PDE," but the experimental results mainly focus on 2D Burgers' equations as PDE, while all other experiments seem to center on coupled ODE systems. Also, the phrase in the abstract, "often confined to specific PDEs under certain constraints," seems somewhat unclear in its meaning. There are different numerical schemes (e.g., FDM, FEM, or more specialized variations) for applying numerical analysis methods to different types of equations. It appears that neural network models have various methods for different types of PDEs. More detailed explanations in the paper title and abstract might be needed to better represent this research.

Author response:
In response to your feedback, we have refined the the second sentence of abstract for better clarification.
> However, they may perform poorly or even be inapplicable when deviating from the conditions for which the models are trained, in contrast to classical PDE solvers that rely on numerical differentiation.

We hope this modification enhances the clarity and better represents the research focus.

Reviewer comment (Novelty):
The paper lacks a clear explanation of how the novelty of the GNRK model differs from existing models. It is essential to clarify where GNRK offers advantages over previous models concerning the combination of numerical analysis and graph neural networks. The novelty mentioned on page 2 and the four categories discussed in Section 3 require a precise comparison with prior research. For instance, the advantage related to the initial condition might not be unique to GNRK but rather a common benefit shared with all models utilizing graph neural networks. Additionally, the statement that learning with changing PDE coefficients is not efficient, as mentioned in Section 3, should be supported by evidence. In light of this, focusing on the advantages of GNRK over other neural network models, considering multiple aspects of operator learning models and models that learn solutions varying over time, and creating a table that clearly illustrates what each model can and cannot do would be beneficial. This table could emphasize where GNRK excels in comparison.

Author response:
We appreciate your attention to the novelty of the GNRK and the need for a clearer explanation of its advantages over existing models.
However, there seems to be a slight misunderstanding regarding the Table. 1, which outlines the applicability and performance of various approaches.
We would like to clarify that the '-' in the table indicates instances where a model cannot be applied to a particular dataset.

Reviewer comment (Writing):
More detailed explanation is needed for the phrase 'less dependent on the form of PDEs or conditions.' at page 1. What does it mean?

Author response:
We have modified the sentence to explicitly state that classical solvers like FEM, FDM are generic, which can find the solution of a PDE regardless of its form, initial and boundary conditions.
> Unlike NN solvers, classical solvers such as finite-element method (FEM), finite-difference method (FDM), and finite-volume method (FVM) can find the solution of the PDE regardless of its form, initial and boundary conditions, given an exact PDE.

Reviewer comment (Theory):
In Section 2.1, it would be helpful to clarify how GNRK can handle non-uniform discretization. For example, if there are data points at 1 second, 2 seconds, and 5 seconds, how can GNRK capture the information that the intervals between 1 and 2 seconds and between 2 and 5 seconds are different?

Author response:
To clarify, the simulation starts with a pre-defined randomized $\{ \Delta t \}$.
The detailed explanation of how GNRK captures nonuniform time intervals is provided in Section 3.
We highlight that the GNRK captures time information using $\Delta t$.
Specifically, $\Delta t$ affects the Runge-Kutta scheme, as illustrated in Fig. 2(a), while the trainable neural network $f_\theta$ remains independent of $\Delta t$.

Reviewer comment (Theory):
In the middle of page 3, there is an assumption about "same degree." Is this assumption absolutely necessary, or can GNRK be applied without it? You've clarified that GNRK is applicable in cases like the left graph in Figure 1, but what about scenarios like the right graph in Figure 1? Can GNRK be used there as well?

Author response:
In short, the assumption of 'same degree' is not necessary for the GNRK.
GNRK can be applied to scenarios where the degrees of nodes are different, as illustrated in the right plot of Fig. 1.
The term `grid' refers to a special case where all degrees are equal, as shown in left plot of Fig. 1.
We have presented experimental results on general graphs, which has varying degrees, in Section 4.2.

Reviewer comment (Experiments):
In the Burgers experiments in Section 4.1, it would be interesting to know how the model in Figure 2 adapts and learns GNRK as nu changes. This is one of the novelties mentioned in Section 3 but is not detailed in the main text.

Author response:
In the GNRK model, $\nu$ is included as one of the input features.
Therefore, as $\nu$ changes, the structure of the GNRK remains constant.
The model adapts to variations in $\nu$ by adjusting the purple $\mathbf{g}$ component in Fig 2(b).

Reviewer comment (Theory):
Traditionally, the Runge-Kutta method comes with strict guarantees in accuracy, convergence, stability, etc. But with a GNN as the estimator in a timeslot, such theoretical properties are lost. So I wonder if it is still necessary or beneficial to utilize the Runge-Kutta scheme with GNN when handling a practical PDE problem.

Author response:
$\Rightarrow$ Classical numerical solvers, including the Runge-Kutta method, provide rigorous guarantees when the governing equation is precisely known.  
The GNRK, on the other hand, offers a unique advantage by allowing training and prediction without prior knowledge of the PDE.  
This is particularly beneficial in scenarios where the exact governing equation is unkown, but relevant observational data and parameters are available.  
While the GNRK may not inherit the theoretical guarantees of classical solvers, its flexibility in handling unknown PDEs makes it a valuable tool in situations where classical solvers might face limitations.

Reviewer comment (Theory):
In the conclusion section, it is claimed that "it is invariant to spatial and temporal discretization". But I don't see how the proposed method can guarantee such invariance capability.

Author response:
$\Rightarrow$ In respose, we modified the phrase to make it clearer as follows:  
`it can be applied to any arbitrary spatial and temporal discretization, allowing ...'  
The claim of invariance to spatial and temporal discretization is based on the inherent properties of the GNRK outlined in Section 3.  
- Spatial discretization: The GNRK handles spatial discretization by considering the coordinates of each grid point as input features for nodes or edges.  
The parameter sharing mechanism of GNN allows GNRK to effectively capture variations in these features, making it adaptable to different spatial discretization schemes.  
- Temporal discretization: The GNRK views time as the time interval $\Delta t$.  
Importantly, this temporal information only affects the Runge-Kutta scheme (Fig 2(a)), and not the trainable neural network $f_\theta$.  
Therefore, the GNRK remains accurate even when dealing with arbitrary $\Delta t$ in temporal discretization.

Reviewer comment (Theory):
I am wondering which numerical method suits the situation better, the Runge-Kutta method or the multi-step method. And I would like to know if the authors have some comments on this.

Author response:
$\Rightarrow$ The choice between the Runge-Kutta method and multi-step method may depends on the specific characteristics of the problem being sovled.  
In the context of the paper, the Runge-Kutta method is used for a representative numerical solver.  
It is important to note that the GNRK approach is designed to be versatile, and there is potential for its application to other numerical solving methods.  
If you have a specific multi-step method in mind or if there are particular aspects you would like us to consider, please provide more details, and we'll be happy to discuss further.

Reviewer comment (Reproducibility):
I don't see how remeshing is handled in the proposed method? Remeshing is a vital part of time-dependent PDEs, especially for fluid problems, hence I think remeshing should be a necessary part in the solution.

Author response:
$\Rightarrow$ The current implementation of the proposed GNRK method does not explicitly incorporate time-dependent remeshing.  
The method operates on non-uniform grids that are initially chosen and remain fixed over time.  
Each prediction at a time step is made independently, allowing the GNRK to handle predictions on different grids.  
Therefore, the GNRK can potentially be combined with a remeshing technique and corresponding interpolation scheme.  
Integrating time-dependent adaptive remeshing is an interesting direction for future research, as it could enhance the accuracy and applicability of the GNRK, especially in scenarios with dynamic or evolving geometries.

Reviewer comment (Experiments):
One of the weaknesses of this work is in its experimental protocol, as it has not been validated on large-scale datasets and problems with challenging geometries and high Reynolds numbers (10^6).

Author response:
Thank you for your observation regarding in the experiment protocol, particularly the absence of validation on large-scale datasets and challenges with complex geometries and high reynolds.
We acknowledge this limitation.
The focus of this study is to introduce a new approach to solving PDEs, and we chose relatively simple problems for the initial comparison.
While we recognize the importance of validating on more complex and large-scale datasets, we hope this work lays the groundwork for future researches in that direction.

Reviewer comment (Theory):
It is mentioned in the paper "It can even generate solutions with higher precision than the training data through the RK order adjustment": Can you explain more how it is possible? Is it just inference or fine-tuning at test time?

Author response:
For clarification, we rephrased the sentence as follows:
> It can even predict solutions with higher precision than the training data by modifying its' structure without retraining or fine-tuning.

During the training phase, the GNRK is trained with a specific order, and the trained $f_\theta$ captures the underlying dynamics.
However, at test time, the GNRK can adapt to different orders by changing the recurrent number in Fig 2(a), without requiring additional training.
This flexibility allows us to generate solutions with higher precision than the training data.
For instance, if the model is trained with an Runge-Kutta order of 1, it can be adjusted to order 4 at the test time, enabling more accurate predictions.
The result of this order adjustment capability is shown in Dataset IV in Section 4.1 and Section 4.2.

Reviewer comment (Experiments):
It would be very important to validate the proposed work on challenging datasets. For instance:  
A- EAGLE: Large scale learning of turbulent fluid dynamics with mesh transformers. About 1 million 2D meshes with 600 different scenes.  
B- AirfRANS: High Fidelity Computational Fluid Dynamics Dataset for Approximating Reynolds-Averaged Navier–Stokes Solutions. Complex geometries, high Reynolds numbers, and challenging to obtain accurate results at the boundary layer.  
C- PDEBench: An Extensive Benchmark for Scientific Machine Learning. For instance considering the following datasets: 3D_cfd, 2D_cfd, darcy_flow.

Author response:
We appreciate your suggestions for challenging datasets, particularly in fluid dynamics research.
As in general response, it's essential to note that the primary focus of this work is on introducing novel approach for AI PDE solvers rather than adapting existing approaches to specific PDEs.
The experiment on 2D Burgers' equation is conducted on a relatively small scale.
The training involves 20 trajectories generated from 10,000 2D grid points, and predictions were made for 50 different trajectories.
While this may not cover the complexities of the large-scale datasets you suggested, it provides a foundation for future exploration and validation on more challenging problems.
Although, your specific suggestions are valuable, and the consideration of such datasets could be a direction for further research to assess the scalability and robustness of the proposed GNRK approach.

Reviewer comment (Experiments):
There exists serious mismatches between the authors' claim and their experiments. The authors claim that GNRK can be generalized to different initial conditions, PDE coefficients, spatial discretization and temporal discretization without retraining in Section 3. However, no experiments are provided to prove the generalization without retraining. They just change the Runge-Kutta order between train and test datasets in each experiment. To demonstrate the model generality, they should train the model in Dataset I and test in Dataset 2.

Author response:
As you highlighted, our assertion is that the GNRK can be generalized to diverse initial conditions (ICs), PDE coefficients, or other variables without retraining. Table 1 displays the test error for each dataset, where training was conducted under specific conditions. For example, Dataset I comprises trajectories originating from various ICs, aiming to evaluate the model's adaptability to ICs beyond those used in training. We intend to revise the main text to enhance clarity and prevent any potential misunderstandings arising from this presentation.
The scenario you propose, where the GNRK is trained on Dataset I (different ICs, same coefficients) and subsequently tested on Dataset II (same ICs, different coefficients), poses a formidable challenge. This is because the model wasn't exposed to any information regarding the varying coefficients in Dataset II during its training phase.

Reviewer comment (Experiments):
The experimental datasets are too simple to provide convincing results. In Appendix B, the initialized state of the 2D Burgers' equation is determined by a quite simple periodic function, which only changes the phase and the center point. Reference to FNO[3], the 2D Burgers' Equation dataset can be initialized using a normal distribution.

Author response:
The chosen initial condition is expressed as $A \sin(2\pi x - \phi_x) \sin(2\pi y - \phi_y) \exp(-(x-x_0)^2 - (y-y_0)^2)$.
It's noteworthy that this formulation introduces changes not just in phase (i.e., the peak locations of the $sin$ function) but also includes asymmetry in the curvature of each peak.
Furthermore, it exhibits steeper slopes along the $x$ and $y$ axes compared to a normal distribution. 
We believe these characteristics add complexity to solving the 2D Burgers' equation.

Reviewer comment (Evaluation):
Higher order of Runge-Kutta leads to higher computational cost, and comparisons between GNRK and other baselines in running time and GPU memory are expected.

Author response:
We acknowledge that a higher order of Runge-Kutta could potentially result in increased computational cost.
To address this concern, we have included additional information in Table. 1, introducing two new rows specifically for the wall clock time and GPU memory usage required for predicting a single trajectory.