PDE-PFN: Prior-Data Fitted Neural PDE Solver

AI summary

Overview

Research area: Scientific machine learning (SciML) / neural PDE solvers; specifically prior-data fitted networks (PFNs) and in-context learning applied to partial differential equations.

Technical level: Advanced. The paper assumes familiarity with neural operators, physics-informed neural networks, posterior predictive inference, and Transformer attention mechanisms.

Scope: The paper introduces PDE-PFN, a Transformer-based, prior-data-fitted model that predicts PDE solution fields by in-context learning from noisy, approximate pre-training data, and evaluates it on 2D convection–diffusion–reaction equations, shallow water equations, and compressible Navier–Stokes equations.

What This Paper Is About

Existing scientific foundation models for PDEs are typically pre-trained on many equations and then fine-tuned downstream, but they inherit weaknesses from their backbones: they are not robust to noisy inputs and they often require rigid input formats such as uniform grids or fixed spatial locations. PDE-PFN instead borrows the prior-data fitted network idea from tabular and language modeling: the model is pre-trained to fit a posterior predictive distribution over PDE solutions, so that at inference it solves a new instance purely by conditioning on observed (coordinate, value) pairs, with no gradient updates, no task-specific fine-tuning, and no symbolic equation information. The goal is a single pre-trained checkpoint that generalizes across heterogeneous PDEs while remaining robust when the pre-training data are approximate rather than exact.

Key Contributions

1. A PFN-style neural PDE solver. PDE-PFN adapts a PFN architecture with self-attention within the context set and cross-attention from queries to context, and is trained to approximate the posterior predictive mean (PPM) of PDE solutions rather than the full posterior predictive distribution, which the authors argue would be computationally intractable.

2. Pre-training on deliberately noisy priors. Instead of exact numerical solutions, the model is pre-trained on approximate solutions generated by physics-informed neural networks (PINNs), described as "diverse but not necessarily exact priors." On the CDR pre-training set the generated PINN priors exhibit an average L2 relative error of 0.02280. The paper reports that the model often improves over these noisy approximate pre-training data.

3. Architectural extensions of the original PFN. A Fourier feature embedding is added before the encoder, extending the original PFN design that was limited to tabular data, and the attentive graph filter (AGF) structure is adopted for the Transformer layers to reduce computational complexity and improve representational capacity. An orthogonality regularization term (L_AGF), scaled by the number of context points N_D, is added to the MSE objective, giving the final loss L = L_α + (1/N_D) L_AGF.

4. Equation-free, shape-flexible inference. The model predicts solutions directly from observed quantities such as velocity and pressure without symbolic equation input, accepts variable-sized sets of coordinate–value observations as context, and takes target coordinates as queries, so it does not require a fixed observation grid.

Main Findings

• Within-family generalization on seen coefficients: On the family of 2D CDR equations with the pre-training coefficient dictionary (βx, βy, νx, νy, ρ1, ρ2, ρ3 ∈ {0.0, 0.5, 1.0}), both "Ours" and "Ours (PINN)" outperform the baselines on spatiotemporal interpolation and temporal extrapolation without task-specific fine-tuning. For spatiotemporal interpolation, L1 Abs is 0.01218 for Ours, 0.01507 for Ours (PINN), and 0.01922 for DeepONet; L2 Rel is 0.01643, 0.02057, and 0.02479 respectively. For temporal extrapolation, L1 Abs for Ours is 0.01474 versus 0.02731 for the next-best baseline F-FNO; L2 Rel is 0.02261 for Ours versus 0.03430 for F-FNO and 0.15580 for Poseidon.

• Coefficient generalization to unseen parameters: On unseen coefficients — interpolation at {0.25, 0.75} and extrapolation at {1.25, 1.5} — both model versions outperform the baselines on both tasks. Under coefficient extrapolation, L1 Abs on temporal extrapolation is 0.01027 for Ours and 0.02063 for Ours (PINN), compared with 0.02652 for FNO and 0.14661 for Poseidon.

• Cross-PDE transfer to the shallow water equations: Using the SWE dataset from PDEBench, with 700 of 1,000 samples for training and 300 held out as the unseen test set, both model versions outperform the baselines on seen and unseen data. Ours achieves the best spatiotemporal interpolation result (L2 Rel of 0.00739 on the seen set) while Ours (PINN) performs best on temporal extrapolation (L1 Abs of 0.01418 on the seen set, 0.01448 on the unseen set).

• Cross-PDE transfer to compressible Navier–Stokes equations: On CNSE operator learning across t = 0.25, 0.5, 0.75, and 1.0, both versions of the model outperform the baselines on the two relative error metrics, with L2 Rel of 0.54833 (Ours) and 0.54369 (Ours (PINN)) versus 1.09567 for DeepONet. The paper notes that DPOT's pre-training corpus already includes the CNSE dataset from PDEBench, yet the proposed models still perform better than DPOT on both relative metrics; DPOT's performance is lower than in its original paper because the task here was extended to predicting up to four time steps ahead.

• Robustness to noisy priors: The PINN-prior variant remains competitive despite being pre-trained on approximate solutions, which the paper presents as evidence of robustness to imperfect pre-training data.

• Zero-gradient-update in-context inference: A single pre-trained PDE-PFN checkpoint predicts query values from the observed context without task-specific gradient updates or fine-tuning. Unlike ICON-LM, which requires separate task/operator-level demonstrations before prediction, PDE-PFN uses only the available observations as the conditioning context.

• Theoretical framing: Theorem 2.1 establishes the existence of a neural network approximating the posterior predictive distribution; Theorem 3.1 shows that under squared loss the population-optimal deterministic predictor is the posterior predictive mean, and that under unbiased target noise the same population minimizer equals the clean-solution conditional mean given the noisy context.

Methodology in Plain English

The model is trained in an episodic way. For each training example, a PDE instance is drawn by sampling coefficients from a parameter space, and one set of points is designated as the "context" — coordinates paired with their solution values — while another set is designated as the "queries," where only the coordinates are given and the values are masked out. The Transformer attends within the context (self-attention) and from the queries to the context (cross-attention), and is trained to output the solution values at the query locations. The loss is simply the mean squared error between predicted and target values at the query points, plus a small regularization term for the AGF layers. Because minimizing squared error drives the predictor toward the conditional mean of the target given the context, the authors argue this trains the model to approximate the posterior predictive mean over PDE solution fields.

Crucially, the "prior data" used for training is not exact. The ground-truth solution family U is difficult to obtain, so the authors substitute a PINN prior D: a collection of approximate solutions from independently trained PINNs, intended to have a similar distribution to the exact data. At test time the model conditions only on the observations of the target instance and predicts at requested coordinates — no equation symbols, no retraining, and no requirement that the observations sit on a grid.

The empirical program has three phases: (1) set up datasets and baselines; (2) pre-train on the family of 2D CDR equations and test on seen coefficients and unseen coefficients (interpolation and extrapolation) with no fine-tuning; (3) fine-tune on the shallow water equations and the compressible Navier–Stokes equations to test cross-PDE transfer. Baselines evaluated are described in the text as six baseline models and listed as DeepONet, FNO, F-FNO, A-FNO, LNO, Poseidon, and DPOT; Poseidon and DPOT are used with publicly released pre-trained weights and fine-tuned, while the others are trained and evaluated separately per dataset and task. The pre-training data uses 6 time steps within t ∈ [0.0, 0.5] at intervals of 0.1, and the CDR coefficient dictionary (βx, βy, νx, νy, ρ1, ρ2, ρ3 ∈ {0.0, 0.5, 1.0}) yields 2,187 unique PDEs.

Why This Matters

Impact on research. The paper demonstrates that a single model can be pre-trained on approximate, noisy PDE solution data and then solve heterogeneous equations in-context, without gradient updates at inference and without being told the governing equations. The paper states this is the first work to address a wide range of reaction terms within the family of CDR equations using a single unified model. It also argues that the realistic problem of imperfect pre-training data — familiar from large language models — has been largely overlooked in the scientific foundation model literature.

Real-world applications cited or implied by the paper:

• Semiconductor manufacturing, where the paper notes observational data are often collected from irregularly distributed sensors rather than uniform grids.
• Settings where the underlying physical equations are incomplete or entirely unknown, which the paper gives as motivation for excluding explicit equations from the model input.
• Cross-PDE trajectory prediction, as demonstrated on the compressible Navier–Stokes equations operator-learning task over four future time steps.
• Generally adapting one pre-trained checkpoint to new PDE families through a reduced fine-tuning schedule, as shown on the shallow water equations using half as many epochs as the task-specific baselines.

The paper does not enumerate specific industries beyond semiconductor manufacturing.

Industry relevance. The combination of equation-free inference, tolerance for irregular observations, tolerance for noisy priors, and no per-instance retraining maps onto industrial settings where sensors are sparse or irregular and the physics are imperfectly known. The paper reports that PDE-PFN can beat a foundation model (DPOT) whose pre-training corpus already contained the CNSE dataset, suggesting that PFN-style prior fitting is competitive with large-scale pre-training for these benchmarks.

Future Directions

• Examining continuous coefficient priors and broader initial- and boundary-condition distributions, rather than the discrete coefficient dictionaries used here.
• Systematic scaling studies with respect to model size, prior size and diversity, context size, and distributional distance.
• Residual and conservation diagnostics, plus structured distribution shifts, to assess physical consistency of the predictions.
• Probabilistic prediction under ambiguous contexts, extending the current deterministic posterior-mean regression.
• Additional evaluations reported as supplementary: operator learning on the Airfoil dataset with an irregular mesh structure and on the Darcy Flow dataset across different features, along with a matched GNOT comparison and its attribution limitations in Appendix H.17.

Target Audience

Researchers and practitioners in scientific machine learning, neural operator development, and physics-informed modeling, particularly those working on foundation models for PDEs or on inverse and simulation problems with irregular sensor data. The paper is written for a reader already comfortable with PFNs, in-context learning, and neural operator benchmarks; readers seeking an introductory treatment of PDE solvers will find the theory and notation sections demanding.

Authors’ abstract

Motivated by the success of large language models (LLMs) with broad generalizability and robustness to noisy or unreliable pre-training data, we explore in-context prediction of PDE solution fields using approximate pre-training data. In addition, inspired by the posterior predictive mean-based inference mechanism of the in-context learning in prior-data fitted networks (PFNs), we propose PDE-PFN, a prior-data fitted neural solver that approximates the posterior predictive mean of PDE solutions via in-context learning. PDE-PFN builds on a PFN architecture with self- and cross-attention mechanisms of Transformer and is pre-trained on noisy approximate solutions generated by physics-informed neural networks, serving as diverse but not necessarily exact priors. Through experiments on a range of two-dimensional PDEs, we demonstrate that PDE-PFN achieves empirical generalization across heterogeneous equations, robustness under noisy priors, and zero-gradient-update in-context inference capability. PDE-PFN achieves competitive or improved predictive accuracy relative to the evaluated baselines while supporting flexible context and query representations.