[Paper Review] Message Passing Neural PDE Solvers

J. Brandstetter, Daniel E. Worrall+1 more|arXiv (Cornell University)|Feb 7, 2022

Model Reduction and Neural Networks61 citations

TL;DR

Introduces MP-PDE, a fully neural, autoregressive PDE solver based on graph neural networks that generalizes across PDEs, domains, and boundary conditions, with stability techniques like the pushforward trick and temporal bundling.

ABSTRACT

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.

Motivation & Objective

• Motivate end-to-end neural PDE solvers that generalize across resolution, topology, geometry, boundary conditions, discretization, and dimensionality.
• Propose a neural message-passing architecture that encompasses classical solvers as special cases (FDM/FVM/WENO).
• Address autoregressive training instability with the pushforward trick and temporal bundling to improve zero-stability.
• Enable generalization across multiple PDEs within a class and support testing on unseen PDE coefficients at inference.

Proposed method

• Model the computational grid as a graph and use a graph neural network in an Encode-Process-Decode framework.
• Represent the PDE solver as a sequence of learned message-passing updates that mimic classical derivative/flux computations.
• Use an equation-embedding vector to inject PDE parameters and boundary conditions into the network.
• Apply a decoder that outputs multiple future time steps jointly, enforcing consistency for stability.
• Train with a combined loss that includes a one-step term and a stability (pushforward) term to mitigate distribution shift.
• Employ temporal bundling to predict several time steps in parallel, reducing rollout error and computation.

Experimental results

Research questions

• RQ1 Can a neural message-passing solver generalize across different PDEs within a class and across varying domain geometries and boundary conditions?
• RQ2 Do stability-oriented training strategies (pushforward and temporal bundling) improve long-rollout accuracy and mitigate autoregressive errors?
• RQ3 To what extent can the MP-PDE solver recover classical solver behavior (FDM/FVM/WENO) and generalize to irregular grids and higher dimensions (1D/2D)?
• RQ4 How does conditioning on PDE coefficients (θPDE) affect generalization when solving unseen parameter configurations?
• RQ5 What is the performance of MP-PDE on multi-physics-like problems (e.g., Burgers, heat, KdV, and Navier–Stokes-inspired setups) compared to state-of-the-art neural operators and classical solvers?

Key findings

• The MP-PDE solver generalizes across resolutions, domain topologies, geometries, boundary conditions, and dimensions (1D and 2D).
• Neural message passing can represent and reproduce classical solvers such as finite differences, finite volumes, and WENO within its architecture.
• Stability-focused training (pushforward trick) and temporal bundling significantly improve rollout stability and reduce error accumulation in autoregressive predictions.
• MP-PDE outperforms WENO5 and an autoregressive FNO-RNN on unseen PDEs within a family and benefits from PDE-parameter conditioning (θPDE).
• On irregular grids and non-periodic boundaries, MP-PDE can outperform pseudospectral methods, and can generalize across boundary conditions when θPDE features are provided.
• Preliminary 2D experiments on Navier–Stokes–inspired flows (PhiFlow data) indicate scalability to higher dimensions.