Message Passing Neural PDE Solvers
MP-PDE unifies classical local solvers with message passing and improves autoregressive stability via pushforward training.
Key Findings
Methodology
Message Passing Neural PDE Solver (MP-PDE) represents grid cells as graph nodes and advances solutions autoregressively. An Encoder embeds recent solution histories, coordinates, time, and PDE attributes; an MLP-based Processor performs local message passing; a temporal CNN Decoder predicts K future increments. Temporal Bundling reduces solver calls, while the Pushforward Trick trains on model-induced inputs, treating rollout instability as a distribution-adaptation problem related to zero-stability.
Key Results
- Across 2,096 one-dimensional trajectories, MP-PDE was evaluated on nondiffusive Burgers, variable-diffusion Burgers, and mixed PDE parameters. For E3 at (nt,nx)=(250,40), accumulated MSE was 3.70, compared with 20.90 for FNO-RNN, 5.98 for FNO-PF, and 15.94 for WENO5.
- For the shock-forming E1 task, MP-PDE errors at nx=100, 50, and 40 were 1.55, 1.67, and 1.47. A 250-step rollout took only 0.09, 0.08, and 0.08 seconds, respectively, while shocks remained recognizable at low resolution.
- Ablations showed that Pushforward training substantially improved survival ratios. Gaussian perturbations helped stability less effectively and reduced accuracy. Temporal Bundling also improved FNO, producing 0.54 error for FNO-PF on E1.
Significance
The paper addresses a long-standing tension between the flexibility of learned models and the reliability of numerical solvers. Classical schemes are accurate but usually tailored to a PDE, mesh, geometry, and boundary condition; neural operators can be global and resolution-aware but are often trained for one equation and may extrapolate poorly in time. MP-PDE provides a graph-based interface for irregular domains and conditions on equation parameters. Its broader significance is conceptual: learnable local operators can retain the structural intuition of numerical analysis while adapting from data. It therefore offers a route toward reusable scientific simulators, although it should not yet be interpreted as a universal replacement for validated numerical methods.
Technical Contribution
The first contribution is an explicit representational bridge: weighted message aggregation can express finite differences, finite volumes, and WENO5-like local constructions, requiring roughly one, two, and three message-passing layers in the paper's interpretation. Second, the Pushforward Trick generates perturbations from the solver's own predecessor predictions and backpropagates only through the final unrolled step, approximating direct control of perturbation amplification. Third, Temporal Bundling predicts several future slices per call. The residual Decoder, uk+ℓ=uk+(tk+ℓ−tk)dℓ, preserves a consistency-motivated update form while allowing nonlinear, input-adaptive time integration.
Novelty
The central novelty is not merely applying a GNN to a PDE; earlier work had already used learned graph dynamics. Rather, MP-PDE combines a fully neural local solver, explicit coverage of classical discretization families, equation conditioning, and a rollout-stability objective based on the model's pushforward distribution. Compared with hybrid methods that retain a handcrafted solver and learn only coefficients, MP-PDE replaces the computational components themselves. Compared with neural operators, it performs local autoregressive evolution and is designed for cross-equation generalization.
Limitations
- The evidence is concentrated on one- and two-dimensional fluid-like conservation problems. Generalization to stiff systems, nonconservative equations, high-dimensional turbulence, long-range interactions, and strongly coupled multiphysics remains unestablished.
- Empirical survival curves support the Pushforward method, but no rigorous nonlinear convergence or long-horizon error bound is provided. Training cost, sensitivity to K and message depth, and performance on real measured data are also not systematically quantified.
Future Work
Important directions include three-dimensional complex geometries, unstructured meshes, stiff and multiphysics equations, and hard enforcement of conservation, symmetry, monotonicity, and physical invariants. The community should seek formal stability and generalization guarantees for pushforward training, adaptive neighborhoods and time steps, and scalable sparse message passing. Comparisons with modern neural operators should use matched compute and broad out-of-distribution tests. Real weather, engineering, and materials datasets would test whether gains persist beyond numerically generated trajectories.
AI Executive Summary
Partial differential equations underpin weather prediction, astrophysics, and engineering design, yet their numerical solution usually requires a solver tailored to the equation, mesh, geometry, and boundary conditions. Finite differences, finite volumes, and WENO5 are dependable but specialized. Neural operators such as the Fourier Neural Operator offer broad function mappings, but are commonly trained for one equation and can fail during long autoregressive rollouts.
Brandstetter, Worrall, and Welling propose MP-PDE, which treats a discretized domain as a graph and learns local updates through message passing. An Encoder combines recent solution values, coordinates, time, and equation parameters; an MLP Processor aggregates neighboring information; and a temporal CNN Decoder predicts several future increments at once. The Pushforward Trick exposes training to states generated by the model itself, while Temporal Bundling reduces the number of distribution-shifting rollout calls.
Experiments used 2,096 trajectories from a parameterized one-dimensional PDE family spanning Burgers, heat-like, and KdV-like regimes. On E3 with nx=40, MP-PDE reached accumulated MSE 3.70, versus 5.98 for FNO-PF, 20.90 for FNO-RNN, and 15.94 for WENO5. On shock formation, errors were about 1.5 across nx=100, 50, and 40, with 250-step runtimes near 0.08–0.09 seconds. The method therefore combines classical local structure with learned adaptability. Its unresolved challenges are rigorous stability theory, three-dimensional complex physics, and validation on real-world data.
Deep Analysis
Background
PDE computation has evolved around finite difference, finite volume, pseudospectral, and WENO schemes. FDM is simple but stability-sensitive; FVM conserves integrated quantities but requires conservation form; pseudospectral methods are highly accurate for smooth periodic solutions but fail around nonsmooth structure. Neural operators, including the Fourier Neural Operator of Li et al. (2020a), reduce grid dependence, yet models are often equation-specific. The paper seeks an autoregressive solver that can adapt across resolution, geometry, topology, boundary conditions, and coefficients.
Core Problem
The central difficulty is autoregressive distribution shift. One-step training sees ground-truth uk, whereas inference receives its own imperfect prediction, so small errors can amplify over many steps. Fixed-grid architectures also struggle with irregular sampling, changing geometry, and boundary representation. The desired solver must remain accurate and stable over long rollouts while retaining fast inference and generalizing across a PDE family rather than memorizing one equation.
Innovation
MP-PDE makes four linked innovations. First, it replaces handcrafted local numerical components with learned graph message passing. Second, Pushforward training feeds back states produced by the model and computes loss only on the final unrolled step, targeting the relevant test distribution. Third, Temporal Bundling predicts several future times simultaneously, reducing error-propagation opportunities. Fourth, θPDE conditions the same network on α, β, γ and boundary attributes, enabling interpolation across equations rather than separate models per PDE.
Methodology
- �� Graph representation: cells become nodes and local neighborhoods become edges, supporting irregular meshes, geometry, topology, and dimensions.
- �� Encoding: f⁰ᵢ=ε([uᵏ⁻ᴷ:ᵏᵢ,xᵢ,tᵏ,θPDE]) embeds solution history, position, time, and equation metadata.
- �� Processing: mᵐᵢⱼ=φ(fᵐᵢ,fᵐⱼ,uᵢ−uⱼ,xᵢ−xⱼ,θPDE), followed by fᵐ⁺¹ᵢ=ψ(fᵐᵢ,Σmᵐᵢⱼ,θPDE). MLPs learn messages and updates.
- �� Decoding: a shallow temporal CNN outputs dᵢ¹:ᴷ, and the residual update is uk+ℓᵢ=ukᵢ+(tk+ℓ−tk)dℓᵢ.
- �� Training: two rollout steps are unrolled, but gradients are cut through the first; K future slices are predicted synchronously.
Experiments
The main family was ∂tu+∂x(αu²−β∂xu+γ∂xxu)=δ, with periodic domain [0,16), 200 spatial cells for reference data, and 200 time points over [0,4]. The dataset contained 2,096 trajectories, downsampled to (nt,nx)=(250,100), (250,50), and (250,40). Reference solutions used WENO5 for convection, fourth-order finite differences for remaining terms, and adaptive RK4 time integration. Baselines were WENO5, FDM, pseudospectral methods, FNO-RNN, and FNO-PF. Metrics included accumulated MSE, runtime, and survival time under a normalized L1 threshold of 0.1.
Results
MP-PDE achieved E3 errors of 4.26, 3.74, and 3.70 for nx=100, 50, and 40, versus FNO-RNN errors of 10.16, 14.49, and 20.90 and FNO-PF errors of 5.69, 5.39, and 5.98. In E1, MP-PDE errors were 1.55, 1.67, and 1.47, while FNO-RNN reached 11.93, 29.98, and 10.44. Runtime was approximately 0.08–0.09 seconds for 250 steps. Equation conditioning mattered most in E3; Pushforward outperformed no perturbation and Gaussian noise.
Applications
The method is suited to repeated time evolution in transport, fluid, and conservation systems. Engineering groups could train it on high-fidelity simulations and use it for rapid parameter sweeps, design optimization, control, and interactive forecasting. Graph inputs are attractive for irregular domains and changing geometry. Deployment should nevertheless retain a trusted numerical solver for calibration, conservation checks, boundary validation, and detection of out-of-distribution states.
Limitations & Outlook
The study does not establish performance on large-scale 3D turbulence, stiff equations, strongly coupled multiphysics, or experimentally measured trajectories. Reference data are generated by numerical schemes, so learning may inherit their discretization bias. Pushforward improves empirical rollout survival but lacks a formal long-term error guarantee. Local message passing may limit distant propagation, while temporal bundle size K, processor depth, and neighborhood design introduce task-specific choices. Future work should add physics constraints, adaptive sparsity, conservation projections, and theory.
Plain Language Accessible to non-experts
Imagine a city whose streets carry traffic. Each small map square is a neighborhood, and roads connect neighboring squares. Every neighborhood reports how many cars it has, where it is, and how crowded nearby neighborhoods are. MP-PDE is a learned traffic coordinator: it listens to local reports and predicts how traffic will move in the next moments.
Traditional numerical methods are fixed traffic rules. Finite differences compare nearby locations; finite volumes count cars entering and leaving a block; WENO5 chooses more reliable nearby information when traffic changes sharply. The paper shows that a message-passing network can represent these local rules, then adjust them from examples. It is therefore not simply ignoring mathematics; it turns local mathematical updates into trainable components.
Long-term prediction is difficult because mistakes snowball. If the coordinator makes a small error, the next prediction treats that error as reality. Pushforward training deliberately exposes the model to its own imperfect states. Temporal Bundling lets it plan several future moments in one call, reducing repeated opportunities for drift. In experiments, it preserved shock-like structures even on coarse grids and adapted to changing equation parameters, while remaining much faster than repeated high-resolution simulation. It is promising, but complex three-dimensional cities—and real physical systems—still require careful testing.
ELI14 Explained like you're 14
Picture a water-flow video game. The map is divided into tiny squares, and every square stores information about how much water it has and where the water is moving. At every turn, you must predict the next turn. Old-school programs use carefully written rules. They can be excellent, but changing the map, the borders, or the kind of water may require a new program.
MP-PDE uses a team of tiny robots, one per square. Each robot talks to nearby robots, shares its value and position, and then updates its guess. A neural network learns how those conversations should affect the next water state. It can also receive the “settings” of the equation, so the same team can handle different kinds of flow.
Here is the tricky part: one wrong guess can snowball. If turn one is slightly wrong, turn two uses that wrong answer, and soon the whole game may look ridiculous. The Pushforward Trick makes training include the network’s own imperfect guesses. Temporal Bundling is like planning several game turns at once, so the system does not repeatedly restart from a shaky guess.
The researchers trained on 2,096 trajectories. In a mixed-parameter test with only 40 spatial cells, MP-PDE had error 3.70, compared with 20.90 for FNO-RNN. It also preserved shock waves and completed 250 steps in about 0.08 seconds. Awesome—but not magic! Three-dimensional turbulence, complicated physics, and absolute reliability still need much more testing.
Glossary
Message Passing Neural Network
A graph model in which each node receives information from neighboring nodes and updates its hidden state. It is a learnable form of local interaction and aggregation.
MP-PDE uses it to approximate local spatial operators on discretized PDE domains.
Finite Difference Method
A numerical method that replaces derivatives with weighted differences between nearby grid values. It is efficient but can be sensitive to discretization and stability conditions.
The paper interprets FDM stencils as shallow weighted message passing.
Finite Volume Method
A method that updates the average quantity in each cell through fluxes crossing cell boundaries. Its integral formulation naturally supports conservation.
The paper relates FVM flux updates to multi-layer message passing.
Pushforward Trick
A training strategy that uses states generated by the model as perturbed inputs, matching the distribution encountered during inference. It addresses autoregressive distribution shift.
The authors unroll two steps and backpropagate only through the second step.
Temporal Bundling
Predicting multiple future time slices in one network call rather than advancing exactly one step. It reduces calls and opportunities for error accumulation.
The temporal CNN Decoder outputs K future increments.
Zero-stability
The numerical requirement that small input perturbations do not grow excessively over time. It is distinct from accuracy but necessary for reliable long rollouts.
Pushforward training is motivated as directly reducing perturbation amplification.
Open Questions Unanswered questions from this research
- 1 Can empirical survival improvements from Pushforward training become rigorous nonlinear long-horizon error bounds? The paper motivates the connection to zero-stability but does not prove a general convergence theorem.
- 2 Does local message passing remain efficient for 3D complex geometry, stiff multiphysics, and long-range interactions? This requires large real-world datasets, unstructured meshes, and conservation-focused evaluations.
- 3 How should neighborhood size, processor depth, bundle length K, and adaptive time steps be selected automatically for robust transfer across PDE families?
Applications
Immediate Applications
Rapid fluid parameter sweeps
Engineering teams can generate trajectories with WENO5 or another trusted solver, then condition MP-PDE on α, β, γ and new initial states. The learned model can accelerate design screening, optimization, and interactive prediction, provided extreme cases are checked against high-fidelity simulation.
Surrogate simulation on irregular meshes
A complex region can be represented as a graph whose nodes and edges encode local geometry and boundary information. This supports fast approximate transport, heat, or flow prediction without rebuilding a fixed-grid architecture for every geometry.
Long-term Vision
General-purpose scientific simulation models
A future system could jointly encode equations, materials, geometry, boundaries, and initial conditions, enabling transferable simulators across physical domains. Major obstacles include strict conservation, uncertainty calibration, 3D scale, and safe out-of-distribution behavior.
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.