PDE-PFN: Prior-Data Fitted Neural PDE Solver

Motivation and approach

• PDE-PFN uses PFN-style prior fitting for context-conditioned predictions, addressing noisy-input and input-rigidity drawbacks attributed to scientific foundation models.The paper evaluates within-family task and coefficient generalization, cross-PDE transfer after fine-tuning, and flexible context/query inputs.Source: §1 · p. 2

Within-family generalization and cross-PDE transfer

• Across evaluated settings, PDE-PFN achieved competitive or improved predictive accuracy relative to task-specific baselines and fine-tuned scientific foundation models.The evaluation includes within-family task generalization and cross-PDE transfer after fine-tuning.Source: §1 · p. 2

Flexibility on input shape without prior physical knowledge

• PDE-PFN accepts variable-sized coordinate–value observations and target-coordinate queries without requiring symbolic equations or a fixed observation grid.The introduction notes that irregularly distributed sensors are common in settings such as semiconductor manufacturing.Source: §1 · p. 2

Robustness to noisy prior data

• PDE-PFN shows strong predictive performance in the evaluated noisy-prior settings using approximate PINN solutions.Noisy and unreliable pre-training data are presented as an underexplored issue for scientific foundation models.Source: §1 · pp. 2–3

Zero-gradient-update in-context inference

• A single pre-trained checkpoint predicts query values from observed context without task-specific gradient updates or fine-tuning.The method adapts to a target PDE instance through in-context conditioning, without a separate demonstration set.Source: §1 · p. 3

Posterior predictive inference

• The framework treats observed coordinate–solution pairs as samples from a family of distributions indexed by an unknown parameter, then updates a prior using observed data to form a posterior and posterior predictive distribution.The parameter prior is a probability measure over the parameter space.Source: §2 · p. 3

Posterior predictive consistency

• For a well-behaved prior, the posterior predictive distribution converges toward the true distribution as data size increases, with cited work noting strong consistency in well-specified scenarios.The background identifies H as the Hellinger distance used in the convergence statement.Source: §2 · p. 3

Approximating the posterior predictive distribution

• Theorem 2.1 states that a neural network exists that approximates the posterior predictive distribution.The full theorem claim and proof are deferred to Appendix A.Source: §2 · p. 3

Posterior predictive mean

• Because modeling the full predictive distribution can be computationally intractable, the paper instead targets the posterior predictive mean for deterministic PDE prediction.The background describes this as the focus of the following section.Source: §2 · p. 3

Model schematic

• Figure 2 depicts context self-attention and query-to-context cross-attention, with training contexts and queries drawn from the prior and test inputs drawn from the ground-truth distribution.The model predicts solution values at the queried points.Source: Figure 2 · p. 4

• The setup represents spatiotemporal coordinates as (x, t), PDE solutions as outputs, and context/query sets as i.i.d. samples from solution distributions.Source: §3 · p. 4
• MSE training makes the population-optimal deterministic predictor the posterior predictive mean conditioned on context, including the clean-solution conditional mean under unbiased target noise.Source: §3 · pp. 4–5
• Fourier feature inputs and AGF-based Transformer layers adapt PFN architecture to PDE data while targeting lower computational complexity and greater representational capacity.Source: §3 · pp. 4–5
• The benchmark parameterizes two-dimensional CDR equations with convection, diffusion, and multiple nonlinear reaction terms, including cases with zero coefficients.Source: §3 · p. 5
• Because exact prior data are hard to obtain, training uses collections of PINN-predicted approximate solutions across a practical coefficient subspace.Source: §3 · p. 5
• The model performs posterior predictive mean inference through in-context learning, conditioning on the generated prior p(D).Source: §3 · pp. 4–5
• Training samples PDE parameters i.i.d. and minimizes queried-point MSE; an AGF orthogonality regularizer is scaled by context-point count.Source: §3 · pp. 5–6

4 Experiments

• Experiments progress from dataset and baseline setup to pre-training-family tasks without fine-tuning and heterogeneous-PDE transfer after fine-tuning.Table 1 is the experiment roadmap; its task labels include spatiotemporal interpolation and temporal extrapolation.Source: §4 · p. 6, Table 1 · p. 6

4.1 Experimental Setup

• The pre-training family contains 2,187 unique 2D CDR equations; generated PINN priors have an average L2 relative error of 0.02280.Six operator-learning baselines are listed, with training and evaluation procedures varying by task.Source: §4.1 · p. 6

4.2 Pre-training on the Family of 2D CDR Equations and Evaluation

• PDE-PFN takes observed coordinate–solution pairs as context and predicts masked solution values at query coordinates.The study compares versions pre-trained on numeric priors and PINN priors.Source: §4.2 · p. 7

4.2.1 Task Generalization for Seen coefficients

• On seen coefficients, both model versions outperform baselines on spatiotemporal interpolation and temporal extrapolation without task-specific fine-tuning.The PINN-prior results support robustness to approximation errors in the evaluated pre-training data.Source: §4.2.1 · p. 7, Table 2 · p. 7

4.2.2 Task Generalization for Unseen coefficients

• Without fine-tuning, both model versions outperform baselines on interpolation and temporal extrapolation with unseen coefficients, covering coefficient interpolation and extrapolation.Unseen coefficients are either intermediate values absent from pre-training or larger values than those used in pre-training.Source: §4.2.2 · p. 7, Table 3 · p. 8

4.3 Cross-PDE Transfer after Fine-tuning

• Cross-PDE transfer experiments fine-tune PDE-PFN on SWE for interpolation and extrapolation, and on CNSE for trajectory prediction.The SWE and CNSE datasets are from PDEBench.Source: §4.1 · p. 6, §4.3 · p. 8

4.3.1 Cross-PDE Transfer to Shallow Water Equations

• Both model versions outperform baselines on SWE spatiotemporal interpolation and temporal extrapolation for seen and unseen datasets.Ours performs best on interpolation, while Ours (PINN) performs best on extrapolation.Source: §4.3.1 · p. 8, Table 4 · p. 8

4.3.2 Cross-PDE Transfer to Compressible Navier-Stokes Equations

• Both PDE-PFN versions outperform baselines on two relative-error metrics for compressible Navier–Stokes trajectory prediction after fine-tuning.The evaluated task predicts up to four time steps ahead; DPOT’s original evaluation predicted one step ahead.Source: §4.3.2 · p. 9, Table 5 · p. 9

In-context learning

• Transformer in-context learning has been linked to emulating Bayesian predictors and linear models, performing Bayesian inference, and adapting algorithms to input sequences.Prior work also reports robustness across function classes and distribution shifts, and connections to implicit optimization dynamics.Source: §5 · p. 9

Context-query neural operators

• GNOT supports multiple input functions and irregular meshes, while UPT supports queries at arbitrary spatiotemporal locations.Both are context-query neural-operator architectures for prediction from flexible inputs and coordinate queries.Source: §5 · p. 9
• PDE-PFN emphasizes PFN-style episodic prior fitting over PDE instances and context/query sets, enabling one CDR checkpoint to support multiple prediction settings without task-specific updates.Cross-attention itself is not unique to PDE-PFN; a matched GNOT comparison and attribution limitations are reported in Appendix H.17.Source: §5 · p. 9

Scientific foundation models

• PDE-related foundation models include PDEformer for 1D PDE solving and inverse problems, ICON-LM for multimodal in-context operator learning, and ICON work on PDE generalization and few-shot learning.Other studies train a single model across diverse PDEs, but many rely on symbolic PDE information, task/operator-level demonstrations, or fine-tuning.Source: §5 · pp. 9–10

Conclusion

• PDE-PFN directly approximates the mean of the posterior predictive distribution of PDE solutions via in-context learning.The method is presented as a scientific machine-learning approach.Source: §6 · p. 10
• Experiments demonstrate CDR task and coefficient generalization without gradient updates, cross-PDE transfer after fine-tuning, flexible context/query inputs without symbolic equations at inference, and robustness in evaluated noisy-prior settings.The conclusion describes these as experimental findings across the studied settings.Source: §6 · p. 10
• The findings support PFN-style prior fitting as a promising approach to equation-free in-context prediction over PDE solution fields.Future work includes broader coefficient, initial- and boundary-condition, scaling, distribution-shift, physical-consistency, and probabilistic-prediction studies.Source: §6 · p. 10

Theorem A.1

• Theorem A.1 is described as establishing approximation of the full posterior predictive distribution as context grows, under its stated assumptions.This is not a convergence result for the deterministic PDE-PFN implementation.Source: A The Proof of Theorems · pp. 16–17

Theorem A.2

• The theorem’s posterior-consistency condition and equality of the specified coordinate distributions are among the stated conditions for its almost-sure conclusion.The supplied passage states these conditions but does not establish them for arbitrary PDE settings.Source: A The Proof of Theorems · pp. 16–17

Theorem A.1 proof

• The proof outline uses a Hellinger-distance bound via KL divergence and posterior consistency to argue that the relevant region vanishes asymptotically.The supplied proof passages describe partitioning by whether Hellinger distance exceeds ε and applying a triangle inequality.Source: A The Proof of Theorems · pp. 16–17

Theorem A.2

• Under squared loss, the population-optimal deterministic predictor is the conditional mean of the target solution given the query coordinate and observed context.The stated setup uses zero-mean observation noise independent across coordinates and independent of the latent PDE instance.Source: A The Proof of Theorems · p. 17

Scope of the results

• The result does not imply recovery of clean context observations or removal of arbitrary bias or correlated errors in PINN-generated priors.Its observation-noise setup specifies zero-mean, independent noise.Source: A The Proof of Theorems · pp. 17–18

Corollary A.3

• With nested context sets, the population-optimal squared loss cannot increase as more observations are provided.For noise-free query targets or evaluation against the clean solution, clean-solution prediction error can approach zero if the latent PDE instance becomes identifiable as context grows.Source: A The Proof of Theorems · p. 18

B.1 Family of 2D Convection-Diffusion-Reaction Equations

• The 2D CDR experiments use analytic solutions when available and spectral-method solutions otherwise, with 21 sampled times and 2,187 coefficient combinations.A separate PINN is trained for each of the 2,187 distinct equations to generate PINN priors.Source: §B.1 · p. 18

B.2 2D Shallow Water Equation

• The shallow-water dataset contains water-height simulations from 1,000 initial conditions, generated with PyClaw and subsampled to 21 time steps.The original simulation domain is [-2.5, 2.5]^2 over [0.0, 1.0] at resolution 128 × 128 × 101.Source: §B.2 · p. 19

B.3 2D Compressible Navier-Stokes Equation

• The compressible Navier–Stokes dataset contains four-feature velocity, pressure, and density fields from 1,000 initial conditions, subsampled to 64 × 64 × 5.The source simulations use a 128 × 128 × 21 grid, periodic boundaries, and Mach number 1.0 with both viscosity coefficients set to 0.1.Source: §B.3 · p. 19

B.4 Airfoil

• The Airfoil dataset represents compressible Euler-flow trajectories on 5,233 irregular mesh nodes, with node positions, types, velocity, pressure, and density.It spans 101 time steps from 0.0 to 4.8, spaced by 0.24, and the prediction operator targets four future time instances.Source: §B.4 · pp. 19–20

B.5 Darcy Flow

• Darcy Flow is evaluated as a time-independent elliptic operator mapping viscosity fields to steady-state solutions on a 128 × 128 grid.The dataset contains 2,000 samples, and its source term is fixed at 0.1.Source: §B.5 · p. 20

B.6 3D Diffusion-Reaction Equations

• The 3D diffusion-reaction experiments use 256 coefficient combinations on a 16 × 16 × 16 × 21 resolution.The coefficients νx, νy, νz, and ρ are drawn from {0.0, 0.5, 1.0, 1.5}.Source: §B.6 · p. 20

B.7 3D Gray-Scott

• The 3D Gray-Scott task predicts the two-species state five time steps ahead from its current state.The dataset uses 32 × 32 × 32 grids and the specified APEBench pattern regime f = 0.04, k = 0.06; training spans [0.0, 1000.0], while validation and testing span [0.0, 4000.0].Source: §B.7 · pp. 20–21

C PINN Used in Prior Generation

• PINN priors are generated using the PINN introduced by Raissi et al., with a loss defined over initial-condition, collocation, and boundary points.The reported configuration took approximately 373 seconds per 2D CDR coefficient combination and used up to 4,552 MB of memory.Source: C PINN Used in Prior Generation · pp. 21–22

D Detailed Description of the Baselines

• DeepONet, FNO, F-FNO, A-FNO, Poseidon, and DPOT provide distinct operator-learning or foundation-model baselines.Their approaches include branch-and-trunk networks, Fourier-based operators, multiscale transformers, and Fourier attention.Source: D Detailed Description of the Baselines · p. 22
• For spatiotemporal interpolation, only DeepONet was compared because the other baselines cannot process mesh-structured inputs; all models received matched context locations.Poseidon, DPOT, and PDE-PFN were trained for half the epochs used for other models, with the pretrained models initialized from released weights.Source: D Detailed Description of the Baselines · p. 22
• Neural operator baselines other than DeepONet generally require grid-structured inputs and predetermined target locations, while DeepONet requires fixed-coordinate inputs.PDE-PFN accepts variable-sized coordinate–value contexts and coordinate queries, though accuracy depends on prior coverage and context informativeness.Source: Table 7 · p. 22

E Model Architecture

• The architecture comprises an encoder, a Transformer block, and a decoder.The decoder maps Transformer latent representations to predicted solution values at queried coordinates.Source: E Model Architecture · pp. 22, 24, Figure 3 · p. 23
• Fourier feature embedding augments inputs with high-frequency components to help MLPs capture complex solution patterns.Coordinate inputs use sinusoidal mappings with frequencies ω_i = 2π/i; setting the FFE dimension to zero retains raw inputs only.Source: E Model Architecture · p. 23
• The encoder separately embeds context domains, context solutions, and query domains, then combines context-domain and context-solution representations for the Transformer.Inputs include spatial and temporal coordinates; query solution values are masked, and Fourier features are optional.Source: E Model Architecture · p. 23
• The Transformer block uses attentive graph filter layers, which learn to reweight attention’s spectral components while avoiding vanilla Transformer quadratic complexity.An orthogonality regularizer is included for the SVD matrices U and V.Source: E Model Architecture · pp. 23–24
• Self-attention operates among contexts, while cross-attention lets queries attend to contexts; the decoder’s MLP produces queried solution values.The model uses separate parameters for self-attention and cross-attention.Source: E Model Architecture · p. 24

H.1 2D CDR Equations with systematically Noisy Prior Data

• Across tested zero-mean Gaussian and uniform prior corruptions, accuracy is largely maintained; uniform-noise training is best on all interpolation metrics.On temporal extrapolation, uniform-noise training is best on L∞relative error and second best on two metrics; Gaussian-noise training remains competitive.Source: §H.1 · p. 25, Table 8 · p. 25

H.2 Test Time Evaluation Given PINN Prior in the Family of 2D CDR Equations

• With noisy PINN priors supplied at test time, Ours achieves the best performance, while Ours (PINN) trails DeepONet on two interpolation relative errors but leads on L1 absolute error.The authors attribute the relative-error result to Ours (PINN) potentially producing outputs closer to its PINN prior than to the analytic solution.Source: §H.2 · p. 25, Table 9 · p. 26

H.3 Test-Time Prediction with Biased Context Noise in the 2D CDR Family

• For biased Gaussian context noise, degradation remains relatively minor below 3% bias, while error increases more noticeably at approximately 3% bias or higher.The test uses σ = 0.01 noise standard deviation and evaluates L2 relative error for temporal extrapolation.Source: §H.3 · p. 26, Table 10 · p. 26

H.4 Joint Coefficient and Temporal Extrapolation in the 2D CDR Family

• PDE-PFN performs best when coefficient and temporal extrapolation are required simultaneously, while baselines fare worse in this joint setting.The temporal domain and coefficient domain both include unseen ranges.Source: §H.4 · p. 26, Table 11 · p. 26

H.6 2D CDR Equations with Random Initial Condition Sampled from Uniform Distribution

• PDE-PFN remains strong with changed or randomly sampled initial conditions, including best performance across all three temporal-extrapolation metrics for pixelwise uniform initial conditions.For a changed fixed initial condition it remains best; under random initial conditions, Ours (PINN) leads on two interpolation metrics while Ours is second on two.Source: §H.6 · p. 27, Table 13 · p. 27, §H.5 · p. 27

H.7 Operator Learning in Airfoil

• Both model versions outperform baselines on two metrics for Airfoil irregular-mesh operator learning and Darcy Flow mapping between viscosity and steady-state solution.The Airfoil task predicts solutions at four time steps from the initial condition; Darcy Flow transfers between different physical quantities.Source: §H.7 · p. 28, Table 14 · p. 28, §H.8 · p. 28

H.9 Spatiotemporal Interpolation and Temporal extrapolation in 3D Diffusion-Reaction equations

• For 3D diffusion–reaction equations, PDE-PFN wins every interpolation metric and two temporal-extrapolation metrics; increasing test context size also consistently lowers all three errors.The context-size trend is consistent with non-increasing population-optimal MSE, but does not establish monotonic improvement for every finite trained model or convergence of the full posterior distribution.Source: §H.9 · p. 28, Table 16 · p. 29, §H.11 · p. 29

H.10 Operator Learning in 3D Gray-Scott equation

• In 3D Gray-Scott operator learning, both PDE-PFN versions outperform all baselines on every metric, with Ours first and Ours (PINN) second.The results cover one-step prediction in a complex 3D system, including the model trained with noisy PINN priors.Source: §H.10 · p. 29, Table 17 · p. 29

H.12 Ablation Study

• The full model is consistently better on temporal extrapolation, while ablated variants can occasionally do slightly better on interpolation.Under PINN-prior evaluation, the full model is best on both tasks and both ablations degrade; AGF and rational activation are most beneficial for extrapolation and robustness.Source: §H.12 · pp. 30–31, Table 18 · p. 30, Table 19 · p. 30

H.13 Computational Cost Analysis related to the Number of Context and Query Points

• Runtime and memory show approximately linear scaling with context or query length in the profiled configurations.Measurements use BF16 mixed precision and fixed 4,096 query or context points when varying the other length; they do not establish asymptotic complexity.Source: §H.13 · p. 31, Table 20 · p. 31, Table 21 · p. 31

H.14 Context Identifiability in the CDR Family

• Replacing identical initial contexts with informative later observations increases separation recovery and reduces prediction errors, though substantial error remains.Variation across observation selections indicates that observation location matters; flexible context input cannot resolve ambiguity absent from observations.Source: §H.14 · p. 31, Table 22 · p. 32

H.15 Sensitivity to Context Time

• Errors are similar across tested context snapshots but do not fall monotonically as context approaches the query times.Snapshots at 0.60 and 0.65 lie outside the training-time domain, unlike those at 0.45 and 0.50, which may contribute to the pattern.Source: §H.15 · p. 32, Table 23 · p. 32

H.16 Periodic Boundary Consistency

• PDE-PFN shows lower periodic mismatch in this evaluation, indicating approximate consistency rather than exact boundary enforcement or PDE-residual guarantees.Endpoint pairs are queried off-grid without retraining or explicit boundary-condition constraints.Source: §H.16 · p. 32, Table 24 · p. 32

H.17 Comparison with GNOT

• PDE-PFN has lower errors than GNOT on all three metrics in 10 of 11 settings; the exception is interpolated-coefficient CDR interpolation, where GNOT leads on L2 and L∞ errors.PDE-PFN has lower L1 error in that exception, and the comparison does not isolate PFN-style prior fitting because training formulation and architecture both differ.Source: §H.17 · pp. 33–34, Table 25 · p. 33, Table 26 · p. 33

I.1 Detailed Experimental Results of PDE-PFN with Standard Deviation

• 5-seed results show consistently stable PDE-PFN performance across seen, interpolated, and extrapolated coefficients and both task types.Table 28 reports mean ± standard deviation over five random seeds.Source: §I.1 · p. 34, Table 28 · p. 34

I.2 Qualitative results

• For CDR, PDE-PFN closely reconstructs spatial patterns while DeepONet deviates more at later query times; for SWE, it better preserves water-height profiles than DeepONet in interpolation and FNO in extrapolation.Figures 5 and 6 show the CDR and SWE comparisons, respectively.Source: §I.2 · p. 34, Figure 5 · p. 34, Figure 6 · p. 35

J.1 Reproducibility

• The authors state that the code will be publicly available as supplementary material upon publication.They report that experimental procedures, data sources, and algorithms are documented in the paper.Source: §J.1 · p. 34

J.2 Environments

• Experiments were conducted using an NVIDIA RTX A6000; Python and package details are in the code environment file.Source: §J.2 · p. 35

J.3 Cost Comparison on the Main Experiments

• Table 29 compares model parameter counts, inference time, and GPU memory use on the 2D CDR temporal-extrapolation test process.Both PDE-PFN versions use the same hyperparameters and are reported under the single label Ours.Source: §J.3 · p. 35, Table 29 · p. 35

J.4 Hyperparameters

• Table 30 reports training, validation, and test data-point counts; unseen-coefficient CDR evaluations use testing only, without training or validation.Source: §J.4 · p. 35, Table 30 · p. 36
• Best hyperparameters were selected by validation performance, and the appendix lists configurations for the experiments.For fine-tuning when the predicted variable count differs, the encoder’s first linear layer and decoder’s final linear layer are reinitialized.Source: §J.4 · pp. 35–36, Table 31 · p. 36, Table 32 · p. 37

K Impact Statement

• The authors report no clear negative societal implications, while acknowledging potential future applications in climate modeling, biomedical simulations, and engineering systems.They also state that experiments complied with relevant laws and ethical guidelines and that their process is transparent and reproducible.Source: K Impact Statement · pp. 36, 38