Source-linked AI summary
Message Passing Neural PDE Solvers
Johannes Brandstetter, Daniel Worrall, Max Welling
TL;DR
Accurate PDE simulation requires solvers that generalize across problem structures and remain stable during long autoregressive rollouts. The paper develops a fully neural message-passing solver with pushforward-based stability training, representing classical schemes and supporting varied settings. Experiments report reduced error explosion and broad flexibility, while the method requires high-quality groundtruth data and lacks typical accuracy guarantees.
Problem
Existing PDE solvers must address many structural requirements, while neural autoregressive solvers can suffer instability from distribution shift during rollouts.
Method
The paper uses a graph-based neural message-passing solver with temporal bundling and the pushforward trick to encourage zero-stability.
Results
The method reduces error explosion in autoregressive training and generalizes across spatial, temporal, geometric, boundary, dimensional, and equation-parameter settings.
Takeaways & Limitations
Neural message passing can flexibly represent classical PDE solvers while supporting generalization across varied discretizations and domain properties.
Takeaways & Limitations
Training requires high-quality groundtruth data, and the learned solver lacks the accuracy guarantees typical numerical solvers provide.
Abstract
from arXiv · showhide
The numerical solution of partial differential equations (PDEs) is difficult, having led to a century of research so far. Recently, there have been pushes to build neural--numerical hybrid solvers, which piggy-backs the modern trend towards fully end-to-end learned systems. Most works so far can only generalize over a subset of properties to which a generic solver would be faced, including: resolution, topology, geometry, boundary conditions, domain discretization regularity, dimensionality, etc. In this work, we build a solver, satisfying these properties, where all the components are based on neural message passing, replacing all heuristically designed components in the computation graph with backprop-optimized neural function approximators. We show that neural message passing solvers representationally contain some classical methods, such as finite differences, finite volumes, and WENO schemes. In order to encourage stability in training autoregressive models, we put forward a method that is based on the principle of zero-stability, posing stability as a domain adaptation problem. We validate our method on various fluid-like flow problems, demonstrating fast, stable, and accurate performance across different domain topologies, equation parameters, discretizations, etc., in 1D and 2D.
1 INTRODUCTION
The paper targets a general-purpose PDE solver that can accommodate diverse structural requirements without relying on separately handcrafted methods. It proposes an end-to-end neural message-passing solver with classical numerical schemes represented as special cases and explicit autoregressive-stability techniques.
- PDE solvers must satisfy diverse user, structural, and implementation requirements, including accuracy, efficiency, geometry, boundary conditions, dimensionality, stability, and invariants.
- The paper proposes an end-to-end fully neural PDE solver based on neural message passing.
- Neural message passing can represent finite differences, finite volumes, and WENO schemes as special cases.
- Temporal bundling and the pushforward trick are introduced to encourage zero-stability when training autoregressive models.
- The solver accepts new PDE coefficients at test time and is designed to generalize across multiple PDEs within a class.
2 BACKGROUND AND RELATED WORK
The background presents classical and neural approaches for solving temporal PDEs, emphasizing how discretization choices affect stability, conservation, accuracy, and generalization. It distinguishes autoregressive solvers, which iteratively update states, from neural operator methods, which learn mappings from initial conditions to solutions.
- Classical solvers: Temporal PDEs can be transformed into ordinary differential equations by spatially discretizing the domain and approximating spatial derivatives.
- Classical solvers: Finite differences use stencil-based derivative approximations but can be unstable unless spatial and temporal discretizations are carefully controlled.
- Classical solvers: Finite volumes estimate boundary fluxes for conservation-form equations and exactly conserve the integral of the solution over each cell.
- Classical solvers: Pseudospectral methods compute derivatives in Fourier space with exponential accuracy for smooth solutions on periodic domains and regular grids.
- Neural solvers: Neural operator methods learn mappings from initial conditions to solutions, while finite-dimensional variants remain grid-dependent and infinite-dimensional variants address that limitation.
- Neural solvers: Autoregressive methods iteratively apply a temporal update, but prior hybrid approaches preserve classical solver graphs and are reported as difficult to train.
3 METHOD
The method combines a message-passing neural architecture with training strategies designed to address distribution shift and instability in autoregressive PDE rollouts. It also predicts multiple future timesteps per call and can represent classical numerical solver operations.
- Training framework: The training framework targets instability caused by distribution shift in autoregressive solvers.One-step training can accumulate small prediction errors during longer rollouts, causing divergence from the ground truth.
- Training framework: The pushforward trick adds a stability loss by perturbing training inputs with states produced by a preceding solver step.The solver is unrolled for two steps, while errors are backpropagated only from the final step.
- Training framework: Temporal bundling predicts K future timesteps synchronously, reducing solver calls and the number of distribution shifts during a rollout.The model maps u0 to (u1, u2, ..., uK) together rather than predicting only u1.
- Architecture: The MP-PDE architecture represents the computational domain as a graph whose nodes are grid cells and whose edges define local neighborhoods.This representation supports variation in grid regularity, resolution, topology, geometry, boundary modeling, and dimensionality.
- Architecture: Message passing can represent finite-difference, finite-volume, and WENO-style derivative computations.WENO5 forms adaptively weighted averages of finite-volume estimates computed over different neighborhoods.
4 EXPERIMENTS
The experiments test MP-PDE across equation families, resolutions, grids, boundary conditions, and spatial dimensions. Results show accurate low-resolution shock modeling, generalization across equations and boundary conditions, and improved autoregressive stability from pushforward training.
- Interpolating between PDEs: MP-PDE outperforms WENO5 and FNO-RNN in accuracy, while FNO-PF performs best on E1 and MP-PDE generalizes better on E2 and E3.E1 models Burgers’ shock formation; E2 varies diffusion; E3 mixes equation parameters.
- Temporal bundling and pushforward: Autoregressive rollouts increasingly diverge after approximately 8 s, as reflected by survival-ratio measurements.The long-trajectory evaluation unrolls solvers for 1000 timesteps.
- Shock formation: MP-PDE accurately models shock waves and preserves small and large shocks even at low resolutions.Figure 4 also illustrates boundary-condition modeling during 1D rollouts.
- Temporal bundling and pushforward: Pushforward training produces much higher survival times than training without pushforward or with Gaussian perturbations.Gaussian perturbations improve stability but reduce accuracy by injecting noise into the system.
- Irregular grids and boundary conditions: MP-PDE remains accurate at low resolutions where pseudospectral solvers break and generalizes across Dirichlet, Neumann, and mixed boundary conditions.Boundary-condition generalization is stronger when boundary information is supplied through θPDE features.
- 2D experiments: The experiments extend MP-PDE to two-dimensional smoke-flow simulations on a 32 × 32 grid.The 2D data come from the PHIFLOW2 fluid-simulation toolkit.
5 CONCLUSION
The paper introduces a fully neural MP-PDE solver that representationally contains several classical numerical methods and uses pushforward training with temporal bundling to address autoregressive error explosion. It demonstrates flexibility across many discretization, geometry, boundary, dimensionality, and solution-space settings, while noting data and accuracy limitations.
- MP-PDE is a fully neural solver that representationally contains finite difference, finite volume, and WENO schemes.
- Temporal bundling and the pushforward trick reduce error explosion in autoregressive training, including for a state-of-the-art neural operator.
- MP-PDE solvers generalize across spatial resolution, timescale, sampling regularity, topology, geometry, boundary conditions, dimensionality, and solution-space smoothness.
- Training requires high-quality ground-truth data, whose generation was the project’s toughest part.
- The model lacks the accuracy guarantees typically provided by classical numerical solvers.
6 REPRODUCIBILITY STATEMENT
The authors describe dataset verification, implementation documentation, validation of the proposed training methods, and public code release. Their experiments compare WENO outputs with analytical solutions for selected Burgers-equation cases.
- Datasets were generated by the authors and extensively cross-checked for correctness.
- For experiments E1, E2, and E3, the implemented WENO scheme was compared against two analytical Burgers-equation solutions from the literature.
- The paper documents its architecture and implementation details, including temporal bundling and the pushforward method.
- Temporal bundling and pushforward were validated on both MP-PDE and the Fourier Neural Operator method.
- The authors provide their codebase publicly for reproducibility.
7 ETHICAL STATEMENT
The supplied material frames MP-PDE solvers as potentially useful for computationally expensive scientific simulations, with possible implications for computational shortcuts and carbon-footprint reduction. It also describes classical reconstruction machinery underlying numerical PDE solvers.
- MP-PDE solvers may help create computational shortcuts for expensive simulations in areas such as weather forecasting, astronomy, and molecular modeling.
- Reducing computational cost is connected in the paper to potentially reducing the carbon footprint of simulations.
- ENO and WENO use multiple stencils to estimate derivatives near discontinuities while favoring substencils that avoid the discontinuity.
- WENO reconstruction combines overlapping substencils using nonlinear weights and smoothness indicators.
- Finite-difference spatial discretization approximates derivatives with conservative differences and numerical fluxes.
C.3 COMPARING WENO SCHEME TO ANALYTICAL SOLUTIONS
Two analytical Burgers-equation cases compare WENO with analytical solutions, with finite differences additionally evaluated in the first case. WENO matches the analytical solutions, while FDM becomes inaccurate or fails as diffusion decreases.
- The first case is defined by analytical Burgers-equation boundary and initial conditions used for comparison with numerical schemes.
- At ν = 0.005 in the first analytical case, WENO models the analytical solution perfectly, whereas fourth-order FDM fails to capture the shock accurately.
- For lower ν values in the first case, the difference between WENO and FDM becomes stronger.
- At ν = 0.002 in the second analytical case, WENO again models the analytical solution perfectly.
- In the second analytical case, the FDM scheme fails completely.
D PSEUDOSPECTRAL METHODS FOR WAVE PROPAGATION ON IRREGULAR
The wave-propagation experiments use pseudospectral ground truth on irregular grids and evaluate Dirichlet and Neumann boundary conditions. Boundary interactions are modeled through reflected wave packets with condition-dependent sign behavior.
- Dirichlet boundaries impose u = 0, whereas Neumann boundaries impose ∂xu = 0.
- Wave packets bounce between boundaries, retaining their sign for Neumann conditions and changing sign for Dirichlet conditions.
- Wave solutions are obtained on irregular grids using pseudospectral solvers for Dirichlet and Neumann boundary conditions.
E EXPLICIT RUNGE-KUTTA METHODS
Explicit Runge-Kutta methods advance temporal PDE discretizations through multiple stages using specified coefficients. Their coefficients are organized in Butcher tableaux.
- A Runge-Kutta method computes stage derivatives from the current time, timestep, solution, and previously computed stages.
- Each method requires a stage count s and coefficients aij, bi, and ci.
- The coefficients are conventionally arranged in a Butcher tableau.
F EXPERIMENTS
The experiments assess MP-PDE solvers across one- and two-dimensional equations, resolutions, parameters, boundary conditions, and implementation settings. Results show accurate shock, dispersive, wave, and smoke-inflow rollouts, while the appendix details the architecture and training setup.
- Experiments and setup: The studied equations include Heat, Burgers, KdV, and Kuramoto-Shivashinsky equations, with their flux terms summarized in Table 3.
- Implementation details: The architectures combine node-wise encoding, six message-passing layers, and a convolutional decoder, with zero-stability training using a maximum two-step unroll.
- Implementation details: MP-PDE experiments use 6 message-passing neighbors in E1–E3 and 8 neighbors in 2D experiments, while irregular-grid wave experiments use k-NN neighborhoods.
- E1, E2, E3: MP-PDE solvers reconstruct dispersive wiggles at very low resolution that numerical solvers cannot capture, while lower diffusion produces faster Burgers shock formation.
- WE1, WE2, WE3: MP-PDE solvers accurately model wave propagation on irregular grids with Dirichlet and Neumann boundaries and capture smoke inflow over 100 timesteps.
G ARCHITECTURE ABLATION AND COMPARISON TO CNNS
The ablation and CNN comparisons examine graph-based architectural flexibility, decoder choices, PDE and boundary-condition features, and convolutional baselines. Results show benefits from PDE features and shared-weight convolutional decoding, while the 1D-CNN remains a strong but less general baseline.
- Architecture: The graph architecture is designed to generalize across sampling regularity, resolution, topology, geometry, boundary modeling, and dimensionality, while containing FDM, FVM, and WENO function classes.
- Ablation design: The ablations test GNN versus CNN representational power, shared-weight 1D convolution versus an MLP decoder, and PDE-coefficient features.
- PDE features: Injecting θPDE features has a stronger effect when more equation-specific parameters are available and strengthens generalization across boundary conditions.
- Decoder ablation: Replacing the shared-weight 1D-CNN decoder with a standard MLP slightly degrades performance, likely because shared weights improve temporal modeling.
- CNN comparison: The 1D-CNN performs much better than the 2D-CNN, but its generalization across different PDEs becomes harder in experiment E3.