RandNet-Parareal: a time-parallel PDE solver using Random Neural Networks

Parallel-in-time (PinT) techniques have been proposed to solve systems of time-dependent differential equations by parallelizing the temporal domain. Among them, Parareal computes the solution sequentially using an inaccurate (fast) solver, and then "corrects" it using an accurate (slow) integrator that runs in parallel across temporal subintervals. This work introduces RandNet-Parareal, a novel method to learn the discrepancy between the coarse and fine solutions using random neural networks (RandNets). RandNet-Parareal achieves speed gains up to x125 and x22 compared to the fine solver run serially and Parareal, respectively. Beyond theoretical guarantees of RandNets as universal approximators, these models are quick to train, allowing the PinT solution of partial differential equations on a spatial mesh of up to $10^5$ points with minimal overhead, dramatically increasing the scalability of existing PinT approaches. RandNet-Parareal's numerical performance is illustrated on systems of real-world significance, such as the viscous Burgers' equation, the Diffusion-Reaction equation, the two- and three-dimensional Brusselator, and the shallow water equation.

Paper

References (82)

Scroll for more · 38 remaining

Similar papers

Peer review

Reviewer awZW6/10 · confidence 2/52024-07-07

Summary

This paper proposes a new method for sequentially predicting and correcting numerical simulations by introducing random neural networks. The RandNets are single-layer feed-forward neural networks and only the output layer is for training. The numerical experiments have shown the improved training efficiency and scalability of the proposed method compared to the baseline models.

Strengths

- This paper is well-written and well-organized. The theoretical guarantee of RandNets-Parareal is provided. - Different types of PDEs have been used to evaluate the model performance.

Weaknesses

- The studied problem is domain-specific, which might not be of general interest to the scientific machine-learning community. - In Sec 5 Numerical Experiments, this paper provides the comparison of speed-ups/runtimes. It would be better to also show the varying solution accuracy with different setups. - Sections 2 and 3 might be shortened and some of these parts can be moved to the Appendix. The Robustness study (Appendix C) can be moved to the main text since it highlights the benefits of the proposed method.

Questions

In lines 7-8, is it a character "x" or a symbol "\times" for `x125`? Probably a symbol is preferred.

Rating

6

Confidence

2

Soundness

3

Presentation

3

Contribution

2

Limitations

N/A.

Authorsrebuttal2024-08-11

Follow-up to Reviewer awZW's comment

Thanks for your comment on our rebuttal and W2. We would appreciate if you could read what we wrote as a comment to Reviewer FREq about the impact of our paper on AI, which we report below for your convenience. We hope tis may reduce some of your reservations on the significance of our work. Properly communicating the relevance and impact of our method is very important for us. Our work exemplifies an instance when ML methods allow the advancement of various fields of science and engineering where solving (O/P)DEs is needed, see also our response to Ref FREq to his further question. We mentioned these applications in our initial rebuttal as our submission belongs to the Primary Area of *ML for Physical Sciences*. However, AI is also one of the fields that could greatly benefit from progresses regarding efficient ODE solvers, and thus our method. This is why we gladly provide examples of how our approach can also assist ML and AI, adding these points in the camera-ready version of the paper. ODEs are a crucial building block of some relevant techniques, such as *Diffusion models (DMs)*, *Neural ODEs*, *Optimal control and reinforcement learning*, and *Optimization for ML models*, to name a few. Below, we provide an example of how our work could lead to substantial advancements for DMs. - _Solvers for DMs_: In the context of DMs, the continuous time reversal process can be described by a probability flow ODE, defined by the score function, usually approximated with deep (convolutional) NNs (note that the ODE formulation offers significant advantages over SDEs in high dimensions). Faster ODE solvers can improve sampling speed, yielding faster image synthesis. The main existing directions in the literature involve using classical sequential solvers [30], developing faster and dedicated ones such as DDIM [24], DDPM [29], DPM-Solver [25], Heun [26], and parallelization of the autoregressive sampling process of DMs (by using Picard-Lindelöf iterations for ODE/SDE) [27,28]. The latter emerged mainly due to the need for other solutions to avoid the bottlenecks of sequential solvers. To the best of our knowledge, modulo this particular parallelized sampling, PinT schemes have not yet been used in this context. **This is the gap this paper could fill, as the existing successful sequential dedicated solvers could be embedded into our proposed RandNet-Parareal**. We expect this to be straightforward, since the goal is to collect samples via solving the corresponding (diffusion) ODEs at $[0,T]$, where RandNet-Parareal immediately finds its purpose. We refer to our latest comment to Reviewer FREq for an additional discussion on the GPU implementation which could further improve the efficiency of our method. We hope these clarifications provide you with evidence of the expected impact/implications of our method in the context of generative AI models and beyond. **References used above:** [24] J. Song, C. Meng, S. Ermon. Denoising diffusion implicit models. 2020. [25] C. Lu, Y. Zhou, F. Bao, J. Chen, C. Li, J. Zhu. DPM-solver: A fast ODE solver for diffusion probabilistic model sampling in around 10 steps. 2022. [26] T. Karras, M. Aittala, T. Aila, S. Laine. Elucidating the design space of diffusion-based generative models. 2022. [27] A. Shih, S. Belkhale, S. Ermon, D. Sadigh, N. Anari. Parallel sampling of diffusion models. 2023. [28] Z. Tang, J. Tang, H. Luo, F. Wang, T.-H. Chang. Accelerating parallel sampling of diffusion models. 2024. [29] J. Ho, A. Jain, P. Abbeel. Denoising diffusion probabilistic models. 2020. [30] Y. Song, J. Sohl-Dickstein, D.P. Kingma, A. Kumar, S. Ermon, B. Poole. Score-based generative modeling through stochastic differential equations. 2021.

Reviewer nFhC6/10 · confidence 3/52024-07-08

Summary

This paper proposes a method to accelerate the simulation of partial differential equations (PDEs) by converting them into systems of ordinary differential equations (ODEs). It then utilizes a framework that merges the random neural network and the parareal approach, termed RandNet-Parareal. For validation, three complex systems of PDEs are considered: the Burgers equation, the diffusion-reaction equation, and the shallow water equation. Results show gains achieved in terms of computational cost compared to other methods.

Strengths

1. The proposed method accelerates neural PDE simulations by combining random neural networks with parareal concepts. 2. The considered PDEs are prototypical and complicated and exhibit the efficacy of the proposed approach.

Weaknesses

1. The authors have not mentioned the trade-off between accuracy and computational cost, which is essential for high-fidelity simulations. 2. State-of-the-art comparisons are lacking, restricting the evaluation of the proposed method with random neural network-based PDE solvers. 3. The method is based on conventional numerical solvers, and hence, extending its applicability to complex geometries in higher dimensions may suffer from the curse of dimensionality. 4. It is often the case that random neural networks do not scale well to deep networks, and not much gain is achieved when using a deep network. The paper does not discuss extending the method to deep neural networks and what gains could be achieved.

Questions

1. The proposed method's performance and compared methods are not quantified through a metric. It would be interesting to see the trade-off between accuracy and computational cost. 2. The method is not compared with other neural PDE solvers, specifically random neural network-based PDE solvers aiming to accelerate PDE simulations. Assessing the gain achieved compared to the methods proposed in the literature would be interesting. 3. As the method depends on traditional mesh-based numerical solvers, can the authors provide insights on how the method would perform in complicated geometries, particularly in higher dimensions? 4. Can the authors comment on the performance of their method when using a deep network instead of a shallow network?

Rating

6

Confidence

3

Soundness

2

Presentation

3

Contribution

3

Limitations

The authors have mentioned the limitations of the work related to coarse numerical solvers and stiff systems. There are no social or ethical issues concerning the paper.

Reviewer Eu5i7/10 · confidence 4/52024-07-08

Summary

The authors introduce a numerical algorithm that computes the solution to a large system of ordinary differential equations (ODE) "parallel in time". The main idea of the solver is based on the existing "Parareal", which introduces parallelism by running a sequential, fast, and inaccurate ODE solver and then corrects it in parallel with a more accurate correction scheme. The novelty in this manuscript is that this correction scheme is done using random neural networks, i.e., networks where the hidden weights and biases are chosen at random and then fixed. The benefit of this is that training time for the network is reduced significantly, as only the last layer must be approximated, which is possible using a linear solve. The authors demonstrate this efficiency and the parallel in time property on several nonlinear, time-dependent partial differential equations, discretized to obtain the required large ODE system.

Strengths

Using random neural networks in a parallel-in-time setting is (to my knowledge) a new approach, and helps to mitigate issues with earlier versions of algorithms (e.g. the GP version of ParaReal, as presented). To that end, the work is a combination of well-known techniques to solve a problem in a new way. The authors discuss and demonstrate how their approach differs from earlier work, mostly from ParaReal with GP, fine-scale solvers, and a classical ParaReal approach. The paper is very well written, with clear descriptions of the algorithm and numerical results, and even includes scaling results for multi-core systems. The authors also include a robustness study of the approach in the appendix.

Weaknesses

1) There is no theoretical analysis beyond restating existing results from the literature. For example, there is no complexity analysis for the random network setting, which should have been relatively straightforward (given that it involves only a single linear solve, for which complexity results are available). 2) The numerical scenarios chosen in the paper are not particularly challenging on their own. It seems to me that the parallel in time setting itself is what makes them challenging, not the high number of degrees of freedom (e.g., 10^5 is not very large), nor the particularly high spatial dimension (one and two), or the complexity of the PDE (mildly nonlinear). Of course, for parallel-in-time algorithms these are challenging problems nonetheless, but the examples do not demonstrate why parallel in time in general makes more sense than just solving the system with a better algorithm (e.g. higher order, other basis functions, etc.). For example: it is not reasonable that solving Burgers equation on a 1D line takes 13 hours on the "fine scale" (table 1). This looks like an unnecessary difficulty for a problem that should be simple to solve, especially with only 128 degrees of freedom (variable d). 3) It seems to me that the main novelty of this paper is to replace the (numerically suboptimal) Gaussian process with a neural network where the internal weights are randomly chosen. It is not clear to me why this particular change is any better than improving the numerics for the GP setting, or using any other possible solution (e.g., just a linear mapping as approximation, or running a fine-scale solver on multiple scales, or using polynomial regression, ...). It is probably not possible to perform these additional experiments during the rebuttal time, but the manuscript does not even state any other possibilities that were ruled out. 4) Similarly, the state of the art focuses very much on ParaReal, and not on any other parallel in time solver. While this may be acceptable if ParaReal was the most advanced PinT solver available, it is not clear (because it is not stated) if that is the case in general (or at least for the chosen examples). A few minor issues: 1) I would avoid using "nns" for nearest neighbors (l.35), because it is confusing with "NNs" (neural networks). 2) l.87: it is strange to say "24 core processors, 48 cores". It is clear that 24*2=48; I hope that is what the authors mean.

Questions

1) The authors stress that training a GP is fundamentally more expensive than training the random network. However, I fail to see this: if we do not compute a full kernel matrix of N by N (N being the number of data points), but just randomly choose M center points and then solve a linear system of N by M, the complexity is exactly the same as in the random network case. The only difference is that instead of M neurons we have M kernel functions (k(x_i, *)), essentially forming a "radial basis neural network". This has been discussed at length in the GP literature, and many methods for this exist (e.g. Nystroem kernel approximations, inducing points methods, etc.). Why do the authors only compare to the full kernel case, which obviously is subpar in performance? 2) Why do the authors compare to clearly sub-optimal numerical solvers in the PDE setting? I may have missed the reasoning why even the simplest PDE settings take so much time to solve. 3) Is Parareal the only parallel-in-time solver available? The paper focuses mostly on different Parareal settings and approaches, but not on different parallel-in-time solution methods. 4) How does the high ambient-space dimension d of the state U affect the nearest neighbor search?

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors are missing a discussion that "nearest neighbor" algorithms do not work well in high-dimensional settings (as given here, with $d \gg 1000$ in many cases). They also do not discuss the suboptimal scaling of the linear solution for the random neural networks (caused by the linear solve necessary for the outer weights).

Reviewer FREq5/10 · confidence 2/52024-07-12

Summary

The paper introduces their method RandNet-Parareal, which is a method to solve differential equations. Their method can be categorized as a Parallel-in-time technique that aims at parallelizing solvers in the temporal domain. They extend the Parareal algorithm by using RandNets which learn the difference between a coarse and a fine solver. They show experimentally a speed-up of their method compared to previous methods and provide theoretical guarantees for their method.

Strengths

The paper evaluates their method on three different PDEs. The authors show that their method yields a significant speed-up. The authors provide theoretical guarantees of their method. The authors claim that the algorithm can be used as a convenient out-of-the-box algorithm.

Weaknesses

Paper Structure: The authors provide experimental details within the introduction (lines 85-90) that should belong to the experimental section. The paper is not structured as most papers in the conference. I think Section 2 should be put into a Section called “Problem Description”. Section 3 should be included in Section 2 or into a “Related Work” section. The paper misses a related work section. To me, it is questionable if this paper is of interest to the audience of the conference. The paper replaces a subpart of a numerical algorithm with a two-layer Neural Network where only the last layer is trained. Therefore, while the contribution is sound the impact of the work to other fields of machine learning is very limited.

Questions

1. Can you put your paper into a broader perspective? This paper only compares methods that are based on the Parareal algorithm. Are there also alternatives and how do they compare to your method? 2. What is the “Fine” algorithm in Table 1 and in Fig. 1? I could not find its definition. 3. How can the Runtime in Figure 1 be negative? Is log(Runtime) shown in this plot? 4. Is there a way to measure the solution quality of the solver? Is the solution quality fully defined by the converged initial condition $\epsilon$ and therefore the same for all methods? If the solution quality can be quantified, how does the solution quality of your method compare to other methods? 5. Are the subscripts $k$ and $i$ dropped in Equation 6? If so this is a bit confusing since these are used in the previous Section and then dropped without a notice.

Rating

5

Confidence

2

Soundness

3

Presentation

2

Contribution

2

Limitations

The authors addressed some limitations of their work.

Authorsrebuttal2024-08-12

Follow-up answer to Reviewer FREq's comments and raised score

We are glad that our rebuttal comments were helpful in addressing your concerns. We are also grateful for reflecting this in your score, and for your further suggestions. We notice that there are still some lingering misconceptions about our work, which we believe are important to point out. Our work does not "present an improvement to a numerical ODE solver". Instead, it improves existing _Parallel-in-Time_ (PinT) methods for solving ODEs and PDEs, as written in the paper abstract, summarized by the other reviewers, and discussed at length during the rebuttal. We do, ultimately, solve an (O/P)DE with a numerical scheme, but our proposed scheme (1) solves DEs by parallelizing in time, belonging thus to a different class than that of the "simple" numerical ODE solvers; (2) models the correction term using particularly suited **random neural networks**, allowing us to achieve significant speed-up compared to existing PinT techniques; (3) can incorporate existing (O/P)DE solvers in its architecture; (4) is not meant to improve a specific ODE solver, but instead to make PinT techniques more effective and, hence, more appealing to practitioners in the field. Moreover, our new approach (even if considered to be an improvement to an ODE solver, an oversimplification which we disagree with) is far from "solely relevant to obtaining training data for PDE machine learning methods". Numerical methods for solving DEs have profoundly impacted scientific progress since the last century, extending far beyond the generation of training data generation. In fact, once again, **our method can incorporate existing PDE machine learning solvers instead of training data for them.** Independently on that, landmark achievements such as weather forecasting, space exploration, and advancements in nuclear energy are just a few examples of the broad and transformative influence of numerical DE solvers. Commenting on the relevance of designing efficient PinT solvers beyond AI, the following non-exhaustive list contains some other areas of natural sciences, social sciences, and engineering where the availability of efficient (O/P)DE solvers could be a game-changer leading to extremely important advances: - *Atmospheric Dynamics:* Accurate solutions can mitigate climate change, enhance disaster preparedness, and optimize agriculture. - *Biomedical Processes:* Real-time simulations of biological processes, such as tumor growth and blood flow, can aid in personalized treatment planning and targeted interventions. - *Earthquake Engineering:* Refined modeling of seismic wave propagation and structural responses enhances earthquake predictions and leads to more resilient building designs. - *Financial Modeling:* Improved pricing accuracy and real-time risk assessment can stabilize financial markets and reduce systemic risk. We believe our work is another step toward providing faster numerical solutions for DEs in any field where this may be required, whether in the natural sciences, engineering, machine learning, or artificial intelligence. We thank you for raising the score and hope that our further comments help convey the main scope of our contributions.

Reviewer FREq2024-08-08

Answer to Rebuttal

**AQ1:** I appreciate the authors clarification on the background and encourage the authors to add a corresponding section into the paper. **AQ2:** I think this should be defined more clearly in the experiment section. **AQ3:** The y-axis label is incorrect, as it is denoted in hours but should be log10(hours). It would be better to scale the y-axis logarithmically while keeping the y-axis ticks in hours, as this is easier to interpret than log10(hours). **AQ4:** The authors should not drop indices without mentioning it, and they should add a footnote for clarification. **AW2:** While this may not be a NeurIPS requirement, the reviewers should still consider the potential impact of the method in AI when assigning ratings. The reviewers noted the method's potential applicability to ODE solvers used in machine learning, such as in diffusion models. The paper would benefit from a section discussing this potential application in more detail. **Further Question:** Are PinT schemes used to solve the reverse ODE in Diffusion Models? If not, which ODE solvers are used? What would be necessary to apply the authors' method in this setting? Likely it would need to be implemented on a GPU? Demonstrating the benefit of the authors' method in the context of diffusion models would significantly increase the relevance of this paper to the NeurIPS community.

Authorsrebuttal2024-08-09

Follow-up answer to Reviewer FREq's comments

**AQ1-3**: Thank you for your feedback. We will carefully take it into account when preparing the final version. We welcome any comment on our response to **Q4** on the solution quality of the solver. **AQ5: Dropping indices without prior mention.** Please note that we specify that on Line 195 in [paper], which reads "Prior to that [i.e., defining RandNet for Parareal], we define how RandNets work in a general setting with input $\boldsymbol{U}$". Indices are introduced later, upon combining RandNets with Parareal, in line 237. Should you find this still unclear to the reader, we could change it to "Prior to that, we define how RandNets work in a general setting with input $\boldsymbol{U}$, *before going back to the input of interest $\boldsymbol{U}_i^k$ within the Parareal framework*". **AW2 and Further Q: On the potential impact of RandNet-Parareal on ML/AI, and, specifically, on Diffusion Models (DMs).** Thank you for engaging in further conversation about this matter. Properly communicating the relevance and impact of our method is very important for us. Our work exemplifies an instance when ML methods allow the advancement of various fields of science and engineering where solving (O/P)DEs is needed, see also our response to W2 above, and W1 of Ref awZW. We mentioned these applications as our submission belongs to the Primary Area of *ML for Physical Sciences*. However, AI is also one of the fields that could greatly benefit from progresses regarding efficient ODE solvers, and thus our method. This is why we gladly provide examples of how our approach can also assist ML and AI, adding these points in the camera-ready version of the paper. ODEs are a crucial building block of some relevant techniques, such as *DMs*, *Neural ODEs*, *Optimal control and reinforcement learning*, and *Optimization for ML models*, to name a few. - _Solvers for DMs_: In the context of DMs, mentioned in your answer, the continuous time reversal process can be described by a probability flow ODE, defined by the score function, usually approximated with deep (convolutional) NNs (note that the ODE formulation offers significant advantages over SDEs in high dimensions). Faster ODE solvers can improve sampling speed, yielding faster image synthesis. The main existing directions in the literature involve using classical sequential solvers [30], developing faster and dedicated ones such as DDIM [24], DDPM [29], DPM-Solver [25], Heun [26], and parallelization of the autoregressive sampling process of DMs (by using Picard-Lindelöf iterations for ODE/SDE) [27,28]. The latter emerged mainly due to the need for other solutions to avoid the bottlenecks of sequential solvers. To the best of our knowledge, modulo this particular parallelized sampling, PinT schemes have not yet been used in this context. **This is the gap this paper could fill, as the existing successful sequential dedicated solvers could be embedded into our proposed RandNet-Parareal**. We expect this to be straightforward, since the goal is to collect samples via solving the corresponding (diffusion) ODEs at $[0,T]$, where RandNet-Parareal immediately finds its purpose. - _GPU Implementation_: There are two increasingly sophisticated approaches for implementing RandNets-Parareal on GPUs: (i) parallel implementation of the solvers (model parallelism); (ii) parallel implementation of Parareal (in-time parallelism). For (i), both the solver and the RandNet computation can be carried out on the GPU, see e.g.[31]. In-time parallelism would then be implemented across GPUs. Otherwise, under (ii), both model and time parallelism can be implemented concurrently on the same hardware [32]. The latter guarantees a better use of resources at the cost of increased implementation complexity. We hope these clarifications provide you with evidence of the expected impact/implications of our method in the context of generative AI models and beyond. **References used above:** [24] J. Song, C. Meng, S. Ermon. Denoising diffusion implicit models. 2020. [25] C. Lu, Y. Zhou, F. Bao, J. Chen, C. Li, J. Zhu. DPM-solver: A fast ODE solver for diffusion probabilistic model sampling in around 10 steps. 2022. [26] T. Karras, M. Aittala, T. Aila, S. Laine. Elucidating the design space of diffusion-based generative models. 2022. [27] A. Shih, S. Belkhale, S. Ermon, D. Sadigh, N. Anari. Parallel sampling of diffusion models. 2023. [28] Z. Tang, J. Tang, H. Luo, F. Wang, T.-H. Chang. Accelerating parallel sampling of diffusion models. 2024. [29] J. Ho, A. Jain, P. Abbeel. Denoising diffusion probabilistic models. 2020. [30] Y. Song, J. Sohl-Dickstein, D.P. Kingma, A. Kumar, S. Ermon, B. Poole. Score-based generative modeling through stochastic differential equations. 2021. [31] A.N.Budko, M. Möller, C. W. J. Lemmens. Time integration parallel in time. 2020. [32] A.Q. Ibrahim, S. Götschel, D. Ruprecht. Parareal with a physics-informed neural network as coarse propagator.

Reviewer FREq2024-08-12

Raised Score

I appreciate the authors' efforts in addressing the concerns raised in the rebuttal. Based on this, I am increasing my score from 4 to 5. The paper presents an improvement to a numerical ODE solver, which is solely relevant to obtaining training data for PDE machine learning methods. To enhance the paper's relevance to the AI community, I suggest two improvements: 1. Provide a GPU implementation of the proposed method. 2. Include a comparison between this method and other ODE solvers commonly used in simulating the reverse process of diffusion models. These additions would help demonstrate the practical applications and potential advantages of the method in AI-related contexts.

Reviewer Eu5i2024-08-10

I highly appreciate the additional results, both experimental and theoretical. I very much like the general idea of the paper but was unsure about the soundness in these two aspects, both of which were addressed. I will raise my score to 7. * "Should we add [the new experiments] to Suppl. Material?" Yes, please.

Authorsrebuttal2024-08-11

Follow-up to Reviewer Eu5i's comment

We are glad you appreciated our additional theoretical and numerical results. We thank you again for your constructive criticisms, we believe the revised paper will greatly benefit from them. We are also extremely grateful that you have upgraded your score notably.

Reviewer awZW2024-08-10

Response for rebuttal

Thanks for your rebuttal, especially the response to **W2**. I still have some reservations about the significance of the work. I will be maintaining my score.

Area Chair ACQd2024-08-13

Last day for discussion

Reviewer nFhC, today is the last day for discussion. I hope you can take a moment to respond to the authors' rebuttal.

Reviewer nFhC2024-08-13

Raising my score

Thank you so much for providing detailed answers. Considering all the answers, I have raised my score from 5 to 6. It would be nice to include trade off results in the revised paper.

Program Chairsdecision2024-09-25

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC