Message Passing Neural PDE Solvers

TL;DR

This work presents MP-PDE, a neural, autoregressive PDE solver based on neural message passing that unifies and generalizes across resolution, topology, geometry, boundary conditions, and discretization. It introduces stability-oriented training tricks—the pushforward method and temporal bundling—to mitigate error accumulation in rollouts, and demonstrates robustness across 1D and 2D problems, including irregular grids and multiple boundary conditions. By showing that classical solvers like FDM, FVM, and WENO can be represented as MP updates, the approach blends traditional numeric methods with modern neural computation, achieving fast, stable, and accurate results while generalizing to unseen PDEs within a family. The work also provides reproducibility details and discusses limitations and future directions, including potential probabilistic numerics integration and improved symmetry exploitation for better generalization.

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.

Figures (13)

• Figure 1: (a) Left: Neural operators perform a direct mapping from initial conditions to a solution at time $t$. Right: Autoregressive models on the other hand compute the solution at time $t$ based on the computed solution at a fixed time offset before. (b) Our autoregressive solver outputs multiple time slices on every call, to reduce number of solver calls and therefore error propagation speed.
• Figure 2: Different training strategies. Left: One-step training only predicts solutions one step into the future. Middle: Unrolled training predicts $N$ steps into the future. Right: Adversarial training predicts $N$ steps into the future, but only backprops on the last step.
• Figure 3: Schematic sketch of our MP-PDE Solver.
• Figure 4: Top: Exemplary 1D rollout of shock formation at different resolutions. The different colors represent PDE solutions at different timepoints. Both the small and the large shock are neatly captured and preserved even for low resolutions; boundary conditions are perfectly modeled. Bottom: Exemplary long 2D rollout of shock formations over 1000 timesteps. Different colors represent PDE solutions at different space-time points.
• Figure 5: Survival times at E1. Rollout for long trajectories of 8 s (left), pushforward (pf) ablation (right). The ablation compares survival times at resolutions $n_x=100$ (solid) and $n_x=50$ (dashed) against survival times using pushforward (no pf), no pushforward but putting Gaussian noise ($\sigma = 0.01$), pushforward but without cutting the gradients (pf gradients).