35 Generative Models for PDE-Based Simulation

The Grand Challenge: Simulating Nature

“Since Newton, mankind has come to realise that the laws of physics are always expressed in the language of differential equations.” - Steven Strogatz

Imagine you are an engineer designing the wing of a next-generation aircraft. Before a single rivet is hammered, before a single sheet of aluminium is cut, you must answer a deceptively simple question: will this wing fly? More precisely, you need to know the pressure distribution, the lift-to-drag ratio, the onset of turbulence over the wing surface at cruising altitude, and how these quantities change under gusts, varying angles of attack, and different flight Mach numbers. The physical laws governing these phenomena-the compressible Navier–Stokes equations-have been known since the nineteenth century. Writing them down takes a single line. Solving them, however, can consume millions of CPU-hours on the world's largest supercomputers.

This is the central paradox of computational physics: the equations are simple, but their solutions are staggeringly complex. A turbulent flow at a Reynolds number of ${10}^6$ contains eddies spanning six orders of magnitude in spatial scale, and the computational grid must resolve every one of them for a faithful simulation. Direct Numerical Simulation (DNS) of the Navier–Stokes equations at industrially relevant Reynolds numbers remains out of reach for all but the simplest geometries, even on exascale hardware.

And yet, nature solves these equations effortlessly. Every ocean wave, every wisp of smoke curling from a candle, every heartbeat pushing blood through your arteries is an instantaneous, massively-parallel computation performed by the universe itself. The question that drives this chapter is: can we train neural networks to approximate the solutions of partial differential equations (PDEs) at a fraction of the computational cost of traditional solvers, while faithfully capturing the rich, multi-scale physics?

The answer, as we shall see, is a resounding-and increasingly precise-yes.

Why this chapter, and why now?

The marriage of deep learning and differential equations is not entirely new. Physics-Informed Neural Networks (PINNs), introduced by Raissi, Perdikaris, and Karniadakis in 2019, showed that one can embed PDE residuals directly into a neural network's loss function and obtain mesh-free approximate solutions. But PINNs, for all their elegance, are single-instance solvers: you train a network to solve one specific PDE with one specific set of initial and boundary conditions. Change the initial temperature distribution, the shape of the domain, or the viscosity coefficient, and you must start over.

The breakthrough of the last five years has been the shift from solving one PDE instance to learning the operator that maps PDE parameters to solutions. Neural operators-most famously the Fourier Neural Operator (FNO) and DeepONet-learn mappings between infinite-dimensional function spaces, so that a single trained model can produce solutions for an entire family of PDEs, evaluated at any desired resolution.

But even neural operators have a fundamental limitation: they are deterministic. They output a single “best guess” solution, typically the conditional mean under a squared-error loss. For many physical systems, especially turbulent flows, chaotic dynamics, and ill-posed inverse problems, the conditional mean is not a physically meaningful quantity. It is a blurred average of many possible states, much as the average face of all humans looks like no one in particular.

The two generative architectures that have proven most powerful in this context are Transformers and diffusion models:

• Transformers bring global self-attention over discretised fields, enabling the resolution of long-range spatial dependencies and the handling of irregular, non-Cartesian geometries that defeat convolutional approaches.

• Diffusion/score-based models cast PDE solving as an iterative denoising process governed by stochastic differential equations, naturally producing distributions over solution fields and enabling plug-and-play conditioning on sparse observations and physics constraints.

A roadmap for this chapter

This chapter is designed to be largely self-contained. We assume familiarity with basic PDE concepts (the heat equation, the wave equation, boundary conditions) at the level of a first undergraduate course, and with the deep learning fundamentals covered in Chapter ch:deeplearning of this book. Everything else-PINNs, neural operators, Transformer variants, diffusion formulations on function spaces-is developed from scratch with full mathematical detail, worked examples, and TikZ illustrations.

Here is what lies ahead:

1. A First Look at PDEs and Why They Matter (Section A First Look at PDEs and Why They Matter): We recall the classification of PDEs into elliptic, parabolic, and hyperbolic types, and motivate each class with a flagship physical application. We introduce the Navier–Stokes equations and the associated Clay Millennium Prize Problem-one of the seven greatest unsolved problems in mathematics.

2. The Computational Bottleneck (Section The Computational Bottleneck): We quantify why classical solvers are expensive, what the “curse of dimensionality” means for parametric PDEs, and why repeated-query settings (design optimisation, uncertainty quantification, Bayesian inversion, real-time control) demand fundamentally new approaches.

3. Physics-Informed Neural Networks (Section Physics-Informed Neural Networks): We develop the PINN framework rigorously, derive the physics-informed loss, discuss automatic differentiation for computing PDE residuals, and analyse the failure modes (spectral bias, gradient pathologies, multi-term loss balancing). We cover causality-aware training and gradient-rebalancing remedies.

4. Neural Operator Learning (Section Neural Operator Learning): We formalise the concept of a solution operator $𝒢:𝒜\to 𝒰$ mapping input functions to output functions, develop DeepONet and the Fourier Neural Operator from first principles, and discuss theoretical approximation guarantees.

5. Transformer-Based PDE Solvers (Section Transformer-Based PDE Solvers): We cover Galerkin Transformers, GNOT, Transolver, HAMLET, PiT, physics-attention, PDE-conditional Transformers (Unisolver), and foundation-model pretraining approaches (MPP, CoDA-NO, ICON, PDEformer).

6. Diffusion and Score-Based Models for PDEs (Section Diffusion and Score-Based Models for PDE Simulation): We develop the theory of diffusion models on function spaces, derive physics-informed diffusion (PIDMs), discuss DiffusionPDE, conditional diffusion for forecasting and data assimilation, CoCoGen, and flow-matching approaches for fast sampling.

7. Applications by PDE Family (Sections Poisson and Laplace-Type Elliptic PDEs–Coupled and Multiphysics PDE Systems): We walk through each major PDE class-Poisson/Laplace, Burgers, Navier–Stokes, shallow water, Helmholtz, Maxwell, Stefan, Cahn–Hilliard, Schrödinger, and coupled multiphysics systems-showing how Transformers and diffusion models have been applied, with results, figures, and critical analysis.

8. Turbulence and Generative Modelling (Section Generative Modelling for Turbulence): A dedicated treatment of the turbulence problem, including GAN-era approaches, diffusion for turbulent flow generation, CoNFiLD, GenCFD, and the interplay between spectral accuracy and generative diversity.

9. Inverse Problems and Uncertainty Quantification (Section PDE-Constrained Inverse Problems): We formalise PDE-constrained inverse problems, develop plug-and-play diffusion priors, RED-DiffEq for full waveform inversion, and InverseBench as a community benchmark.

10. Benchmarks and Evaluation (Section Benchmarks and Evaluation Practice): We survey PDEBench, PINNacle, and emerging evaluation standards, discussing metrics beyond $L^2$ (energy norms, spectral fidelity, conservation residuals).

11. Open Problems and the Road Ahead (Section Open Problems and Future Challenges): We discuss the Clay Millennium Prize Problem on Navier–Stokes existence and smoothness, the role of ML in conjecture generation for PDEs, and frontier challenges including discretisation-agnostic generation, real-time turbulence, and generative PDE foundation models.

12. Exercises (Section Pen-and-Paper Exercises): A comprehensive collection of exercises ranging from pen-and-paper derivations to coding projects.

A First Look at PDEs and Why They Matter

Partial differential equations are the mathematical language of continuum physics. Wherever a quantity varies smoothly in space and time-temperature, pressure, electric potential, probability density-its evolution is typically governed by a PDE\@. In this section, we revisit the three fundamental types of second-order linear PDEs (elliptic, parabolic, hyperbolic), introduce their flagship physical incarnations, and build toward the Navier–Stokes equations that will serve as the running example throughout this chapter.

The three archetypes

Consider a general second-order linear PDE in two independent variables $(x,y)$: (General2nd)$$\begin{array}{ccc}& A\frac{{\partial }^{2}u}{\partial x^2}+2B\frac{{\partial }^{2}u}{\partial x\partial y}+C\frac{{\partial }^{2}u}{\partial y^2}+\text{(lower-order terms)}=0.& \text{<span class="tml-eqn"></span>}\end{array}$$ The discriminant ${\mathrm{\Delta }}_{\mathrm{PDE}}=B^2-AC$ classifies the equation:

• Elliptic (${\mathrm{\Delta }}_{\mathrm{PDE}}<0$): steady-state problems, e.g., the Laplace equation ${\nabla }^{2}u=0$. Solutions are smooth and encode equilibrium configurations.

• Parabolic (${\mathrm{\Delta }}_{\mathrm{PDE}}=0$): diffusion and heat conduction, e.g., ${\partial }_{t}u=\kappa {\nabla }^{2}u$. Information propagates at infinite speed but is smoothed over time.

• Hyperbolic (${\mathrm{\Delta }}_{\mathrm{PDE}}>0$): wave propagation, e.g., ${\partial }_{tt}u=c^2{\nabla }^{2}u$. Information propagates at finite speed along characteristics; discontinuities (shocks) can form.

The Navier–Stokes equations: beauty and terror

The incompressible Navier–Stokes equations govern the motion of a viscous, incompressible fluid: (Momentum)$$\begin{array}{cccc}& \frac{\partial 𝐮}{\partial t}+(𝐮\cdot \nabla )𝐮& =-\frac{1}{\rho }\nabla p+\nu {\nabla }^2𝐮+𝐟,& \text{<span class="tml-eqn"></span>}\\ & \nabla \cdot 𝐮& =0,& \text{<span class="tml-eqn"></span>}\end{array}$$ where $𝐮(𝐱,t)\in {ℝ}^d$ is the velocity field, $p(𝐱,t)$ is the pressure, $\rho$ is the (constant) density, $\nu$ is the kinematic viscosity, and $𝐟$ represents external body forces. Equation is the momentum equation (Newton's second law for a fluid element), and Equation is the incompressibility constraint (conservation of mass for a constant-density fluid).

The nonlinear advection term $(𝐮\cdot \nabla )𝐮$ is what makes the Navier–Stokes equations simultaneously beautiful and terrifying. It couples the velocity field to itself, creating feedback loops that can amplify small perturbations into the chaotic, multi-scale structures we call turbulence.

Kolmogorov's theory of turbulence

In 1941, Andrey Kolmogorov proposed a statistical theory of fully-developed turbulence that remains one of the most profound insights in fluid dynamics. His key hypotheses:

1. Energy cascade: Energy is injected into the flow at large scales (the integral scale $\ell$), cascades through a hierarchy of progressively smaller eddies, and is finally dissipated by viscosity at the Kolmogorov microscale $\eta$.

2. Universal scaling: In the inertial range (scales between $\ell$ and $\eta$), the energy spectrum follows the celebrated $-5/3$ power law: (Kolmogorov)$$\begin{array}{ccc}& E(k)\propto {\epsilon }^{2/3}{k}^{-5/3},& \text{<span class="tml-eqn"></span>}\end{array}$$ where $k$ is the wavenumber and $\epsilon$ is the mean energy dissipation rate per unit mass.

The Kolmogorov microscale is given by: (Kolmogorov ETA)$$\begin{array}{ccc}& \eta ={\left(\frac{{\nu }^3}{\epsilon }\right)}^{1/4},& \text{<span class="tml-eqn"></span>}\end{array}$$ and the ratio of the largest to smallest scales goes as $\ell /\eta \sim {\mathrm{Re}}^{3/4}$. In three dimensions, the number of grid points needed for DNS scales as: (DNS COST)$$\begin{array}{ccc}& N_{\text{grid}}\sim {\left(\frac{\ell }{\eta }\right)}^3\sim {\mathrm{Re}}^{9/4}.& \text{<span class="tml-eqn"></span>}\end{array}$$ For an aircraft wing at $\mathrm{Re}\sim {10}^7$, this gives $N_{\text{grid}}\sim {10}^{15.75}\approx 5.6\times {10}^{15}$ grid points-far beyond current computational capabilities.

The Clay Millennium Prize Problem

In the year 2000, the Clay Mathematics Institute announced seven Millennium Prize Problems, each carrying a reward of $1,000,000. The Navier–Stokes problem is among them:

What makes this problem so difficult? The nonlinear advection term $(𝐮\cdot \nabla )𝐮$ can amplify vorticity through the vortex stretching mechanism. In two dimensions, vortex stretching is absent, and global regularity was proven by Ladyzhenskaya in the 1960s. In three dimensions, vortex stretching can potentially concentrate vorticity to a point, creating a singularity. Despite decades of effort by the world's best analysts, neither existence nor blow-up has been proven in general.

The Computational Bottleneck

We have seen that DNS of turbulent flows is prohibitively expensive. But the computational bottleneck extends far beyond turbulence. In this section, we examine three settings where classical PDE solvers become a fundamental limitation, motivating the need for neural surrogates and generative alternatives.

The many-query problem

In many scientific and engineering workflows, one needs to solve not one PDE but a family of PDEs, parametrised by varying inputs:

• Design optimisation: Finding the wing shape that minimises drag involves evaluating the Navier–Stokes equations for thousands of candidate geometries. Each evaluation requires a full PDE solve.

• Uncertainty quantification (UQ): If the material properties (e.g., permeability in a porous medium) are uncertain, one must propagate this uncertainty through the PDE solver. Monte Carlo sampling requires $𝒪({10}^3)$–$𝒪({10}^6)$ PDE solves.

• Bayesian inversion: Estimating subsurface geological properties from seismic measurements requires evaluating the forward PDE model at every step of an MCMC chain.

• Real-time control: Active flow control (e.g., reducing drag on a vehicle) requires evaluating the PDE model in milliseconds, far faster than any classical solver can operate.

In all these settings, the PDE solver is called inside a loop, and the total computational cost is: (Total COST)$$\begin{array}{ccc}& \text{Total cost}=N_{\text{queries}}\times C_{\text{solver}},& \text{<span class="tml-eqn"></span>}\end{array}$$ where $C_{\text{solver}}$ is the cost of a single solve and $N_{\text{queries}}$ can range from hundreds to millions.

The curse of dimensionality

Classical PDE solvers discretise the domain $\mathrm{\Omega }$ into a mesh with $N$ degrees of freedom. For a $d$-dimensional domain with $n$ grid points per dimension, $N=n^d$. The cost of a single solve typically scales as:

• $𝒪(N)$ for explicit time-stepping (e.g., forward Euler),

• $𝒪(N\log N)$ for spectral methods with FFT,

• $𝒪(N^{1+\alpha })$ for sparse direct solvers ($\alpha \approx 0.5$–$1$ depending on dimension and method).

When the PDE is parametric-i.e., we want solutions for all $\theta$ in some parameter space $\mathrm{\Theta }\subset {ℝ}^p$-the effective dimensionality of the problem is $d+p$. For high-dimensional parameter spaces (e.g., uncertain spatially-varying coefficients $\kappa (𝐱)$ discretised on a grid), $p$ can be in the thousands, and traditional methods like sparse grids or polynomial chaos expansions break down.

What classical solvers do well-and where they fall short

Before we dive into neural alternatives, it is important to acknowledge what classical numerical methods excel at:

1. Rigorous error bounds: Finite element methods come with a priori and a posteriori error estimates in well-defined norms ($H^1$, $L^2$, etc.).

2. Conservation properties: Finite volume methods conserve mass, momentum, and energy by construction.

3. Mesh adaptivity: Adaptive mesh refinement (AMR) can concentrate resolution where it is needed (e.g., near shocks or boundary layers).

4. Decades of engineering validation: These methods have been tested and trusted in safety-critical applications (aircraft certification, nuclear reactor design, weather prediction).

Neural surrogates do not (yet) provide rigorous error bounds, do not conserve physical quantities by construction, and have limited track records in safety-critical settings. The goal is not to replace classical solvers but to complement them-using neural models for fast exploration and reserving classical solvers for final validation.

The landscape of neural approaches to PDEs

The neural PDE solving landscape can be organised along two axes:

1. Single-instance vs. operator learning: Does the model solve one specific PDE (PINNs) or learn the solution operator for a family (neural operators)?

2. Deterministic vs. generative: Does the model output a single solution (point estimate) or a distribution over solutions?

Table Table 1 provides a quantitative comparison of the key approaches along several dimensions.

Method	Scope	Output	Mesh-free?	UQ?	Inverse?
PINNs	Single	Deterministic	Yes	Limited	Limited
FNO	Operator	Deterministic	No${}^{†}$	No	No
DeepONet	Operator	Deterministic	Yes	No	No
Transolver	Operator	Deterministic	Yes	No	No
DiffusionPDE	Operator	Generative	No	Yes	Yes
PIDMs	Operator	Generative	No	Yes	Yes
FunDPS	Operator	Generative	Yes	Yes	Yes
${}^{†}$FNO has “zero-shot super-resolution” on uniform grids but does not handle unstructured meshes.

Physics-Informed Neural Networks

We now arrive at one of the most influential ideas in scientific machine learning: using a neural network as a trial function for a PDE, and training it by penalising the PDE residual evaluated at random collocation points. This simple but powerful idea-the Physics-Informed Neural Network (PINN)-was popularised by Raissi, Perdikaris, and Karniadakis (2019) and has since spawned an enormous literature.

The core idea

Suppose we wish to solve a PDE of the general form: (General)$$\begin{array}{ccc}& 𝒩[u](𝐱,t)=0,(𝐱,t)\in \mathrm{\Omega }\times (0,T],& \text{<span class="tml-eqn"></span>}\end{array}$$ subject to initial condition $u(𝐱,0)=u_0(𝐱)$ and boundary condition $ℬ[u](𝐱,t)=0$ for $𝐱\in \partial \mathrm{\Omega }$. Here $𝒩$ is a (possibly nonlinear) differential operator and $ℬ$ encodes boundary conditions (Dirichlet, Neumann, Robin, periodic, etc.).

The PINN approach is breathtakingly simple:

1. Represent $u$ by a neural network: Let $u_{\varphi }(𝐱,t)$ be a multilayer perceptron (MLP) with parameters $\varphi$, taking $(𝐱,t)$ as input and returning the predicted solution value.

2. Compute PDE residuals via automatic differentiation: The partial derivatives $\partial u_{\varphi }/\partial t$, ${\partial }^2{u}_{\varphi }/\partial x_i\partial x_j$, etc. are computed exactly (up to floating-point precision) using the chain rule through the network's computational graph. No finite differences or mesh are needed.

3. Minimise a physics-informed loss: The loss function penalises the PDE residual, boundary residual, and initial condition residual at randomly sampled collocation points: (LOSS)$$\begin{array}{ccc}& ℒ(\varphi )={\lambda }_r{ℒ}_r(\varphi )+{\lambda }_b{ℒ}_b(\varphi )+{\lambda }_0{ℒ}_0(\varphi ),& \text{<span class="tml-eqn"></span>}\end{array}$$ where: (LOSS Residual)$$\begin{array}{cccc}& {ℒ}_r(\varphi )& =\frac{1}{N_r}\sum _{i=1}^{N_r}|𝒩[u_{\varphi }]({𝐱}_i^r,t_i^r){|}^2,& \text{<span class="tml-eqn"></span>}\\ & {ℒ}_b(\varphi )& =\frac{1}{N_b}\sum _{i=1}^{N_b}|ℬ[u_{\varphi }]({𝐱}_i^b,t_i^b){|}^2,& \text{<span class="tml-eqn"></span>}\\ & {ℒ}_0(\varphi )& =\frac{1}{N_0}\sum _{i=1}^{N_0}|u_{\varphi }({𝐱}_i^0,0)-u_0({𝐱}_i^0){|}^2.& \text{<span class="tml-eqn"></span>}\end{array}$$

Automatic differentiation for PDE residuals

The magic of PINNs lies in automatic differentiation (AD). Unlike finite differences, which introduce truncation errors and require a mesh, AD computes exact derivatives through the computational graph.

For a network $u_{\varphi }(𝐱,t)$ with $L$ layers: (MLP)$$\begin{array}{ccc}& u_{\varphi }(𝐱,t)=W_L\sigma (W_{L-1}\sigma (\cdots \sigma (W_1[𝐱;t]+b_1)\cdots )+b_{L-1})+b_L,& \text{<span class="tml-eqn"></span>}\end{array}$$ where $\sigma$ is the activation function. The key observation is that each partial derivative $\partial u_{\varphi }/\partial x_j$ can be computed by the chain rule: (Chainrule)$$\begin{array}{ccc}& \frac{\partial u_{\varphi }}{\partial x_j}=\frac{\partial u_{\varphi }}{\partial h^{(L)}}\cdot \frac{\partial h^{(L)}}{\partial h^{(L-1)}}\cdots \frac{\partial h^{(1)}}{\partial x_j},& \text{<span class="tml-eqn"></span>}\end{array}$$ where $h^{(k)}$ denotes the hidden representation at layer $k$. Higher-order derivatives (needed for PDEs like the heat equation or Navier–Stokes) require multiple applications of this chain rule.

Failure modes of PINNs

Despite their elegance, PINNs suffer from several well-documented failure modes that limit their applicability to challenging PDEs. Understanding these failures is essential both for improving PINNs and for appreciating why generative alternatives are needed.

Spectral bias

Neural networks with smooth activations (tanh, sigmoid, softplus) have an intrinsic preference for low-frequency functions. This spectral bias causes PINNs to learn the low-frequency components of the solution first and converge extremely slowly (or not at all) on high-frequency features such as:

• Sharp shocks in Burgers' equation,

• Rapidly oscillating solutions of the Helmholtz equation,

• Fine-scale turbulent structures.

Mathematically, the Neural Tangent Kernel (NTK) of an MLP with smooth activations has eigenvalues that decay rapidly with frequency. For a target function with significant high-frequency content, the corresponding NTK eigenvalues are small, and the gradient flow converges slowly along those directions.

Multi-term loss balancing

The PINN loss is a weighted sum of multiple terms. The weights ${\lambda }_r$, ${\lambda }_b$, ${\lambda }_0$ are hyperparameters that dramatically affect convergence. If ${\lambda }_b$ is too small, the network may satisfy the PDE in the interior but violate boundary conditions; if ${\lambda }_r$ is too small, the network may interpolate the initial/boundary data but ignore the physics.

Gradient pathologies

A deeper issue is that the gradients ${\nabla }_{\varphi }{ℒ}_r$ and ${\nabla }_{\varphi }{ℒ}_b$ can point in conflicting directions. Stiff PDEs (where the PDE residual varies by many orders of magnitude across the domain) exacerbate this: the optimiser may oscillate between reducing the interior residual and the boundary residual, making net progress extremely slow.

Wang, Teng, and Perdikaris (2021) provided a systematic analysis of these gradient pathologies, showing that the condition number of the effective Hessian of the PINN loss can be astronomically large for stiff systems, leading to training failure.

Causality violation

For time-dependent PDEs, collocation points are typically sampled uniformly over the entire spatiotemporal domain $\mathrm{\Omega }\times [0,T]$. This means the optimiser simultaneously tries to fit the solution at $t=0$ (where the initial condition is known) and at $t=T$ (where the solution must be propagated through the dynamics). The network may “cheat” by fitting later times without correctly propagating from the initial condition, resulting in solutions that are smooth and low-residual but physically meaningless.

Remedies and extensions

The failures of vanilla PINNs have spurred an impressive body of work on improved training strategies. We highlight several:

1. Fourier feature embeddings: Replacing the raw input $(𝐱,t)$ with random Fourier features $[\sin (B𝐱);\cos (B𝐱)]$ (where $B$ is a random matrix) alleviates spectral bias by mapping the input to a higher-dimensional space where all frequencies are equally accessible. This is closely related to the random Fourier features kernel approximation.

2. Adaptive activation functions: Jagtap and Karniadakis (2020) introduced learnable parameters in the activation function (e.g., $\sigma (n\cdot a\cdot x)$ where $n$ is a scaling factor and $a$ is learnable), allowing the network to adapt its frequency response during training.

3. Domain decomposition (XPINNs): The domain $\mathrm{\Omega }$ is divided into subdomains, each with its own PINN, with interface conditions enforced via additional loss terms. This reduces the effective complexity each network must capture.

4. Time-stepping PINNs: Instead of solving over the entire time interval $[0,T]$ at once, the interval is divided into windows $[t_k,t_{k+1}]$, and a PINN is trained sequentially on each window, using the previous solution as the initial condition.

5. Curriculum learning: Start training on an easier version of the PDE (e.g., higher viscosity) and gradually increase the difficulty, allowing the network to build up its representation incrementally.

From single instances to operators: the limits of PINNs

The fundamental limitation of PINNs is that they are single-instance solvers. A PINN trained to solve the Burgers equation with $\nu =0.01/\pi$ and initial condition $u_0(x)=-\sin (\pi x)$ cannot, without retraining, solve the same equation with $\nu =0.1$ or $u_0(x)=\exp (-x^2)$.

This limitation motivates the operator learning paradigm, which we develop in the next section: instead of learning the solution $u$ for fixed parameters $\theta$, we learn the operator $𝒢:\theta \mapsto u$ that maps parameters to solutions across the entire family.

Neural Operator Learning

In the previous section we saw how PINNs embed physics into a neural network's loss function to solve one PDE instance at a time. Each new set of initial conditions, boundary conditions, or coefficients requires retraining from scratch. This is rather like hiring a brilliant chef who can cook exactly one dish: delicious, but impractical for running a restaurant.

The conceptual leap of neural operator learning is to move from learning a single function $u:\mathrm{\Omega }\to ℝ$ to learning the operator $𝒢:𝒜\to 𝒰$ that maps an input function (coefficient field, initial condition, forcing term) to the corresponding solution function. Once $𝒢$ is learned, we can evaluate it on any new input function at inference time with a single forward pass-no retraining, no iterative solver, no mesh.

This section develops the mathematical foundations of operator learning, presents the two most important architectures-DeepONet and the Fourier Neural Operator-derives key theoretical guarantees, and concludes by identifying the fundamental limitations that motivate the generative approaches of the remaining sections.

The Operator Learning Problem

Function spaces and the solution operator

We begin by formalising the setting. Let $\mathrm{\Omega }\subset {ℝ}^d$ be a bounded domain, and consider a parametric family of PDEs: (Parametric)$$\begin{array}{ccc}& {𝒩}_a[u]=0\text{in }\mathrm{\Omega },\text{subject to boundary conditions on }\partial \mathrm{\Omega },& \text{<span class="tml-eqn"></span>}\end{array}$$ where $a:\mathrm{\Omega }\to {ℝ}^{d_a}$ is an input function that parametrises the PDE. The input $a$ could be:

• A spatially varying coefficient field (e.g., the permeability in Darcy flow),

• An initial condition (e.g., $u(𝐱,0)=a(𝐱)$ for the heat equation),

• A forcing term (e.g., the right-hand side $f(𝐱)$ in Poisson's equation),

• A combination of the above.

For each admissible input $a$, the PDE has a unique solution $u=u_a$.¹ The solution operator is the mapping that sends $a$ to $u_a$:

The superscript $†$ denotes the true (data-generating) operator; the goal of operator learning is to find an approximation ${𝒢}_{\theta }\approx {𝒢}^{†}$, parametrised by learnable parameters $\theta$.

The data-driven formulation

In practice, we are given a training dataset of $N$ input–solution pairs: (Dataset)$$\begin{array}{ccc}& 𝒟=\{(a^{(i)},u^{(i)}){\}}_{i=1}^N,u^{(i)}={𝒢}^{†}(a^{(i)}),& \text{<span class="tml-eqn"></span>}\end{array}$$ where each pair is generated by running a classical PDE solver (finite elements, finite volumes, spectral methods) on the input $a^{(i)}$. The functions $a^{(i)}$ are drawn from some distribution $\mu$ on $𝒜$ (e.g., a Gaussian random field), and both $a^{(i)}$ and $u^{(i)}$ are observed on discrete grids.

The operator learning problem is then: (Objective)$$\begin{array}{ccc}& \underset{\theta }{\min }{𝔼}_{a\sim \mu }\left[{\Vert {𝒢}_{\theta }(a)-{𝒢}^{†}(a)\Vert }_{𝒰}^2\right]\approx \frac{1}{N}\sum _{i=1}^N{\Vert {𝒢}_{\theta }(a^{(i)})-u^{(i)}\Vert }^2,& \text{<span class="tml-eqn"></span>}\end{array}$$ where the norm $\Vert \cdot \Vert$ is approximated by a discrete quadrature rule on the observation grid.

Universal approximation for operators

The theoretical foundation for operator learning was laid remarkably early-long before the deep learning revolution-by Chen and Chen in 1995.

This theorem says two profound things:

1. Any continuous operator can be approximated to arbitrary accuracy by a finite-rank architecture involving two sub-networks, one acting on the input function $a$ (evaluated at fixed sensor locations $x_1,\dots ,x_m$) and one acting on the output query point $y$.

2. The input function $a$ need only be observed at finitely many sensor locations-we do not need the entire function, only its values at $m$ representative points.

DeepONet

The Chen–Chen theorem provides an existence result, but it took until 2019–2021 for Lu, Jin, Pang, Zhang, and Karniadakis to transform it into a practical deep learning architecture: DeepONet (Deep Operator Network).

Architecture

DeepONet is a direct realisation of equation . It consists of two sub-networks:

1. Branch network ${𝐛}_{{\theta }_{\mathrm{br}}}:{ℝ}^m\to {ℝ}^p$: Takes as input the values of the input function $a$ at $m$ fixed sensor locations $\{x_1,\dots ,x_m\}$ and outputs a vector of $p$ coefficients $(b_1(a),\dots ,b_p(a))$.

2. Trunk network ${𝐭}_{{\theta }_{\mathrm{tr}}}:{ℝ}^d\to {ℝ}^p$: Takes as input a query location $y\in 𝒴$ and outputs a vector of $p$ basis function values $(t_1(y),\dots ,t_p(y))$.

The output is their inner product:

The beauty of this factorisation is its interpretability: the trunk network learns a set of $p$ basis functions ${\{t_k\}}_{k=1}^p$ for the output space, while the branch network learns the coefficients of the expansion as a function of the input $a$. This is precisely a learned generalisation of classical basis expansions such as Fourier series, proper orthogonal decomposition (POD), and reduced basis methods.

Training procedure

Training DeepONet requires a dataset of input–solution pairs ${\{(a^{(i)},u^{(i)})\}}_{i=1}^N$, where each $u^{(i)}$ is observed at (possibly different) query locations ${\{y_j^{(i)}\}}_{j=1}^{Q_i}$. The loss is the mean squared error over all observations: (LOSS)$$\begin{array}{ccc}& ℒ(\theta )=\frac{1}{N}\sum _{i=1}^N\frac{1}{Q_i}\sum _{j=1}^{Q_i}|{𝒢}_{\theta }(a^{(i)})(y_j^{(i)})-u^{(i)}(y_j^{(i)}){|}^2.& \text{<span class="tml-eqn"></span>}\end{array}$$

Key practical considerations include:

• Sensor placement. The $m$ sensor locations $\{x_1,\dots ,x_m\}$ are fixed across all training samples and must be chosen before training. A uniform grid is a common default, but adaptive or random sensor placement can improve performance for problems with localised features.

• Query point sampling. The query points $\{y_j^{(i)}\}$ can be sampled randomly from $\mathrm{\Omega }$ at each training iteration, acting as a form of data augmentation and stochastic quadrature.

• Data generation. Each training pair requires a full PDE solve. For expensive PDEs, this can be the bottleneck. Common strategies include using coarser meshes for training data, multi-fidelity approaches, or augmenting with physics-informed losses.

• Number of basis functions $p$. Increasing $p$ gives more capacity but also more parameters. In practice, $p\in [50,200]$ works well for many benchmark problems.

Physics-informed DeepONet

Just as PINNs augment data losses with PDE residuals, we can create a physics-informed DeepONet (PI-DeepONet) by adding a residual penalty. Since the trunk network maps a query point $y$ to the output value, we can differentiate the DeepONet output with respect to $y$ using automatic differentiation and penalise the PDE residual: (PI Deeponet)$$\begin{array}{ccc}& {ℒ}_{\text{PI}}(\theta )={ℒ}_{\text{data}}(\theta )+{\lambda }_r\frac{1}{N}\sum _{i=1}^N\frac{1}{M}\sum _{j=1}^M|{𝒩}_{a^{(i)}}[{𝒢}_{\theta }(a^{(i)})](y_j^r){|}^2.& \text{<span class="tml-eqn"></span>}\end{array}$$

This hybrid approach reduces the amount of labelled data needed by injecting known physics, and is particularly valuable when full-field solution data is expensive to generate.

Worked example: DeepONet for Darcy flow

Variants and extensions

Several important extensions of the basic DeepONet architecture have been proposed:

• POD-DeepONet. Replace the learnable trunk with the leading POD modes of the training solutions: $t_k(y)={\varphi }_k(y)$ are fixed, and only the branch network is trained. This reduces the parameter count and can improve accuracy when the solution manifold is well-approximated by a linear subspace.

• MIONet (Multiple-Input Operator Network). Extends DeepONet to handle multiple input functions (e.g., both an initial condition and a forcing term) via separate branch networks whose outputs are combined by a tensor product: (Mionet)$$\begin{array}{ccc}& {𝒢}_{\theta }(a_1,a_2)(y)=\sum _{k_1,k_2}{b}_{k_1}(a_1)c_{k_2}(a_2)t_{k_1,k_2}(y).& \text{<span class="tml-eqn"></span>}\end{array}$$

• Stacked DeepONet. For time-dependent problems, the output $u(𝐱,t)$ at time $t_{n+1}$ is predicted using the solution at $t_n$ as input: ${𝒢}_{\theta }(u^n)\approx u^{n+1}$. This autoregressive approach enables long time horizon predictions, though error accumulation remains a challenge.

Fourier Neural Operator

While DeepONet builds on the Chen–Chen universal approximation theorem, the Fourier Neural Operator (FNO), introduced by Li, Kovachki, Azizzadenesheli, Liu, Bhatt, Stern, and Anandkumar (2021), takes an entirely different-and arguably more elegant-route. The key insight is both simple and powerful: if the PDE solution operator involves convolutions and global interactions, why not parameterise the integral kernel directly in Fourier space?

The key insight: convolutions in Fourier space

Many PDEs-particularly linear ones-have solution operators that can be expressed as integral transforms: (Integral)$$\begin{array}{ccc}& (𝒦v)(x)={\int }_{\mathrm{\Omega }}\kappa (x,y)v(y)\mathrm{d}y,& \text{<span class="tml-eqn"></span>}\end{array}$$ where $\kappa$ is the Green's function or integral kernel. For translation-invariant kernels ($\kappa (x,y)=\kappa (x-y)$), this becomes a convolution, and the convolution theorem tells us: (CONV Theorem)$$\begin{array}{ccc}& ℱ[\kappa \ast v](\xi )=\hat{\kappa }(\xi )\cdot \hat{v}(\xi ),& \text{<span class="tml-eqn"></span>}\end{array}$$ where $ℱ$ denotes the Fourier transform and $\hat{\cdot }$ the Fourier coefficients. Convolution in physical space becomes pointwise multiplication in Fourier space.

This is the founding observation of FNO: instead of learning the kernel $\kappa$ in physical space (which requires $O(N^2)$ parameters for an $N$-point discretisation), we directly learn the Fourier-space weights $R_{\varphi }(\xi )$, which requires only $O(k)$ parameters if we truncate to $k$ Fourier modes.

The Fourier layer

The FNO consists of a sequence of Fourier layers, each performing a spectral convolution combined with a pointwise nonlinearity. A single Fourier layer ${ℱ}_{\ell }$ maps a function $v_{\ell }:\mathrm{\Omega }\to {ℝ}^{d_v}$ to $v_{\ell +1}:\mathrm{\Omega }\to {ℝ}^{d_v}$ as follows:

Let us unpack the spectral convolution step by step:

1. Forward FFT: Compute ${\hat{v}}_{\ell }(\xi )=ℱ[v_{\ell }](\xi )$ for all retained modes $|\xi |\le k_{\max }$.

2. Fourier-space multiplication: For each retained mode $\xi$, apply the learnable weight: ${\hat{w}}_{\ell }(\xi )=R_{{\varphi }_{\ell }}(\xi ){\hat{v}}_{\ell }(\xi )$, where $R_{{\varphi }_{\ell }}(\xi )\in {ℂ}^{d_v\times d_v}$.

3. Mode truncation: Set all Fourier modes with $|\xi |>k_{\max }$ to zero. This is both a regulariser and the source of the model's efficiency.

4. Inverse FFT: Compute $({𝒦}_{\ell }{v}_{\ell })(x)={ℱ}^{-1}[{\hat{w}}_{\ell }](x)$, returning to physical space.

The pointwise linear term $W_{\ell }{v}_{\ell }(x)$ plays a crucial role: it captures local, high-frequency interactions that are lost due to Fourier mode truncation. Together, the spectral convolution handles global structure while the pointwise term handles local features.

Full FNO architecture

The complete FNO pipeline consists of three stages:

1. Lift: A pointwise linear layer lifts the input from the physical dimension $d_a$ to a higher-dimensional hidden representation of dimension $d_v$: (LIFT)$$\begin{array}{ccc}& v_0(x)=Pa(x)+p_0,P\in {ℝ}^{d_v\times d_a},p_0\in {ℝ}^{d_v}.& \text{<span class="tml-eqn"></span>}\end{array}$$

2. Fourier layers: A stack of $L$ Fourier layers (typically $L=4$): (Stack)$$\begin{array}{ccc}& v_0\stackrel{\stackrel{}{{ℱ}_1}}{\to }{v}_1\stackrel{\stackrel{}{{ℱ}_2}}{\to }{v}_2\stackrel{\stackrel{}{{ℱ}_3}}{\to }{v}_3\stackrel{\stackrel{}{{ℱ}_4}}{\to }{v}_4.& \text{<span class="tml-eqn"></span>}\end{array}$$

3. Project: A pointwise MLP projects from the hidden dimension $d_v$ back to the output dimension (typically 1): (Project)$$\begin{array}{ccc}& u_{\theta }(x)=Q{v}_L(x)+q_0,Q\in {ℝ}^{1\times d_v},q_0\in ℝ.& \text{<span class="tml-eqn"></span>}\end{array}$$

Zero-shot super-resolution

One of the most remarkable properties of FNO is its discretisation invariance: because the learnable parameters $R_{\varphi }(\xi )$ live in Fourier space and are indexed by frequency rather than grid point, the same trained model can be evaluated on grids of any resolution.

This zero-shot super-resolution is impossible for grid-based methods like standard CNNs, which have a fixed input/output resolution determined by their architecture.

Complexity analysis

Let us compare the computational cost of different approaches for an operator on a $d$-dimensional domain discretised with $N$ points per dimension (total grid size $N^d$):

FNO for benchmark PDEs

The FNO has been benchmarked on several canonical PDE problems. We highlight two here.

FNO variants

The success of FNO has inspired a rich family of extensions:

• U-FNO. Adds U-Net-style skip connections between Fourier layers at different resolutions, improving multi-scale modelling.

• Geo-FNO. Extends FNO to non-rectangular domains by learning a diffeomorphism that maps the irregular domain to a regular one where FFT can be applied.

• Factorised FNO (F-FNO). Factorises the $d$-dimensional Fourier convolution into $d$ one-dimensional convolutions along each axis, reducing complexity from $O(k_{\max }^d)$ to $O(d\cdot k_{\max })$.

• Adaptive FNO. Learns which Fourier modes to retain rather than using a fixed truncation, adapting the spectral representation to the problem structure.

• Spherical FNO (SFNO). Replaces the Cartesian FFT with spherical harmonic transforms, enabling FNO on the sphere-crucial for weather and climate modelling. This variant has been central to recent breakthroughs in data-driven weather prediction.

Theoretical Foundations

Having seen the practical architectures, we now develop the theoretical machinery that underpins neural operator learning. The central questions are: Can neural operators approximate any continuous operator? How efficiently? And what are the fundamental limits?

Universal approximation revisited

We have already stated the Chen–Chen theorem (Theorem Theorem 2) for DeepONet-type architectures. The analogous result for FNO was established by Kovachki, Lanthaler, and Mishra (2021):

Approximation rates

Universal approximation tells us that neural operators can approximate any continuous operator, but it says nothing about how many parameters are needed. Approximation rate theory addresses this crucial question.

The rate $n^{-s}$ mirrors the decay of the Kolmogorov $n$-width, confirming the intuition that operator learning is easy when the solution manifold is effectively low-dimensional.

The Kolmogorov $n$-width

The Kolmogorov $n$-width is a fundamental concept from approximation theory that quantifies the “intrinsic dimensionality” of a set of functions.

In plain English: $d_n(K)$ is the best worst-case error when approximating elements of $K$ by elements of an optimal $n$-dimensional linear subspace. If $d_n(K)$ decays rapidly (e.g., exponentially), then $K$ can be well-represented by a small number of basis functions, and operator learning is “easy.” If $d_n(K)$ decays slowly (e.g., algebraically), the solution manifold is inherently high-dimensional, and any method-classical or neural-faces fundamental difficulties.

Connection to reduced basis methods

Neural operators can be understood as nonlinear generalisations of classical reduced-order models (ROMs). The connection is illuminating:

Error decomposition

The total error of a neural operator can be decomposed into three sources:

This decomposition guides architecture design: increasing model capacity (more layers, modes, channels) reduces ${\epsilon }_{\text{approx}}$ but may increase ${\epsilon }_{\text{gen}}$, while more data reduces ${\epsilon }_{\text{gen}}$, and better optimisers (learning rate schedules, pre-training, curriculum learning) reduce ${\epsilon }_{\text{opt}}$.

Limitations of Deterministic Neural Operators

Neural operators represent a transformative advance over classical solvers and PINNs for many applications. Yet they have fundamental limitations that no amount of architectural engineering can overcome. Understanding these limitations is essential for appreciating why generative approaches to PDE solving-the subject of the remaining sections of this chapter-are not merely interesting alternatives but, for many problems, a necessity.

Regression to the mean

The most severe limitation of deterministic neural operators is the regression-to-the-mean problem. When trained with a mean squared error (MSE) loss , the optimal predictor is the conditional mean: (COND MEAN)$$\begin{array}{ccc}& {𝒢}_{\theta }^{\ast }(a)=𝔼[u|a]=\int up(u|a)\mathrm{d}\mu (u).& \text{<span class="tml-eqn"></span>}\end{array}$$

For deterministic PDEs with unique solutions, this is perfectly fine-the conditional mean is the unique solution. But for chaotic or turbulent systems, the mapping from input to solution is effectively one-to-many: small perturbations in the input (below the measurement precision) lead to vastly different solutions. In such cases, the conditional mean is a blurry average of many possible states, and is not itself a physically realisable solution.

To make this more precise, consider the vorticity field $\omega$ of a turbulent flow. The energy spectrum $E(k)\sim k^{-5/3}$ (Kolmogorov's $-5/3$ law) characterises the cascade of energy from large to small scales. The conditional mean $\bar{\omega }=𝔼[\omega |a]$ suppresses high-frequency fluctuations, leading to a steeper energy spectrum that violates the $-5/3$ scaling. This is not an artefact of poor training; it is a mathematical consequence of the MSE objective.

No uncertainty quantification

A deterministic neural operator outputs a single point prediction $\hat{u}={𝒢}_{\theta }(a)$ with no accompanying measure of confidence. In safety-critical applications-structural engineering, nuclear reactor design, weather forecasting-practitioners need to know how uncertain the prediction is, not just the prediction itself.

There are post-hoc approaches to add uncertainty estimates to deterministic models:

• Ensemble methods: Train $M$ independent neural operators and use the spread of their predictions as a proxy for uncertainty. However, ensemble diversity may not correlate with true posterior variance, especially in out-of-distribution regions.

• MC Dropout: Apply dropout at inference time and interpret the variation across forward passes as uncertainty. This is computationally cheap but often underestimates uncertainty.

• Bayesian neural operators: Place priors on the network weights and perform (approximate) posterior inference. This is principled but computationally demanding for large-scale models.

None of these approaches produce calibrated, high-dimensional distributions over solution fields. Generating diverse, physically consistent samples from the solution distribution requires a fundamentally different approach: generative modelling.

Fixed discretisation for grid-based methods

While FNO enjoys discretisation invariance (via the Fourier representation), many other neural operators-particularly those based on CNNs, U-Nets, or graph neural networks-are tied to a specific mesh or grid resolution. Changing the resolution requires retraining or introduces interpolation errors. Even for FNO, the zero-shot super-resolution property degrades for large resolution jumps, as discussed in Remark Remark 15.

Training data requirements

Each training pair $(a^{(i)},u^{(i)})$ requires a full PDE solve, which can be expensive. For 3D time-dependent problems (e.g., turbulent Navier–Stokes), a single high-fidelity simulation may require hours on a GPU cluster. Generating thousands of such pairs for training can cost more than the savings gained at inference time, especially if the model is used for only a few evaluations.

The break-even point depends on the application:

• Many-query settings (design optimisation, Bayesian inversion, real-time control): The offline training cost is amortised over thousands or millions of queries, making neural operators highly cost-effective.

• Few-query settings: The upfront data generation and training cost may not be justified. This motivates physics-informed approaches that reduce or eliminate the need for labelled data.

Motivating generative approaches

The limitations above can be summarised in a single observation: deterministic neural operators learn the wrong quantity for stochastic, chaotic, or ill-posed problems. They learn the conditional mean $𝔼[u|a]$, when what we need is the conditional distribution $p(u|a)$.

Exercises

Transformer-Based PDE Solvers

In the preceding sections we developed two families of neural PDE solvers: Physics-Informed Neural Networks (PINNs), which embed the PDE residual into the loss function of a single-instance neural network, and Neural Operators (FNO, DeepONet), which learn mappings between infinite-dimensional function spaces from data. Both families have achieved remarkable successes, yet both suffer from a common architectural limitation: they struggle to capture long-range spatial dependencies in the solution field.

PINNs use multilayer perceptrons (MLPs) whose receptive field is, in principle, global but whose spectral bias causes them to learn low-frequency components first and high-frequency details painfully slowly. The Fourier Neural Operator circumvents this by operating directly in spectral space, but at the cost of being restricted to uniform grids-a serious limitation for the complex, irregular geometries that arise in real-world engineering (airfoils, turbine blades, blood vessels, city-scale wind simulations).

Enter the Transformer. Originally developed for natural language processing, the Transformer's core mechanism-self-attention- computes pairwise interactions between all tokens in a sequence, regardless of their distance. When the “tokens” are discretisation points (or patches, or learned physical states) of a PDE domain, self-attention naturally captures the long-range, nonlocal interactions that PDEs encode: a pressure wave at one end of a pipe affects the velocity field everywhere; a vortex shed from the leading edge of an airfoil influences the wake far downstream.

This section provides a comprehensive and mathematically rigorous treatment of Transformer-based PDE solvers. We begin with the conceptual and mathematical motivations (Section Why Transformers for PDEs?), then proceed through the major architectures in roughly chronological order: the Galerkin Transformer (Section The Galerkin Transformer), GNOT (Section GNOT: General Neural Operator Transformer), Transolver (Section Transolver), HAMLET (Section HAMLET: Graph Transformer Neural Operator), and PiT (Section PiT: Position-induced Transformer). We then cover the emerging wave of PDE foundation models and conditioning strategies (Section PDE-Conditional Transformers and Foundation Models), in-context operator learning (Section In-Context Operator Learning), and finally the PINNsFormer architecture that replaces the MLP backbone of PINNs with a Transformer (Section PINNsFormer).

Why Transformers for PDEs?

To appreciate why Transformers have become the dominant architecture for PDE solving, we must understand three foundational insights: (i) self-attention captures long-range dependencies that local convolutions miss, (ii) attention operates on sets rather than grids, enabling natural handling of irregular meshes, and (iii) attention is intimately related to the integral operators that constitute the backbone of PDE theory.

Self-attention captures long-range dependencies

Consider an elliptic PDE such as the Poisson equation $-{\nabla }^{2}u=f$ on a domain $\mathrm{\Omega }$ with Dirichlet boundary conditions. The solution at any interior point $𝐱$ depends on the source term $f$ and the boundary data everywhere in the domain, through the Green's function: (Green)$$\begin{array}{ccc}& u(𝐱)={\int }_{\mathrm{\Omega }}G(𝐱,𝐲)f(𝐲)\mathrm{d}𝐲+{\oint }_{\partial \mathrm{\Omega }}\frac{\partial G}{\partial n_{𝐲}}(𝐱,𝐲)g(𝐲)\mathrm{d}S(𝐲),& \text{<span class="tml-eqn"></span>}\end{array}$$ where $G(𝐱,𝐲)$ is the Green's function of the Laplacian on $\mathrm{\Omega }$ and $g$ is the boundary data.

This integral representation reveals a fundamental nonlocality: the solution at $𝐱$ depends on the values of $f$ at every point $𝐲\in \mathrm{\Omega }$. A convolutional neural network with a kernel of size $k$ has a receptive field that grows only linearly with depth (or as $𝒪(k^L)$ after $L$ layers), and therefore requires very deep architectures to capture global interactions. The standard self-attention mechanism, by contrast, computes pairwise interactions between all input tokens in a single layer:

In the PDE context, we interpret the $N$ rows as spatial discretisation points (mesh nodes, collocation points, or patches). The output at each point is a function of all other points, mediated by learned similarity weights. This is precisely the discrete analogue of the integral operator in equation .

Irregular mesh handling

A second, equally important advantage is that attention operates on sets of tokens, not on regular grids. A convolutional layer assumes that its input lives on a regular Cartesian lattice; applying a CNN to an unstructured triangular mesh or a point cloud requires either interpolation to a regular grid (losing geometric fidelity) or specialised graph-convolutional layers (which reintroduce locality).

Self-attention has no such constraint. Given $N$ points with arbitrary coordinates $\{{𝐱}_1,\dots ,{𝐱}_N\}\subset {ℝ}^d$, we simply encode each point into a query-key-value triplet and let attention determine the interaction weights. Positional information can be injected through positional encodings-either fixed (sinusoidal, random Fourier features) or learned-so that the model is aware of the spatial locations without requiring a grid structure.

Attention as an integral operator

The deepest reason that Transformers are well suited to PDE solving is that self-attention is a discretised integral operator. Consider a general kernel integral operator of the form (Integral OP)$$\begin{array}{ccc}& (𝒦v)(𝐱)={\int }_{\mathrm{\Omega }}\kappa (𝐱,𝐲)v(𝐲)\mathrm{d}𝐲,& \text{<span class="tml-eqn"></span>}\end{array}$$ where $\kappa :\mathrm{\Omega }\times \mathrm{\Omega }\to ℝ$ is a kernel function and $v:\mathrm{\Omega }\to ℝ$ is the input function. This is a Fredholm integral operator of the second kind-precisely the form that appears in Green's function representations, boundary element methods, and the integral formulations of many PDEs.

Now discretise the domain with $N$ points $\{{𝐱}_1,\dots ,{𝐱}_N\}$ carrying quadrature weights $\{w_1,\dots ,w_N\}$. The discrete approximation is (Discrete Integral)$$\begin{array}{ccc}& (𝒦v)({𝐱}_i)\approx \sum _{j=1}^N{w}_j\kappa ({𝐱}_i,{𝐱}_j)v({𝐱}_j).& \text{<span class="tml-eqn"></span>}\end{array}$$ Comparing with the attention output at position $i$: (ATTN Integral)$$\begin{array}{ccc}& \mathrm{Attn}{(Q,K,V)}_i=\sum _{j=1}^N\underset{{\displaystyle \approx w_j\kappa ({𝐱}_i,{𝐱}_j)}}{\underbrace{\frac{\exp (q_i^{\top }{k}_j/\sqrt{d_k})}{{\sum }_{\ell }\exp (q_i^{\top }{k}_{\ell }/\sqrt{d_k})}}}{v}_j,& \text{<span class="tml-eqn"></span>}\end{array}$$ we see that the softmax-normalised attention weights play the role of the kernel $\kappa$ evaluated at the discretisation points, and the value vectors $v_j$ play the role of the function values.

The quadratic bottleneck and its remedies

The elephant in the room is computational cost. Standard self-attention has $𝒪(N^2)$ time and memory complexity, where $N$ is the number of tokens (discretisation points). For a 2D PDE on a $256\times 256$ grid, $N=\mathrm{65,536}$, and the attention matrix has $4.3\times {10}^9$ entries-far too large for most GPUs.

This has motivated a rich line of work on efficient attention for PDE solving:

• Linear attention: Replace softmax with a kernel factorisation $\varphi {(q_i)}^{\top }\psi (k_j)$ so that ${\mathrm{Attn}}_i=\varphi (q_i{)}^{\top }({\sum }_j\psi (k_j)v_j^{\top }\left)/\right(\varphi (q_i{)}^{\top }{\sum }_j\psi (k_j))$, reducing complexity to $𝒪(N{d}^2)$. This is the approach of the Galerkin Transformer.

• Cross-attention with learned queries: Introduce $M\ll N$ learnable “slice” queries and compute cross-attention $𝒪(NM)$. This is the approach of Transolver.

• Sparse or local attention: Restrict attention to local neighbourhoods plus a few global tokens, as in Longformer-style architectures.

• Downsampling + global attention: Pool the discretisation to a coarser representation, apply global attention, and upsample. This multiscale strategy is used by GNOT and several others.

The Galerkin Transformer

The Galerkin Transformer, introduced by Cao (NeurIPS 2021), was one of the first architectures to draw a rigorous connection between Transformer attention and classical numerical methods for PDEs. The key insight is that standard softmax attention approximates a Petrov–Galerkin projection, and that by choosing a different normalisation, one can obtain an attention mechanism that corresponds to a Galerkin projection-a cornerstone of finite element theory.

Background: the Galerkin method

Recall the classical Galerkin method for solving a linear PDE. Given an operator equation $ℒu=f$ in a Hilbert space $H$, we seek an approximate solution $u_N={\sum }_{j=1}^N{c}_j{\varphi }_j$ in a finite-dimensional subspace $V_N=\mathrm{span}\{{\varphi }_1,\dots ,{\varphi }_N\}$ such that the residual is orthogonal to $V_N$: (Galerkin Classical)$$\begin{array}{ccc}& ⟨ℒu_N-f,{\varphi }_i{⟩}_H=0,i=1,\dots ,N.& \text{<span class="tml-eqn"></span>}\end{array}$$ This yields the linear system $𝐀𝐜=𝐟$ where $A_{ij}=⟨ℒ{\varphi }_j,{\varphi }_i⟩$ and $f_i=⟨f,{\varphi }_i⟩$.

From softmax attention to Galerkin attention

Cao observed that standard softmax attention computes: (Softmax ATTN)$$\begin{array}{ccc}& {\mathrm{Attn}}_i^{\text{softmax}}=\frac{{\sum }_j\exp (q_i^{\top }{k}_j/\sqrt{d})v_j}{{\sum }_j\exp (q_i^{\top }{k}_j/\sqrt{d})},& \text{<span class="tml-eqn"></span>}\end{array}$$ where the normalisation is row-wise: each query $q_i$ produces its own probability distribution over keys. This corresponds to a Petrov–Galerkin method where the test functions (queries) and trial functions (keys) live in different spaces.

The Galerkin Transformer replaces this with two alternative attention mechanisms:

Galerkin-type attention.

Instead of row-wise softmax normalisation, Galerkin-type attention normalises column-wise over the keys, corresponding to a symmetric Galerkin projection where test and trial functions come from the same space: (Galerkin ATTN)$$\begin{array}{ccc}& {\mathrm{Attn}}^{\text{Galerkin}}=\frac{1}{N}Q(K^{\top }V).& \text{<span class="tml-eqn"></span>}\end{array}$$ Here $K^{\top }V\in {ℝ}^{d_k\times d_v}$ is computed first (at cost $𝒪(N{d}_k{d}_v)$), and then multiplied by $Q$ (at cost $𝒪(N{d}_k{d}_v)$), giving an overall linear complexity of $𝒪(N{d}_k{d}_v)$ instead of $𝒪(N^2{d}_v)$.

Fourier-type attention.

An alternative normalisation normalises column-wise over the queries: (Fourier ATTN)$$\begin{array}{ccc}& {\mathrm{Attn}}^{\text{Fourier}}=\frac{1}{N}(Q^{\top }K)V.& \text{<span class="tml-eqn"></span>}\end{array}$$ The name “Fourier” comes from the observation that when the keys and values are Fourier basis functions, this computation reduces to computing Fourier coefficients-a discrete Fourier-like transform.

Architecture details

The full Galerkin Transformer architecture wraps the Galerkin attention mechanism in a standard Transformer encoder-decoder structure:

1. Input encoding: The input function $a(𝐱)\in 𝒜$ (e.g., initial condition, forcing term, or coefficient field) is evaluated at $N$ discretisation points. Each point is encoded as $h_i^{(0)}=\text{MLP}([a({𝐱}_i),{𝐱}_i])$ to produce a $d$-dimensional token.

2. Galerkin attention layers: A stack of $L$ layers, each consisting of Galerkin multi-head attention followed by a feed-forward network with residual connections and layer normalisation. In the multi-head variant, each head $h$ computes: (Galerkin Multihead)$$\begin{array}{ccc}& {\text{head}}_h=\frac{1}{N}{Q}_h(K_h^{\top }{V}_h),Q_h=X{W}_h^Q,K_h=X{W}_h^K,V_h=X{W}_h^V,& \text{<span class="tml-eqn"></span>}\end{array}$$ and the outputs are concatenated and projected: $\text{MHA}(X)=[{\text{head}}_1;\dots ;{\text{head}}_H]W^O$.

3. Output decoding: A pointwise MLP maps the final hidden representation to the predicted solution value at each discretisation point. Crucially, the output points need not coincide with the input points-the decoder can interpolate to any desired query location by using cross-attention between query positions and the encoded representation.

GNOT: General Neural Operator Transformer

While the Galerkin Transformer elegantly reduces attention complexity and connects to classical numerical methods, it was primarily demonstrated on problems with a single input function on a regular grid. Many real-world PDE problems involve multiple heterogeneous inputs-an initial condition, a forcing term, boundary data, and possibly geometric parameters-evaluated on irregular, non-uniform meshes. The General Neural Operator Transformer (GNOT), introduced by Hao, Wang, and co-authors at ICML 2023, addresses these challenges through a carefully designed multi-input, multi-scale architecture.

Problem setup and challenges

GNOT considers the general operator learning problem: given a collection of input functions $\{a_1(𝐱),\dots ,a_P(𝐱)\}$ and scalar parameters $𝝁\in {ℝ}^p$ (e.g., Reynolds number, viscosity), learn the solution operator (GNOT Operator)$$\begin{array}{ccc}& 𝒢:(a_1,\dots ,a_P,𝝁)\mapsto u,& \text{<span class="tml-eqn"></span>}\end{array}$$ where the input functions may be evaluated on different, irregular meshes. For instance, in a fluid–structure interaction problem, the fluid velocity might be known on a tetrahedral mesh while the structural displacement is given on a surface triangulation.

The key challenges are:

1. Heterogeneous inputs: Different input types (functions on meshes, scalar parameters, boundary conditions on surface meshes) must be unified into a single representation.

2. Irregular geometry: The discretisation points are not on a regular grid, precluding the use of FFT-based methods like FNO.

3. Scale: Real-world meshes can have ${10}^4$–${10}^6$ nodes, requiring efficient attention.

Heterogeneous attention

GNOT introduces a heterogeneous normalised attention mechanism that handles multiple input types within a unified framework.

This heterogeneous attention allows the model to attend to different input modalities with different attention patterns while producing a unified output representation.

Multi-scale architecture

To handle large meshes efficiently, GNOT uses a multi-scale approach inspired by the Perceiver architecture:

1. Encoding: Input tokens (from all input types) are projected to a common embedding dimension $d$. Scalar parameters are broadcast to all tokens via an additive bias.

2. Cross-attention downsampling: A set of $M\ll N$ learnable latent tokens $Z\in {ℝ}^{M\times d}$ attends to the full input via cross-attention: (GNOT Cross)$$\begin{array}{ccc}& Z^{\prime }=\mathrm{CrossAttn}(Z,X)+Z,& \text{<span class="tml-eqn"></span>}\end{array}$$ where $X\in {ℝ}^{N\times d}$ is the concatenation of all input tokens. This reduces the sequence length from $N$ to $M$.

3. Self-attention on latent tokens: Multiple layers of self-attention are applied to the $M$ latent tokens, at cost $𝒪(M^{2}d)$-much cheaper than $𝒪(N^2)$ when $M\ll N$.

4. Cross-attention upsampling: The output tokens at the desired query locations attend back to the processed latent tokens: (GNOT Decode)$$\begin{array}{ccc}& Y=\mathrm{CrossAttn}(X_{\text{query}},Z_{\text{final}})+X_{\text{query}}.& \text{<span class="tml-eqn"></span>}\end{array}$$

The total complexity is $𝒪(NM+M^2)$, which is linear in $N$ when $M=𝒪(\sqrt{N})$ or $M=𝒪(1)$.

Handling multiple PDE settings

A distinctive feature of GNOT is its ability to handle multiple PDE settings within a single model. By conditioning on scalar parameters $𝝁$ (e.g., Reynolds number, Mach number), the same GNOT model can produce solutions for:

• Different PDE coefficients (varying viscosity in Navier–Stokes).

• Different domain geometries (by encoding geometry as an input function, e.g., a signed distance field).

• Different boundary conditions (by including boundary data as a separate input stream).

Transolver

The Transolver (Wu, Hao, Wang, et al., ICML 2024 Spotlight) represents a paradigm shift in how Transformers tokenise PDE domains. Rather than treating each mesh node as a token-which leads to $𝒪(N^2)$ attention for $N$ nodes-Transolver introduces the concept of intrinsic physical states: learned, compact representations that capture the underlying physics of different spatial regions. Attention is computed between these physical states rather than between individual mesh nodes, reducing complexity from $𝒪(N^2)$ to $𝒪(NM)$ where $M\ll N$ is the number of physical states.

Motivation: physics-aware tokenisation

Consider a complex flow around an airfoil. From a physics perspective, the domain naturally decomposes into a few distinct regions:

• The boundary layer (thin region of high shear near the airfoil surface).

• The wake region (downstream of the trailing edge).

• The freestream region (far from the airfoil, approximately uniform flow).

• Stagnation and separation regions (near the leading edge and any flow separation points).

Within each region, the physics is relatively homogeneous: the solution has similar local structure and the governing balances (e.g., inertia vs. viscosity) are similar. This motivates grouping nearby mesh nodes with similar physical behaviour into a single “physical state” token, and computing attention at the level of these physical states.

Physics-attention mechanism

The Transolver architecture has two main stages: (i) slice assignment, which maps each of the $N$ mesh nodes to one of $M$ physical states, and (ii) physics-attention, which computes attention among the physical states and then maps back to the original mesh nodes.

Step 1: Adaptive slice assignment.

Given the input representation $X\in {ℝ}^{N\times d}$ (from the previous layer or the input embedding), compute a soft assignment matrix: (Transolver Assign)$$\begin{array}{ccc}& S=\mathrm{softmax}\left(\frac{X{W}_S}{\tau }\right)\in {ℝ}^{N\times M},& \text{<span class="tml-eqn"></span>}\end{array}$$ where $W_S\in {ℝ}^{d\times M}$ is a learnable projection and $\tau >0$ is a temperature parameter. The column $S_{\cdot ,m}$ gives the soft membership weights of all $N$ nodes to the $m$-th physical state.

Step 2: Aggregate to physical states.

The intrinsic physical state for each slice $m$ is computed by a weighted average: (Transolver Aggregate)$$\begin{array}{ccc}& z_m=\frac{{\sum }_{i=1}^N{S}_{im}{x}_i}{{\sum }_{i=1}^N{S}_{im}}\in {ℝ}^d,m=1,\dots ,M.& \text{<span class="tml-eqn"></span>}\end{array}$$ Stacking these gives $Z\in {ℝ}^{M\times d}$, the physics token matrix.

Step 3: Attention among physical states.

Standard multi-head self-attention is applied to the $M$ physics tokens: (Transolver ATTN)$$\begin{array}{ccc}& Z^{\prime }=\mathrm{MHA}(Z)+Z\in {ℝ}^{M\times d}.& \text{<span class="tml-eqn"></span>}\end{array}$$ Since $M\ll N$, this costs only $𝒪(M^{2}d)$.

Step 4: Distribute back to mesh nodes.

The updated physics tokens are distributed back to all mesh nodes using the same assignment matrix: (Transolver Distribute)$$\begin{array}{ccc}& x_i^{\prime }=\sum _{m=1}^M{S}_{im}{z}_m^{\prime },i=1,\dots ,N.& \text{<span class="tml-eqn"></span>}\end{array}$$ A residual connection and feed-forward layer complete the block: ${\hat{x}}_i=\text{FFN}(x_i^{\prime }+x_i)+x_i^{\prime }+x_i$.

Handling complex geometries

A major strength of Transolver is its ability to handle genuinely complex geometries. The learned slice assignment matrix $S$ automatically adapts to the geometry: near a thin boundary layer, many nodes are assigned to the “boundary layer” physical state; in the quiescent freestream, fewer physical states are needed.

Transolver has been demonstrated on:

• Airfoils (AirfRANS dataset): Steady-state RANS simulations around parametrically varying NACA airfoil shapes. Transolver achieves state-of-the-art relative $L^2$ error on pressure and velocity fields.

• Vehicle aerodynamics (DrivAerNet): 3D external aerodynamics around car bodies, with $\sim {10}^5$ surface mesh nodes. The $𝒪(NM)$ scaling makes this tractable where $𝒪(N^2)$ attention would be prohibitive.

• Elasticity (von Mises stress on 3D parts): Predicting stress distributions on complex mechanical parts with unstructured tetrahedral meshes.

HAMLET: Graph Transformer Neural Operator

While Transolver compresses the mesh into learned physical states, an alternative strategy is to explicitly model the mesh as a graph and use graph Transformers that combine the local inductive bias of graph neural networks (GNNs) with the global attention of Transformers. The HAMLET framework (Bryutkin, Hao, et al., ICML 2024) takes this approach, building a modular architecture that separately encodes PDE-specific information and geometric/mesh information before combining them via a graph Transformer.

Graph structure for PDE domains

Given a mesh $ℳ=(𝒱,ℰ)$ with nodes $𝒱=\{1,\dots ,N\}$ and edges $ℰ$ (connecting neighbouring mesh elements), HAMLET constructs a multi-relational graph that encodes:

1. Spatial adjacency: The original mesh connectivity, providing local structural information.

2. Multi-scale connections: Edges connecting nodes at different scales of a coarsened mesh hierarchy (obtained via graph coarsening algorithms), enabling efficient multi-resolution information flow.

3. Boundary-to-interior edges: Explicit connections from boundary nodes to interior nodes, reflecting the influence of boundary conditions on the solution.

Modular encoder architecture

HAMLET uses a modular encoder that separately processes PDE-specific information and geometric information:

PDE info encoder.

The PDE information at each node $i$ (input function values, boundary condition types and values, source terms) is processed by a dedicated MLP: (Hamlet PDE ENC)$$\begin{array}{ccc}& h_i^{\text{PDE}}={\text{MLP}}_{\text{PDE}}\left([a({𝐱}_i),b_i,f({𝐱}_i)]\right)\in {ℝ}^d,& \text{<span class="tml-eqn"></span>}\end{array}$$ where $a$ is the input function, $b_i$ encodes boundary condition type and value (zero for interior nodes), and $f$ is the source term.

Geometry encoder.

The geometric information (node coordinates, edge vectors, element areas) is processed by a separate encoder: (Hamlet GEOM ENC)$$\begin{array}{ccc}& h_i^{\text{geom}}={\text{MLP}}_{\text{geom}}\left([{𝐱}_i,{\text{LocalGeom}}_i]\right)\in {ℝ}^d,& \text{<span class="tml-eqn"></span>}\end{array}$$ where ${\text{LocalGeom}}_i$ encodes local geometric features such as element areas, angles, and curvatures at node $i$.

Fusion.

The two encodings are fused via element-wise multiplication and addition: (Hamlet Fusion)$$\begin{array}{ccc}& h_i^{(0)}=h_i^{\text{PDE}}\odot h_i^{\text{geom}}+h_i^{\text{PDE}}+h_i^{\text{geom}}.& \text{<span class="tml-eqn"></span>}\end{array}$$

Graph Transformer layers

The fused representations are processed by a stack of graph Transformer layers. Each layer combines:

1. Local message passing: A GNN-style update using the mesh adjacency edges, capturing local interactions: (Hamlet Local)$$\begin{array}{ccc}& m_i^{(\ell )}=\sum _{j\in 𝒩(i)}{\alpha }_{ij}^{(\ell )}{W}_V^{(\ell )}{h}_j^{(\ell )},& \text{<span class="tml-eqn"></span>}\end{array}$$ where $𝒩(i)$ is the neighbourhood of node $i$ in the mesh and ${\alpha }_{ij}^{(\ell )}$ are graph attention coefficients.

2. Global attention: A subset of $M$ “virtual nodes” (analogous to GNOT's latent tokens) aggregates global information via cross-attention from all mesh nodes, and then broadcasts back. This provides the global receptive field that local message passing lacks.

3. Multi-relational aggregation: Messages from different edge types are combined with learnable weights: (Hamlet Multirel)$$\begin{array}{ccc}& h_i^{(\ell +1)}=h_i^{(\ell )}+\sum _{r=1}^R{\beta }_r{m}_{i,r}^{(\ell )},& \text{<span class="tml-eqn"></span>}\end{array}$$ where $m_{i,r}^{(\ell )}$ is the message from relation type $r$.

PiT: Position-induced Transformer

The Position-induced Transformer (PiT), introduced at ICML 2024, takes a different approach to making Transformers efficient for PDEs. Rather than reducing the number of tokens (Transolver) or using a graph structure (HAMLET), PiT exploits the positional structure inherent in PDE discretisations to induce a structured, efficient attention pattern.

Key insight: spatial proximity implies interaction strength

The fundamental observation is that in most PDEs, the interaction strength between two points decays with their spatial distance. For elliptic PDEs, the Green's function typically decays as $|G(𝐱,𝐲)|\sim |𝐱-𝐲{|}^{2-d}$ in $d$ dimensions (logarithmically in 2D). For parabolic PDEs, the heat kernel decays exponentially with distance. Even for hyperbolic PDEs where information propagates along characteristics, the relevant interactions are structured by the geometry of the characteristic surfaces.

Position-induced attention

PiT exploits this structure by decomposing the attention computation into two components:

Positional bias.

A positional bias function $b:{ℝ}^d\times {ℝ}^d\to ℝ$ encodes the spatial relationship between points: (PIT BIAS)$$\begin{array}{ccc}& b({𝐱}_i,{𝐱}_j)={\text{MLP}}_{\text{pos}}({𝐱}_i-{𝐱}_j),& \text{<span class="tml-eqn"></span>}\end{array}$$ where ${\text{MLP}}_{\text{pos}}$ is a small MLP that maps relative positions to a scalar bias. This is added to the attention logits: (PIT ATTN)$$\begin{array}{ccc}& A_{ij}=\frac{q_i^{\top }{k}_j}{\sqrt{d_k}}+b({𝐱}_i,{𝐱}_j).& \text{<span class="tml-eqn"></span>}\end{array}$$

Hierarchical attention.

To reduce the $𝒪(N^2)$ cost of computing all pairwise biases, PiT uses a hierarchical decomposition:

1. Fine level: Full attention within local neighbourhoods of radius $r$ (containing $\sim k$ points each).

2. Coarse level: Cluster the domain into $N/k$ supernodes, compute attention among supernodes, and broadcast back.

The total complexity is $𝒪(Nk+{(N/k)}^2)=𝒪(Nk+N^2/k^2)$, which is minimised at $k=N^{2/3}$ giving $𝒪(N^{5/3})$-sub-quadratic.

Connection to numerical PDE methods

PiT's hierarchical attention has deep connections to classical numerical methods:

• Multigrid methods: PiT's fine-coarse decomposition mirrors the V-cycle of multigrid: local smoothing on the fine grid followed by correction on a coarser grid. The key difference is that multigrid restricts and prolongs the solution, while PiT restricts and prolongs the attention.

• $\mathscr{H}$-matrices: The hierarchical matrix ($\mathscr{H}$-matrix) framework in numerical linear algebra represents dense matrices arising from integral equations in a compressed, hierarchical format by exploiting the low-rank structure of off-diagonal blocks. PiT's hierarchical attention can be viewed as a learned $\mathscr{H}$-matrix.

• Fast multipole method (FMM): The FMM computes $N$-body interactions in $𝒪(N)$ time by replacing distant interactions with multipole expansions. PiT's supernode aggregation is analogous to the FMM's multipole expansion.

PDE-Conditional Transformers and Foundation Models

The architectures discussed so far-Galerkin Transformer, GNOT, Transolver, HAMLET, PiT-are trained on data from a specific PDE or a specific family of PDEs with varying parameters. A more ambitious goal is to build foundation models for PDEs: large, pretrained Transformers that can solve many different types of PDEs, potentially even generalising to PDE families unseen during training. This subsection covers the key ideas behind PDE-conditional Transformers and the emerging landscape of PDE foundation models.

Unisolver: conditioning on PDE structure

The Unisolver observes that different PDEs have different mathematical structures (different differential operators, different types of nonlinearities), and that a universal solver must be conditioned on these structural properties. The key innovation is a taxonomy of PDE conditioning strategies:

Unisolver uses a combination of both conditioning types. The PDE type is encoded as a domain-wise embedding, while spatially varying parameters (coefficients, source terms, boundary conditions) are encoded point-wise. This allows a single model to handle a heterogeneous dataset of PDE problems.

Zero-shot inference across PDE families

A remarkable finding is that Unisolver, trained on a diverse dataset of PDE families (heat equations, wave equations, advection-diffusion, reaction-diffusion, Burgers, incompressible Navier–Stokes), can perform zero-shot inference on PDE families not seen during training. For example:

• Trained on heat and advection-diffusion equations, Unisolver achieves non-trivial accuracy on the convection-diffusion equation (which combines features of both).

• Trained on 2D PDEs, Unisolver transfers to 3D versions of the same PDEs by leveraging the coordinate-agnostic nature of the attention mechanism.

This zero-shot capability arises because the point-wise conditioning mechanism allows the model to “read” the local PDE structure from the input, rather than memorising specific PDE types.

Multiple Physics Pretraining (MPP)

The Multiple Physics Pretraining (MPP) framework (McCabe, Blancard, et al., NeurIPS 2024) takes the foundation model idea further by pretraining a large Transformer on a massive, diverse dataset of PDE simulations spanning multiple physics domains.

The key design choices are:

Unified tokenisation.

All PDE solutions are represented as spatial fields on regular or irregular grids. Each spatial patch (e.g., $8\times 8$ region on a 2D grid) is tokenised as a vector, analogous to Vision Transformer (ViT) patches for images. The number of physical channels varies by PDE type (1 for scalar equations, 3 for velocity fields, etc.), and is handled by a per-channel linear projection to a common embedding dimension.

Physics-aware masking.

During pretraining, random patches of the solution field are masked, and the model is trained to reconstruct them-a masked autoencoder (MAE) objective. The masking ratio is high (75%) to force the model to learn meaningful physical correlations rather than interpolation shortcuts.

Scaling laws.

MPP demonstrates that PDE foundation models obey scaling laws analogous to those observed in large language models. As the model size increases from 10M to 1B parameters and the dataset grows from ${10}^4$ to ${10}^6$ PDE simulations, the validation loss decreases as an approximate power law: (MPP Scaling)$$\begin{array}{ccc}& {ℒ}_{\text{val}}\approx C\cdot N_{\text{params}}^{-\alpha }\cdot D^{-\beta },& \text{<span class="tml-eqn"></span>}\end{array}$$ with $\alpha \approx 0.07$ and $\beta \approx 0.12$ (dataset-dependent). This suggests that continued scaling of PDE foundation models will yield predictable improvements.

CoDA-NO: Codomain Attention Neural Operator

The Codomain Attention Neural Operator (CoDA-NO, NeurIPS 2024) introduces a fundamentally different attention mechanism: instead of attending across spatial positions (the domain), CoDA-NO attends across the output channels (the codomain).

Codomain attention is computationally cheap ($𝒪(N{C}^{2}d)$ with $C$ typically small) and captures the inter-channel correlations that are crucial for multi-field PDEs. For example, in Navier–Stokes, the pressure and velocity fields are coupled through the incompressibility constraint $\nabla \cdot 𝐮=0$; codomain attention allows the model to learn these couplings.

CoDA-NO uses codomain attention in combination with spatial attention (either via FNO-style spectral convolutions or Galerkin attention) to capture both spatial and inter-channel dependencies. It is pretrained on a diverse multi-physics dataset and fine-tuned for specific PDE tasks, achieving strong transfer learning performance.

Towards PDE foundation models: a synthesis

The approaches of Unisolver, MPP, and CoDA-NO represent three complementary pillars of PDE foundation models:

1. Unisolver: Conditioning the Transformer on PDE structure (domain-wise + point-wise), enabling a single model to handle diverse PDE types.

2. MPP: Scaling up pretrained Transformers on massive multi-physics datasets, demonstrating predictable scaling laws.

3. CoDA-NO: Introducing codomain attention for inter-channel transfer learning across PDE families.

In-Context Operator Learning

One of the most surprising capabilities of large language models is in-context learning: the ability to perform new tasks at inference time simply by providing a few input-output examples in the prompt, without any gradient updates. Can we achieve the same for PDE solving? The answer is yes, and the resulting framework- In-Context Operator Learning (ICON)-represents a fascinating convergence of language modelling ideas and operator learning.

ICON: learning operators from demonstrations at inference

ICON, introduced by Yang, Liu, and co-authors (PNAS, 2023), trains a Transformer to take as input a sequence of input-output function pairs $\{(a_1,u_1),\dots ,(a_K,u_K)\}$-the “demonstrations”-together with a new query input $a_{K+1}$, and outputs the predicted solution $u_{K+1}$. Crucially, the demonstrations define the operator implicitly: the model must infer what operator maps $a$ to $u$ from the examples alone.

The architecture processes the demonstration sequence by interleaving each input function $a_k$ and output function $u_k$ as tokens, similar to how a language model processes alternating prompt and response tokens. Self-attention over this interleaved sequence allows the model to:

1. Identify the pattern relating inputs $a_k$ to outputs $u_k$ (i.e., infer the operator).

2. Apply the inferred operator to the query input $a_{K+1}$.

PDEformer: PDE as computational graph

While ICON treats the operator as an implicit “black box” to be inferred from demonstrations, the PDEformer (ICLR 2024) takes a more structured approach: it explicitly represents the PDE itself as a computational graph and processes it with a graph Transformer.

The key insight is that any PDE can be represented as a directed acyclic graph (DAG) where:

• Leaf nodes: Input variables ($x$, $t$), parameters ($\nu$, $\kappa$), and the unknown function $u$.

• Internal nodes: Mathematical operations (differentiation $\partial /\partial x$, multiplication, addition, nonlinear functions like $\sin$, $\exp$).

• Root node: The PDE residual $𝒩[u]-f=0$.

LLM-inspired ideas for PDEs

The success of in-context learning for PDEs has inspired a broader programme of importing ideas from large language models:

1. Autoregressive PDE solving: Treating the time axis as a sequence and generating solutions autoregressively, one time step at a time, using causal attention. This mirrors how LLMs generate text token by token.

2. Prompt engineering for PDEs: Providing the model with carefully chosen demonstrations that “steer” the inference towards the desired PDE family or solution regime.

3. Chain-of-thought for PDEs: Instead of directly predicting the final solution, having the model output intermediate quantities (e.g., the pressure field before the velocity field in a Navier–Stokes solve) as “reasoning steps.”

4. Retrieval-augmented PDE solving: Augmenting the model's input with retrieved examples of similar PDEs from a database, analogous to retrieval-augmented generation (RAG) in NLP.

PINNsFormer

All the architectures discussed so far in this section are data-driven: they learn operators from input-output pairs generated by classical solvers. The PINNsFormer (Zhao, Fang, and Karniadakis, 2023) takes a different approach: it replaces the MLP backbone of a Physics-Informed Neural Network (PINN, Section Physics-Informed Neural Networks) with a Transformer, combining the physics-informed loss formulation of PINNs with the representational power of Transformers.

Motivation: limitations of MLP backbones

Recall that a standard PINN parameterises the solution $u_{\varphi }(𝐱,t)$ as an MLP and minimises the PDE residual at collocation points. MLPs suffer from two well-known limitations:

1. Spectral bias: MLPs with smooth activation functions (tanh, sigmoid) learn low-frequency components of the solution first and struggle with high-frequency features (shocks, sharp gradients, oscillations).

2. No temporal structure: The MLP treats time $t$ as just another input coordinate, with no architectural mechanism to capture the causal, sequential nature of time evolution.

PINNsFormer addresses both limitations simultaneously.

Pseudo Sequence Generator

The key challenge in applying a Transformer to PINNs is that collocation points ${\{({𝐱}_i,t_i)\}}_{i=1}^N$ are not naturally sequential. PINNsFormer introduces the Pseudo Sequence Generator (PSG), which converts each collocation point into a short sequence by sampling additional points along the time axis.

The physical intuition is clear: the solution at $(𝐱,t)$ depends on the history of the solution at the same spatial location at earlier times. By providing this temporal context as a sequence, the Transformer can capture causal dependencies through its attention mechanism.

Architecture

The PINNsFormer architecture consists of:

1. Input embedding: Each element of the pseudo-sequence is embedded by a small MLP: (Pinnsformer Embed)$$\begin{array}{ccc}& h_{\ell }^{(0)}={\text{MLP}}_{\text{embed}}(𝐱,t_{\ell })+\text{PE}(t_{\ell })\in {ℝ}^d,& \text{<span class="tml-eqn"></span>}\end{array}$$ where $\text{PE}(t_{\ell })$ is a sinusoidal positional encoding.

2. Transformer encoder: A standard Transformer encoder with $L$ layers of multi-head self-attention and feed-forward networks processes the sequence $[h_0^{(0)},\dots ,h_{L_{\text{seq}}-1}^{(0)}]$. The self-attention uses causal masking (each token can only attend to earlier tokens in the sequence), enforcing the causal structure of time evolution.

3. Output: The hidden state at the last position ($t_{L_{\text{seq}}-1}=t$) is mapped to the predicted solution value: (Pinnsformer Output)$$\begin{array}{ccc}& u_{\varphi }(𝐱,t)={\text{MLP}}_{\text{out}}(h_{L_{\text{seq}}-1}^{(L)}).& \text{<span class="tml-eqn"></span>}\end{array}$$

4. Physics-informed loss: The loss function is exactly the same as for standard PINNs (equation ): PDE residuals are computed via automatic differentiation through the Transformer, and the loss penalises the residuals at collocation points plus boundary and initial condition violations.

Wavelet activation functions

To further combat spectral bias, PINNsFormer introduces wavelet activation functions that replace the standard smooth activations (tanh, GELU) in the feed-forward layers: (Wavelet Activation)$$\begin{array}{ccc}& {\sigma }_{\text{wav}}(x)=w_1{\psi }_1\left(\frac{x-b_1}{a_1}\right)+w_2{\psi }_2\left(\frac{x-b_2}{a_2}\right),& \text{<span class="tml-eqn"></span>}\end{array}$$ where ${\psi }_1$ is the Mexican hat wavelet (negative normalised second derivative of a Gaussian): (Mexican HAT)$$\begin{array}{ccc}& {\psi }_1(x)=\frac{2}{\sqrt{3}{\pi }^{1/4}}(1-x^2)e^{-x^2/2},& \text{<span class="tml-eqn"></span>}\end{array}$$ and ${\psi }_2$ is the Morlet wavelet: (Morlet)$$\begin{array}{ccc}& {\psi }_2(x)=e^{-x^2/2}\cos (5x).& \text{<span class="tml-eqn"></span>}\end{array}$$ The parameters $w_1,w_2,a_1,a_2,b_1,b_2$ are learnable, allowing the model to adaptively select the frequency content of its activations.

Comparative analysis and practical guidance

To help the practitioner navigate the growing zoo of Transformer-based PDE solvers, we provide a comparative summary in Table Table 2.

Method	Attention type	Complexity	Irregular mesh	Multi-PDE
Galerkin Trans.\	Linear (Galerkin)	$𝒪(N{d}^2)$	Partial	No
GNOT	Cross-attn + latent	$𝒪(NM+M^2)$	Yes	Partial
Transolver	Physics-attn	$𝒪(NM)$	Yes	No
HAMLET	Graph + virtual nodes	$𝒪(N|ℰ|+NM)$	Yes	No
PiT	Hierarchical	$𝒪(N^{5/3})$	Yes	No
Unisolver	Conditioned attn	$𝒪(N^2)$ or linear	Yes	Yes
ICON	In-context	$𝒪({(KN)}^2)$	Partial	Yes
PINNsFormer	Causal self-attn	$𝒪(N{L}_{\text{seq}}^2)$	Yes	No

We close this section with a forward-looking exercise.

Diffusion and Score-Based Models for PDE Simulation

“The idea that one can create structure by progressively removing noise is, in retrospect, deeply physical: it is the reverse of the second law of thermodynamics.” - Yang Song

We have now seen two families of neural PDE solvers. Physics-Informed Neural Networks (Section Physics-Informed Neural Networks) embed PDE constraints into the loss function but produce a single deterministic solution per problem instance. Neural operators (Section Neural Operator Learning) learn the mapping from PDE parameters to solutions across entire families, but they too are deterministic: given input parameters $\theta$, they output a single prediction $\hat{u}$.

In this section, we make the leap to generative PDE solvers based on diffusion and score-matching. The core idea is both simple and profound: instead of learning a point estimate $\hat{u}=𝒢(\theta )$, we learn the distribution $p(u|\theta )$ over solution fields conditioned on PDE parameters. Sampling from this distribution produces diverse, physically plausible solution fields-complete with calibrated uncertainty-that respect the governing equations. This is the single most important conceptual shift in this chapter.

Why diffusion models for PDEs?

The mean-regression problem

Consider a turbulent flow governed by the incompressible Navier–Stokes equations at Reynolds number $\mathrm{Re}=\mathrm{10,000}$. Two runs with initial conditions that differ by a single-bit perturbation will diverge exponentially, exploring vastly different regions of the solution manifold. Yet if we train a deterministic neural operator $\hat{u}={𝒢}_{\varphi }(\theta )$ by minimising the mean squared error (MSE) over a training set of $({\theta }_i,u_i)$ pairs, the network learns the conditional expectation: (Condmean)$$\begin{array}{ccc}& {\hat{u}}_{\text{MSE}}(\theta )={𝔼}_{p(u|\theta )}[u]=\int up(u|\theta )\mathrm{d}u.& \text{<span class="tml-eqn"></span>}\end{array}$$ This conditional mean is a weighted average over all possible trajectories. For turbulent flows, it produces an unphysical blurred field: the sharp vortices, thin shear layers, and intermittent structures that characterise real turbulence are washed out.

More concretely, suppose the solution distribution is bimodal: $p(u|\theta )=\frac{1}{2}\delta (u-u_A)+\frac{1}{2}\delta (u-u_B)$, where $u_A$ and $u_B$ are two very different but equally valid turbulent realisations. The conditional mean $\hat{u}=\frac{1}{2}(u_A+u_B)$ resembles neither $u_A$ nor $u_B$. A generative model, by contrast, would sample either $u_A$ or $u_B$ with equal probability, always producing a physically valid field.

Why diffusion, specifically?

Among generative models-GANs, VAEs, normalising flows, autoregressive models, diffusion-why have diffusion models emerged as the dominant paradigm for PDE applications? Several reasons:

1. Iterative refinement parallels multigrid. The reverse diffusion process refines a noisy field over many steps, progressing from coarse, large-scale features to fine details. This coarse-to-fine generation mirrors the multigrid hierarchy used in classical PDE solvers, where one resolves low-frequency error components on coarse grids before refining on fine grids.

2. Stable training. Unlike GANs, diffusion models are trained with a simple denoising objective that does not require adversarial min-max optimisation. Training is stable, does not suffer from mode collapse, and scales well to high-dimensional fields.

3. Plug-and-play conditioning. Diffusion models support posterior sampling via guidance: one can inject physics constraints (PDE residuals, conservation laws, boundary conditions) and observational data (sparse sensors, partial measurements) into the sampling process without retraining the model.

4. Calibrated uncertainty. The stochastic nature of the sampling process naturally yields an ensemble of solutions, whose spread provides a measure of uncertainty.

5. Forward diffusion = heat equation. The forward noising process is itself a PDE-the heat equation in the space of solution fields! This deep connection between diffusion models and PDEs provides both intuition and mathematical tools.

The physical analogy: forward diffusion as the heat equation

Let us make the connection between diffusion models and PDEs precise. Recall that the forward process of a variance-preserving (VP) diffusion model adds Gaussian noise to a data sample $x_0$ according to the stochastic differential equation (SDE): (Forward SDE)$$\begin{array}{ccc}& \mathrm{d}x=-\frac{1}{2}\beta (t)x\mathrm{d}t+\sqrt{\beta (t)}\mathrm{d}{W}_t,t\in [0,T],& \text{<span class="tml-eqn"></span>}\end{array}$$ where $\beta (t)>0$ is the noise schedule and $W_t$ is a Wiener process. The marginal distribution $p_t(x)$ satisfies the Fokker–Planck equation: (Fokkerplanck)$$\begin{array}{ccc}& \frac{\partial p_t}{\partial t}=\frac{1}{2}\beta (t)\nabla \cdot \left(x{p}_t\right)+\frac{\beta (t)}{2}\mathrm{\Delta }{p}_t.& \text{<span class="tml-eqn"></span>}\end{array}$$ When $\beta$ is constant and we drop the drift term, this reduces to the heat equation ${\partial }_{t}p=\frac{\beta }{2}\mathrm{\Delta }p$. The forward diffusion literally heats up the data distribution until it becomes featureless Gaussian noise-thermal equilibrium.

The reverse process, which generates data by starting from noise, can then be understood as anti-diffusion: the creation of structure from chaos, running the heat equation backwards in time. Of course, the heat equation is irreversible in general (information is lost), so the reverse process requires knowledge of the score function ${\nabla }_x\log p_t(x)$, which encodes the structure that must be restored at each time step.

Score-based diffusion on function spaces

PDE solutions live not in ${ℝ}^n$ but in infinite-dimensional function spaces: the velocity field of a fluid is an element of a Sobolev space $H^s(\mathrm{\Omega };{ℝ}^d)$, the temperature field lives in $L^2(\mathrm{\Omega })$, and so on. A truly satisfying theory of diffusion models for PDEs must therefore operate directly on function spaces, independent of any particular discretisation (finite element mesh, spectral coefficients, point cloud).

In this subsection, we develop the mathematical framework for score-based diffusion in infinite dimensions, following the landmark works of Pidstrigach (2022), Kerrigan et al. (2023), and Lim et al. (JMLR, 2025).

The finite-dimensional recap

We begin with a rapid review of the score-based framework in ${ℝ}^n$. Let $p_0(x)$ denote the data distribution (e.g., a distribution over discretised solution fields $x\in {ℝ}^n$). The forward noising process is an Itô SDE: (Forward General)$$\begin{array}{ccc}& \mathrm{d}{x}_t=f(x_t,t)\mathrm{d}t+g(t)\mathrm{d}{W}_t,x_0\sim p_0,& \text{<span class="tml-eqn"></span>}\end{array}$$ where $f:{ℝ}^n\times [0,T]\to {ℝ}^n$ is the drift, $g:[0,T]\to ℝ$ is the diffusion coefficient, and $W_t$ is an $n$-dimensional Wiener process. The marginal density $p_t(x)$ evolves according to the Fokker–Planck equation. At the terminal time $T$, we want $p_T\approx \pi$, a tractable reference distribution (typically $𝒩(0,I_n)$).

The fundamental result of Anderson (1982) states that the forward SDE has a time-reversed counterpart: (Reverse SDE)$$\begin{array}{ccc}& \mathrm{d}{x}_t=[f(x_t,t)-g(t{)}^2{\nabla }_x\log p_t(x_t)]\mathrm{d}t+g(t)\mathrm{d}{\bar{W}}_t,& \text{<span class="tml-eqn"></span>}\end{array}$$ where ${\bar{W}}_t$ is a Wiener process running backwards in time (from $t=T$ to $t=0$), and the crucial extra term $-g{(t)}^2{\nabla }_x\log p_t(x_t)$ is precisely the score function scaled by the diffusion coefficient.

In practice, we approximate the score with a neural network $s_{\varphi }(x,t)\approx {\nabla }_x\log p_t(x)$ trained via denoising score matching: (DSM)$$\begin{array}{ccc}& {ℒ}_{\text{DSM}}(\varphi )={𝔼}_{t\sim 𝒰[0,T]}{𝔼}_{x_0\sim p_0}{𝔼}_{x_t\sim p_{t|0}(x_t|x_0)}[\lambda (t)\Vert s_{\varphi }(x_t,t)-{\nabla }_{x_t}\log p_{t|0}(x_t|x_0){\Vert }^2],& \text{<span class="tml-eqn"></span>}\end{array}$$ where $p_{t|0}(x_t|x_0)$ is the tractable transition kernel of the forward SDE (Gaussian for linear drift), and $\lambda (t)$ is a positive weighting function.

Extension to Hilbert spaces

Now we move to infinite dimensions. Let $\mathscr{H}$ be a separable Hilbert space (for concreteness, think of $\mathscr{H}=L^2(\mathrm{\Omega })$ where $\mathrm{\Omega }\subset {ℝ}^d$ is the physical domain). We wish to define a diffusion process on $\mathscr{H}$ that perturbs function-valued data $u_0\in \mathscr{H}$.

The first obstacle is the definition of a Wiener process. In finite dimensions, $W_t$ is a standard Brownian motion with identity covariance. In infinite dimensions, the analogous object-a $\mathscr{H}$-valued Brownian motion with identity covariance-does not exist as a bona fide $\mathscr{H}$-valued process (its sample paths would have infinite norm almost surely). We must use a $Q$-Wiener process with a trace-class covariance operator $Q$.

The covariance operator $Q$ controls the spatial correlation structure of the noise. A common choice is $Q={(-\mathrm{\Delta }+{\tau }^{2}I)}^{-\alpha }$ for some $\alpha >d/2$ and $\tau >0$, yielding a Matérn-like spatial covariance that decays with distance and penalises high-frequency noise.

With the $Q$-Wiener process in hand, we can define the forward SDE on $\mathscr{H}$: (Forward Hilbert)$$\begin{array}{ccc}& \mathrm{d}{u}_t=f(u_t,t)\mathrm{d}t+g(t)\mathrm{d}{W}_t^Q,u_0\sim {\mu }_0,& \text{<span class="tml-eqn"></span>}\end{array}$$ where ${\mu }_0$ is a probability measure on $\mathscr{H}$ (the data distribution over PDE solution fields), $f:\mathscr{H}\times [0,T]\to \mathscr{H}$ is a (possibly nonlinear) drift operator, and $g:[0,T]\to ℝ$ is a scalar diffusion coefficient.

Score matching in function space

How do we train a score network $s_{\varphi }(u,t)\approx s_t(u)$ when $u$ lives in a function space? The denoising score matching objective generalises naturally. For the Ornstein–Uhlenbeck forward process $\mathrm{d}{u}_t=-\frac{1}{2}\beta (t)u_t\mathrm{d}t+\sqrt{\beta (t)}\mathrm{d}{W}_t^Q$, the transition kernel is Gaussian on $\mathscr{H}$: (Transition Hilbert)$$\begin{array}{ccc}& u_t|u_0\sim 𝒩(\alpha (t)u_0,\sigma (t{)}^{2}Q),& \text{<span class="tml-eqn"></span>}\end{array}$$ where $\alpha (t)=\exp (-\frac{1}{2}{\int }_0^t\beta (s)\mathrm{d}s)$ and $\sigma {(t)}^2=1-\alpha {(t)}^2$.

The conditional score in the eigenbasis of $Q$ is: (COND Score Hilbert)$$\begin{array}{ccc}& {⟨{\nabla }_u\log p_{t|0}(u_t|u_0),e_k⟩}_{\mathscr{H}}=-\frac{{⟨u_t-\alpha (t)u_0,e_k⟩}_{\mathscr{H}}}{\sigma {(t)}^2{\lambda }_k}.& \text{<span class="tml-eqn"></span>}\end{array}$$ The function-space denoising score matching loss is: (DSM Hilbert)$$\begin{array}{ccc}& {ℒ}_{\text{DSM}}^{\mathscr{H}}(\varphi )={𝔼}_{t,u_0,u_t}[\lambda (t)\Vert s_{\varphi }(u_t,t)-{\nabla }_u\log p_{t|0}(u_t|u_0){\Vert }_{\mathscr{H}}^2].& \text{<span class="tml-eqn"></span>}\end{array}$$

This discretisation agnosticism is a major practical advantage: one can train on coarse grids (cheap) and generate on fine grids (expensive but needed only at inference time).

Mathematical rigor: well-posedness

We collect the key conditions under which the infinite-dimensional framework is well-posed.

DiffusionPDE: joint diffusion over coefficients and solutions

One of the most striking applications of diffusion models to PDEs is DiffusionPDE (Huang et al., NeurIPS 2024), which proposes learning the joint distribution $p(a,u)$ over PDE coefficients $a$ (e.g., a spatially varying diffusivity field) and solution fields $u$, simultaneously. This is a conceptually radical idea: rather than viewing the PDE as a fixed constraint that maps $a\mapsto u$, we treat both $a$ and $u$ as random fields drawn from a joint distribution that encodes the PDE relationship implicitly.

Problem formulation

Consider a parametric PDE of the form: (Parametric PDE)$$\begin{array}{ccc}& 𝒩[u;a](𝐱)=0,𝐱\in \mathrm{\Omega },& \text{<span class="tml-eqn"></span>}\end{array}$$ where $a:\mathrm{\Omega }\to ℝ$ is a spatially varying coefficient (e.g., permeability in Darcy flow, thermal conductivity in heat transfer) and $u:\mathrm{\Omega }\to ℝ$ is the solution. Given a distribution $p(a)$ over coefficient fields and the deterministic map $a\mapsto u(a)$ defined by the PDE, there is a joint distribution $p(a,u)$ supported on the manifold $\{(a,u):𝒩[u;a]=0\}$.

DiffusionPDE learns this joint distribution using a single diffusion model. The data is represented as a multi-channel field: (Multichannel)$$\begin{array}{ccc}& z(𝐱)=(a(𝐱),u(𝐱))\in {ℝ}^2,& \text{<span class="tml-eqn"></span>}\end{array}$$ and the diffusion model is trained on paired samples $(a_i,u_i)$ where $u_i$ is obtained by solving the PDE with coefficient field $a_i$.

Sparse observation guidance

At inference time, DiffusionPDE can reconstruct full coefficient and solution fields from sparse observations. Suppose we observe the solution $u$ at a sparse set of sensor locations $𝒮=\{{𝐱}_1,\dots ,{𝐱}_m\}$ with values $y_j=u({𝐱}_j)+{\eta }_j$ (where ${\eta }_j$ is measurement noise). We wish to sample from the posterior: (Posterior)$$\begin{array}{ccc}& p(a,u|y_{𝒮})\propto p(a,u)p(y_{𝒮}|u).& \text{<span class="tml-eqn"></span>}\end{array}$$

This posterior sampling is achieved via classifier guidance applied during the reverse diffusion process. At each denoising step $t$, the unconditional score ${\nabla }_z\log p_t(z_t)$ is augmented with a likelihood gradient: (Guided Score)$$\begin{array}{ccc}& {\tilde{s}}_t(z_t)=\underset{\text{unconditional score}}{\underbrace{{\nabla }_{z_t}\log p_t(z_t)}}+\underset{\text{observation guidance}}{\underbrace{w_{\text{obs}}{\nabla }_{z_t}\log p(y_{𝒮}|z_t)}},& \text{<span class="tml-eqn"></span>}\end{array}$$ where $w_{\text{obs}}$ is a guidance weight. For Gaussian observation noise with variance ${\sigma }_{\text{obs}}^2$: (OBS Likelihood)$$\begin{array}{ccc}& \log p(y_{𝒮}|z_t)\approx -\frac{1}{2{\sigma }_{\text{obs}}^2}\sum _{j=1}^m|y_j-{\hat{u}}_0({𝐱}_j;z_t,t){|}^2,& \text{<span class="tml-eqn"></span>}\end{array}$$ where ${\hat{u}}_0(\cdot ;z_t,t)$ is Tweedie's estimate of the clean field from the noisy state $z_t$ at diffusion time $t$: ${\hat{u}}_0=𝔼[u_0|u_t=z_t]$.

PDE constraint guidance

A key innovation of DiffusionPDE is using the PDE residual itself as an additional guidance signal. Since the generated fields $(a,u)$ should satisfy the PDE, we define a PDE-residual likelihood: (PDE Likelihood)$$\begin{array}{ccc}& \log p_{\text{PDE}}(z_t)\approx -\frac{w_{\text{PDE}}}{2}\Vert 𝒩[{\hat{u}}_0;{\hat{a}}_0]{\Vert }_{\mathrm{\Omega }}^2,& \text{<span class="tml-eqn"></span>}\end{array}$$ where ${\hat{a}}_0$ and ${\hat{u}}_0$ are Tweedie estimates of the clean coefficient and solution fields. The total guided score becomes: (Total Guidance)$$\begin{array}{ccc}& {\tilde{s}}_t(z_t)={\nabla }_{z_t}\log p_t(z_t)+w_{\text{obs}}{\nabla }_{z_t}\log p(y_{𝒮}|z_t)+w_{\text{PDE}}{\nabla }_{z_t}\log p_{\text{PDE}}(z_t).& \text{<span class="tml-eqn"></span>}\end{array}$$

Architecture and training

DiffusionPDE uses a U-Net backbone with the following design choices:

• Multi-channel input. The coefficient field $a$ and solution field $u$ are concatenated as channels of a single spatial field. For systems with multiple solution components (e.g., velocity $𝐯$ and pressure $p$ in Navier–Stokes), all components are stacked.

• Time embedding. The diffusion time $t$ is encoded via sinusoidal positional embeddings and injected into each residual block, following the standard DDPM architecture.

• Training data. Paired samples $(a_i,u_i)$ are generated by sampling $a_i$ from a prior (e.g., a Gaussian random field) and solving the PDE numerically to obtain $u_i$.

• Loss. The standard $ϵ$-prediction loss: $ℒ(\varphi )={𝔼}_{t,z_0,ϵ}[\Vert {ϵ}_{\varphi }(z_t,t)-ϵ{\Vert }^2]$, where $z_t=\alpha (t)z_0+\sigma (t)ϵ$ and $ϵ\sim 𝒩(0,I)$.

Physics-informed diffusion models

While DiffusionPDE applies physics constraints at inference time (via guidance), an alternative approach embeds physics directly into the training objective. This is the strategy of Physics-Informed Diffusion Models (PIDMs), introduced by Bastek and Kochmann (ICLR 2025).

Motivation: PDE residuals as virtual observables

Recall the PINN philosophy: the PDE residual $𝒩[u](𝐱)=0$ is treated as a constraint that must hold at every point in the domain. PIDMs import this idea into the diffusion framework. At each training step, the denoised estimate ${\hat{u}}_0$ (obtained from the current noise prediction ${ϵ}_{\varphi }(u_t,t)$) should approximately satisfy the PDE. We can therefore add a physics loss to the denoising objective: (PIDM LOSS)$$\begin{array}{ccc}& {ℒ}_{\text{PIDM}}(\varphi )=\underset{\text{denoising loss}}{\underbrace{{ℒ}_{\text{DSM}}(\varphi )}}+\underset{\text{physics loss}}{\underbrace{{\lambda }_{\text{phys}}{𝔼}_{t,u_0,u_t}[\Vert 𝒩[{\hat{u}}_0(u_t,t;\varphi )]{\Vert }_{\mathrm{\Omega }}^2]}},& \text{<span class="tml-eqn"></span>}\end{array}$$ where ${\hat{u}}_0(u_t,t;\varphi )=\frac{1}{\alpha (t)}(u_t-\sigma (t){ϵ}_{\varphi }(u_t,t))$ is the Tweedie estimate of the clean sample.

Computing PDE residuals through the denoiser

A subtlety arises: to evaluate $𝒩[{\hat{u}}_0]$, we need spatial derivatives of ${\hat{u}}_0$. Since ${\hat{u}}_0$ depends on $u_t$ (a discretised field on a grid), these derivatives must be computed via finite differences on the grid or via spectral methods. The gradient of the physics loss with respect to $\varphi$ flows through the denoiser: (PIDM GRAD)$$\begin{array}{ccc}& {\nabla }_{\varphi }\Vert 𝒩[{\hat{u}}_0]{\Vert }^2=2⟨𝒩[{\hat{u}}_0],{𝒩}^{\prime }[{\hat{u}}_0]\frac{\partial {\hat{u}}_0}{\partial \varphi }{⟩}_{\mathrm{\Omega }},& \text{<span class="tml-eqn"></span>}\end{array}$$ where ${𝒩}^{\prime }[{\hat{u}}_0]$ is the linearisation (Fréchet derivative) of the PDE operator and $\frac{\partial {\hat{u}}_0}{\partial \varphi }$ comes from backpropagating through the noise-prediction network. For PDEs involving second-order spatial derivatives, one must compose spatial differentiation operators with the neural network Jacobian, which can be expensive but is readily handled by modern automatic differentiation frameworks.

Time-dependent physics weighting

The balance between the denoising loss and the physics loss is controlled by ${\lambda }_{\text{phys}}(t)$. Bastek and Kochmann (2025) propose a schedule that increases the physics weight at low noise levels (small $t$), where the denoised estimate ${\hat{u}}_0$ is close to a clean field and the PDE residual is meaningful, and decreases it at high noise levels (large $t$), where ${\hat{u}}_0$ is still heavily corrupted and enforcing the PDE is premature: (PIDM Schedule)$$\begin{array}{ccc}& {\lambda }_{\text{phys}}(t)={\lambda }_0(1-\sigma (t)/\sigma (T){)}^{\gamma },& \text{<span class="tml-eqn"></span>}\end{array}$$ where ${\lambda }_0>0$ is the base weight and $\gamma >0$ controls the steepness of the annealing.

Theoretical justification

The PIDM objective can be justified from a Bayesian perspective. Consider the generative model as defining a prior $p_{\varphi }(u)$ over solution fields. The PDE constraint $𝒩[u]=0$ defines a likelihood $p_{\text{PDE}}(u)\propto \exp (-\frac{1}{2{\sigma }_{\text{PDE}}^2}\Vert 𝒩[u]{\Vert }^2)$, where ${\sigma }_{\text{PDE}}$ reflects our tolerance for residual violations. The PIDM objective approximately maximises the marginal likelihood: (PIDM Bayes)$$\begin{array}{ccc}& \log p_{\text{phys}}=\log \int p_{\varphi }(u)p_{\text{PDE}}(u)\mathrm{d}u,& \text{<span class="tml-eqn"></span>}\end{array}$$ with respect to $\varphi$. The physics loss acts as a regulariser that biases the learned distribution toward PDE-satisfying fields.

Conditional diffusion for forecasting and data assimilation

Many PDE applications require not unconditioned generation but conditional generation: given some partial information (initial conditions, sparse sensor data, boundary measurements, low-resolution fields), produce a high-fidelity solution field. In this subsection, we develop the theory and practice of conditional diffusion models for PDE forecasting and data assimilation.

Conditioning strategies

There are two broad strategies for conditioning a diffusion model:

1. Training-time conditioning (“conditional model”). The score network $s_{\varphi }(u_t,t,c)$ takes the conditioning information $c$ (e.g., initial condition, boundary data, PDE parameters) as an additional input. The model directly learns the conditional score ${\nabla }_u\log p_t(u|c)$.

2. Inference-time conditioning (“post-hoc guidance”). The score network $s_{\varphi }(u_t,t)$ is trained unconditionally. At sampling time, Bayes' rule is applied: (Bayes Guidance)$$\begin{array}{ccc}& {\nabla }_u\log p_t(u|c)=\underset{\text{unconditional score}}{\underbrace{{\nabla }_u\log p_t(u)}}+\underset{\text{likelihood score}}{\underbrace{{\nabla }_u\log p_t(c|u)}}.& \text{<span class="tml-eqn"></span>}\end{array}$$ The likelihood score is estimated from the conditioning data (e.g., via Tweedie's formula and a forward model).

Each strategy has trade-offs:

Autoregressive diffusion for time-stepping

For time-dependent PDEs, a natural strategy is autoregressive diffusion: train a conditional diffusion model $p_{\varphi }(u^{(n+1)}|u^{(n)})$ that generates the next time step $u^{(n+1)}$ given the current state $u^{(n)}$, then roll out iteratively: (Autoregressive)$$\begin{array}{ccc}& u^{(0)}\stackrel{\stackrel{}{\text{sample}}}{\to }{u}^{(1)}\stackrel{\stackrel{}{\text{sample}}}{\to }{u}^{(2)}\stackrel{\stackrel{}{\text{sample}}}{\to }\cdots & \text{<span class="tml-eqn"></span>}\end{array}$$

The score network takes as input the noised target $u_t^{(n+1)}$ concatenated with the (clean) conditioning state $u^{(n)}$: (Autoregressive Score)$$\begin{array}{ccc}& s_{\varphi }(u_t^{(n+1)},t,u^{(n)})\approx {\nabla }_{u_t^{(n+1)}}\log p_t(u^{(n+1)}|u^{(n)}).& \text{<span class="tml-eqn"></span>}\end{array}$$

Data assimilation with diffusion models

Data assimilation is the process of combining a dynamical model (typically a PDE) with sparse, noisy observations to estimate the full state of a system. Classical methods (Kalman filters, ensemble Kalman filters, 4D-Var) require either linearity or tangent-linear adjoints of the forward model. Diffusion models offer a fundamentally different approach: learn the distribution over states, then condition on observations during sampling.

The data assimilation problem can be formalised as posterior sampling: (DA Posterior)$$\begin{array}{ccc}& p(u|y_{1:K})\propto p(u)\prod _{k=1}^{K}p(y_k|{\mathscr{H}}_k[u]),& \text{<span class="tml-eqn"></span>}\end{array}$$ where $u$ is the full spatiotemporal state, $y_k$ is the $k$-th observation, and ${\mathscr{H}}_k$ is the observation operator (e.g., spatial interpolation to sensor locations). The prior $p(u)$ is defined by the trained diffusion model, and the likelihood terms $p(y_k|{\mathscr{H}}_k[u])$ are applied via guidance.

Applications to weather and climate

Diffusion models have recently shown remarkable promise in weather and climate applications:

• GenCast (DeepMind, 2024): A diffusion model for medium-range weather forecasting that generates ensemble forecasts at 0.25${}^{\circ }$ resolution. GenCast matches or exceeds the European Centre for Medium-Range Weather Forecasts (ECMWF) operational ensemble on many skill scores, while producing forecasts in seconds rather than hours.

• CorrDiff (NVIDIA, 2024): Corrective diffusion for downscaling coarse weather model outputs to high resolution. Uses a conditional diffusion model where the condition is the low-resolution forecast and the target is the high-resolution analysis field.

• SeedsDiff (2024): Seed-based diffusion for ensemble generation in numerical weather prediction, using diffusion to perturb initial conditions and generate diverse ensemble members.

These applications demonstrate that diffusion models can serve as powerful emulators of the atmospheric PDE system (the primitive equations), with the added benefit of natively producing calibrated ensemble forecasts.

CoCoGen: physically consistent score-based generation

CoCoGen (Constraint-Conditioned Generation) addresses a fundamental question: how can we ensure that samples from a diffusion model exactly (or very nearly) satisfy known physical constraints-conservation laws, divergence-free conditions, constitutive relations-rather than merely being “close” in an average-error sense?

The consistency problem

Even well-trained diffusion models can produce samples that violate physical constraints. For example, a diffusion model for incompressible fluid flow might generate velocity fields $𝐯$ that are not exactly divergence-free: $\nabla \cdot 𝐯\ne 0$. While the violations may be small in a statistical sense, they can be physically catastrophic-mass is not conserved, which leads to unphysical pressure fields and ultimately unstable dynamics if the generated fields are used as initial conditions for further simulation.

Incorporating discretised PDE information

CoCoGen incorporates physical constraints into the score model in two complementary ways:

1. Constraint-aware architecture. The score network is designed so that its output automatically satisfies certain constraints. For example, for divergence-free fields, the score network outputs a stream function $\psi$ and the velocity field is computed as $𝐯=\nabla \times \psi =(-{\partial }_y\psi ,{\partial }_x\psi )$, which is divergence-free by construction (since $\nabla \cdot (\nabla \times \psi )=0$).

2. Constraint projection during sampling. At each denoising step, the generated field is projected onto the constraint manifold. For a linear constraint $𝐂u=𝐛$, this is an orthogonal projection: (Cocogen PROJ)$$\begin{array}{ccc}& u_{\text{proj}}=u-{𝐂}^{†}(𝐂u-𝐛),& \text{<span class="tml-eqn"></span>}\end{array}$$ where ${𝐂}^{†}$ is the Moore–Penrose pseudoinverse. For nonlinear constraints, iterative methods (Newton's method, Gauss–Newton) are applied.

Forward and inverse PDE tasks

CoCoGen can handle both forward and inverse PDE problems:

• Forward problem. Given PDE parameters $\theta$ (initial/boundary conditions, coefficients), sample from $p(u|\theta )$ using the PDE as a constraint: $𝒞[u]=𝒩[u;\theta ]=0$. The diffusion model provides a data-driven prior over solution fields, and the constraint projection ensures PDE satisfaction.

• Inverse problem. Given sparse observations of the solution $u$, infer the PDE parameters $\theta$. Here the constraint is that the observed values match: $𝒞[\theta ]=\mathscr{H}[u(\theta )]-y=0$, where $u(\theta )$ is the PDE solution for parameters $\theta$.

Physics-informed diffusion for flow reconstruction

One of the most compelling applications of diffusion models in computational physics is the reconstruction of high-fidelity turbulent flow fields from low-fidelity or sparse data. Traditional approaches (interpolation, RANS-based augmentation) fail to capture the small-scale structures that dominate turbulent transport. Diffusion models, by learning the distribution of DNS (Direct Numerical Simulation) fields, can hallucinate physically plausible fine-scale structures that are consistent with the observed large-scale data.

Problem setup

The setup is a conditional generation problem. Let $u_{\text{HR}}\in {ℝ}^{N_{\text{HR}}\times d}$ denote a high-resolution (HR) flow field (e.g., DNS of turbulence on a ${512}^3$ grid) and $u_{\text{LR}}\in {ℝ}^{N_{\text{LR}}\times d}$ a low-resolution (LR) or sparse observation (e.g., a coarsened ${32}^3$ grid, or scattered PIV measurements). We wish to sample from: (SR Posterior)$$\begin{array}{ccc}& p(u_{\text{HR}}|u_{\text{LR}}).& \text{<span class="tml-eqn"></span>}\end{array}$$

This is exactly the super-resolution problem, but in the context of PDE solutions rather than natural images. The key difference is that PDE solutions must obey physical constraints (momentum conservation, energy spectra, boundary conditions), so the “hallucinated” fine-scale details must be physically consistent.

Architecture: conditional U-Net with physics loss

Shu et al. (JCP, 2023) propose a conditional diffusion model for turbulent flow reconstruction with the following architecture:

1. Conditioning. The LR field $u_{\text{LR}}$ is upsampled to the HR grid (e.g., via bicubic interpolation) and concatenated with the noised HR field $u_t$ as input channels to the U-Net.

2. Denoising backbone. A standard U-Net with residual blocks, group normalisation, and sinusoidal time embedding.

3. Physics-informed loss. In addition to the standard denoising loss, a spectral energy constraint is applied: (Spectral LOSS)$$\begin{array}{ccc}& {ℒ}_{\text{spec}}(\varphi )=𝔼[\Vert E({\hat{u}}_0;\kappa )-E_{\text{DNS}}(\kappa ){\Vert }^2],& \text{<span class="tml-eqn"></span>}\end{array}$$ where $E(\cdot ;\kappa )$ is the energy spectrum as a function of wavenumber $\kappa$, and $E_{\text{DNS}}(\kappa )$ is the target spectrum from DNS data.

Results: JCP 2023

Key results from Shu et al. (JCP, 2023) on 2D and 3D turbulent flows:

• 4$\times$ super-resolution of 2D turbulence: Starting from a ${64}^2$ LR field, the diffusion model generates ${256}^2$ HR fields that match the DNS energy spectrum up to the Nyquist wavenumber of the HR grid. By contrast, bicubic interpolation misses the inertial range entirely, and a deterministic U-Net (trained with MSE) produces overly smooth fields.

• Point-wise error reduction: The diffusion model achieves 40–55% lower point-wise $L^2$ error compared to deterministic baselines on held-out DNS snapshots.

• Structural similarity: Vorticity PDFs, structure functions, and two-point correlations of diffusion-generated fields closely match DNS statistics, while deterministic baselines underestimate the tails of the vorticity PDF (i.e., miss extreme events).

• Uncertainty quantification: By drawing multiple samples from the conditional distribution $p(u_{\text{HR}}|u_{\text{LR}})$, the model provides calibrated uncertainty estimates: regions of high uncertainty correspond to small-scale structures that are poorly constrained by the LR data.

Flow matching and fast sampling for PDE solvers

A major practical limitation of diffusion models is the large number of denoising steps required for high-quality generation (typically hundreds to thousands of neural function evaluations, or NFEs). For PDE solvers that must be invoked repeatedly (e.g., in optimisation loops, real-time control, or ensemble forecasting), this cost can be prohibitive. Flow matching and related techniques offer a path to few-step or even one-step generation.

From diffusion to flow matching

Flow matching (Lipman et al., 2023; Albergo and Vanden-Eijnden, 2023; Liu et al., 2023) replaces the stochastic reverse-time SDE with a deterministic ODE-a continuous normalising flow (CNF). The idea is to learn a time-dependent velocity field $v_{\varphi }(u,t)$ such that the ODE: (CNF)$$\begin{array}{ccc}& \frac{\mathrm{d}{u}_t}{\mathrm{d}t}=v_{\varphi }(u_t,t),u_0\sim \pi =𝒩(0,I),t\in [\mathrm{0,1}],& \text{<span class="tml-eqn"></span>}\end{array}$$ transports the base distribution $\pi$ to the data distribution $p_{\text{data}}$. The key insight is that one can construct a conditional vector field $v_t(u|u_1)=(u_1-(1-{\sigma }_{\min })u)/(1-t)$ for each data sample $u_1$ that linearly interpolates between noise $u_0$ and data $u_1$, and train $v_{\varphi }$ to match the marginal vector field $v_t(u)=𝔼[v_t(u|u_1)|u_t=u]$.

Rectified Flow-Based Operator (RFO)

The Rectified Flow-Based Operator (RFO) proposed by Chen and Morozov (ICLR 2025) adapts flow matching specifically for PDE operator learning. The key idea is to learn a flow that transports a conditional noise distribution $\pi (\cdot |\theta )$ to the solution distribution $p(u|\theta )$, where $\theta$ represents PDE parameters.

The RFO framework has several advantages for PDE applications:

1. Unified forward/inverse. For the forward problem, $\theta$ is given and we sample $u$. For the inverse problem, we jointly learn flows for both $u$ and $\theta$, or fix $u$ (from observations) and invert for $\theta$.

2. Few-step generation. Rectified flows produce nearly straight transport paths, so the ODE can be solved in as few as 1–5 Euler steps, compared to 50–1000 steps for standard diffusion.

3. Neural operator backbone. $v_{\varphi }$ can be parameterised as a neural operator (e.g., FNO or Transformer), enabling resolution-agnostic generation.

Truncation and iterative refinement

A further acceleration technique is truncated flow matching. Rather than transporting from pure noise ($t=0$) all the way to data ($t=1$), one can start from a rough deterministic prediction $u_{\text{det}}$ (obtained cheaply from a neural operator like FNO) and only transport from $u_{\text{det}}$ to the refined solution. This “warm start” eliminates the need to traverse the full noise-to-data path.

Formally, define a truncated interpolant: (Truncated)$$\begin{array}{ccc}& u_t=(1-t)u_{\text{det}}+t{u}_{\text{GT}}+\sigma (t)ϵ,t\in [\mathrm{0,1}],& \text{<span class="tml-eqn"></span>}\end{array}$$ where $u_{\text{GT}}$ is the ground truth and $\sigma (t)={\sigma }_0(1-t)$ is a small noise schedule that vanishes at $t=1$. The flow matching model learns to refine $u_{\text{det}}$ toward $u_{\text{GT}}$, requiring only 1–3 ODE steps for high accuracy.

Physics-guided distillation

Another route to few-step generation is distillation: training a student model to reproduce the multi-step output of a teacher diffusion model in a single forward pass. Phys-Instruct combines distillation with physics constraints:

1. Teacher: A pre-trained diffusion model for PDE solutions, using 1000-step DDPM sampling.

2. Student: A lightweight network $G_{\psi }(z,\theta )$ that maps noise $z$ and PDE parameters $\theta$ to a solution in a single forward pass.

3. Distillation loss: The student is trained to match the teacher's output distribution: (Distill)$$\begin{array}{ccc}& {ℒ}_{\text{distill}}(\psi )={𝔼}_{z,\theta }[\Vert G_{\psi }(z,\theta )-u_{\text{teacher}}(z,\theta ){\Vert }^2].& \text{<span class="tml-eqn"></span>}\end{array}$$

4. Physics loss: An additional PDE residual penalty: (Distill PHYS)$$\begin{array}{ccc}& {ℒ}_{\text{phys}}(\psi )={𝔼}_{z,\theta }[\Vert 𝒩[G_{\psi }(z,\theta );\theta ]{\Vert }^2].& \text{<span class="tml-eqn"></span>}\end{array}$$

The physics loss ensures that the student does not just copy the teacher's samples but produces solutions that actually satisfy the PDE-a form of “physics-grounded distillation” that can sometimes outperform the teacher in terms of PDE residual.

Latent neural PDE solver with flow matching

A promising recent direction combines flow matching with latent space representations. Instead of operating directly on the high-dimensional field $u\in {ℝ}^{N\times d}$, one first trains an autoencoder to map $u$ to a low-dimensional latent code $z=\text{Enc}(u)\in {ℝ}^m$ with $m\ll N$, and then applies flow matching in the latent space: (Latent FM)$$\begin{array}{ccc}& \frac{\mathrm{d}{z}_t}{\mathrm{d}t}=v_{\varphi }(z_t,t,\theta ),z_0\sim 𝒩(0,I_m),& \text{<span class="tml-eqn"></span>}\end{array}$$ followed by decoding: $\hat{u}=\text{Dec}(z_1)$.

Mathematical deep dive: convergence of diffusion PDE solvers

We now present several rigorous results on the convergence properties of diffusion-based PDE solvers.

Approximation theory for score networks

This lemma shows that the quality of generation degrades gracefully with score estimation error, and the degradation is controlled by the integral of $g{(t)}^2$-the total amount of noise injected during the forward process.

PDE residual bounds for diffusion samples

Summary and comparative analysis

We conclude this section with a comparative overview of the methods discussed.

Exercises

Poisson and Laplace-Type Elliptic PDEs

We begin our tour of PDE families with the simplest and, in many ways, the most fundamental: elliptic equations. These are the steady-state workhorses of mathematical physics-they describe equilibrium configurations where nothing changes in time, yet the spatial structure can be extraordinarily rich. If you want to know the temperature distribution in a metal plate after it has reached thermal equilibrium, or the electrostatic potential in a region with fixed charges, or the pressure field driving groundwater through a porous rock formation, you solve an elliptic PDE\@.

Elliptic problems are the natural starting point for benchmarking neural PDE solvers because they lack the additional complications of time stepping, shock formation, and chaotic sensitivity. Get an elliptic solver right, and you have a solid foundation for the harder problems that follow.

The Poisson and Laplace equations

The Laplace equation defines harmonic functions: functions whose value at any point equals the average over any surrounding sphere. This mean-value property implies remarkable smoothness-harmonic functions are infinitely differentiable, even analytic. The Poisson equation adds a source term that “pushes” the solution away from being harmonic.

This well-posedness result is important for neural PDE solvers: it tells us that the mapping $f\mapsto u$ is a bounded linear operator, and therefore a neural operator has a well-defined, stable target to approximate.

Darcy flow: the parametric elliptic benchmark

The Poisson equation is useful for pedagogy, but most real-world elliptic problems involve variable coefficients. The canonical example-and arguably the single most important benchmark in the neural operator literature-is Darcy flow.

Why has Darcy flow become the benchmark for operator learning? Several reasons conspire:

1. Parametric richness. The permeability field $a(𝐱)$ can be drawn from diverse distributions-smooth Gaussian random fields, piecewise-constant media, fractal geometries, channelised structures. Each draw defines a different PDE instance, creating a natural family of problems.

2. Well-posedness. As long as $a$ is bounded away from zero, the Lax–Milgram theorem guarantees existence and uniqueness, so the operator $a\mapsto u$ is well-defined.

3. Practical importance. Darcy flow arises in groundwater hydrology, petroleum reservoir simulation, carbon sequestration modelling, and filtration engineering.

4. Scalable ground truth. High-resolution finite element solutions on structured grids can be computed efficiently, providing plentiful training data.

Neural operator results on Darcy flow

The Fourier Neural Operator (FNO) of Li et al. (2021) achieved a landmark result on Darcy flow: a single trained network, evaluated in under 10 milliseconds on a GPU, matched the accuracy of a finite element solver operating at the same resolution, while generalising across the entire distribution of permeability fields.

Let us recall the key numbers. On the standard $85\times 85$ Darcy benchmark (permeability fields drawn from a thresholded Gaussian random field with length scale $l$), FNO achieves a relative $L^2$ error of approximately $1.08\%$ when trained on 1000 samples. In comparison:

Diffusion models for elliptic inverse problems

Forward Darcy flow-given $a$, find $u$-is well-posed. The inverse problem-given sparse, noisy observations of $u$, recover $a$-is drastically ill-posed. Many permeability fields can produce similar pressure observations, and the mapping from data to parameters is discontinuous. This is precisely the kind of problem where generative models shine: instead of seeking a single “best” $a$, we sample from the posterior distribution $p(a|𝐲)$ where $𝐲$ collects the observed pressure values.

Traditional approaches to this inverse problem use Markov chain Monte Carlo (MCMC) methods, which are accurate but require thousands of forward solves-each taking seconds to minutes on a fine mesh. The total computational cost can be prohibitive for high-dimensional permeability fields.

Score-based priors for inverse problems

The key insight of diffusion-model approaches is to replace the explicit prior $p(a)$ with a learned score function. A score-based diffusion model is trained on samples from the prior distribution of permeability fields (which can be cheaply generated from a Gaussian process). Once trained, the model provides the score ${\nabla }_a\log p_t(a)$ at all noise levels $t$, enabling posterior sampling via: (Score)$$\begin{array}{ccc}& {\nabla }_a\log p_t(a|𝐲)={\nabla }_a\log p_t(a)+{\nabla }_a\log p_t(𝐲|a).& \text{<span class="tml-eqn"></span>}\end{array}$$ The first term comes from the pretrained diffusion model. The second term-the likelihood score-can be approximated using a differentiable forward solver or a neural surrogate.

DiffusionPDE on Darcy flow

Huang et al. (2024) introduced DiffusionPDE, a framework that goes beyond the standard “learn prior, condition on likelihood” paradigm. DiffusionPDE jointly models the PDE coefficient field and the solution field as a coupled system, using a single diffusion model that denoises both simultaneously.

The architecture treats the input as a multi-channel “image”: channel 1 is the permeability field $a(𝐱)$, and channel 2 is the pressure field $u(𝐱)$. During training, noise is added to both channels, and a U-Net-based denoising network learns to predict the clean pair $(a,u)$ from the noisy version. During inference, the model can perform:

1. Forward problems: Given a clean $a$ (no noise on channel 1) and pure noise on channel 2, the denoising process “generates” the solution $u$.

2. Inverse problems: Given sparse observations of $u$ (partial noise on channel 2) and pure noise on channel 1, the model simultaneously recovers both $a$ and the full $u$.

3. Sparse recovery: Given both $a$ and sparse $u$ observations, the model infills the missing pressure data.

The key advantage of the joint approach is that the model learns the physical coupling between $a$ and $u$: it internalises the constraint that $-\nabla \cdot (a\nabla u)=f$, even though this PDE is never explicitly imposed in the loss. On the Darcy flow benchmark, DiffusionPDE achieves state-of-the-art results for inverse recovery with as few as 1% observed pressure values.

Transformer operators on Poisson variants

While FNO operates in Fourier space (limiting it to periodic or rectangular domains), Transformer-based operators can handle arbitrary geometries by treating the discretised field as a sequence of tokens-one per grid point or mesh node.

The Galerkin Transformer (Cao, 2021) replaces the softmax attention kernel with a linear (Galerkin-type) attention mechanism: (ATTN)$$\begin{array}{ccc}& \text{Attn}(Q,K,V)=\frac{1}{n}Q(K^{\top }V),& \text{<span class="tml-eqn"></span>}\end{array}$$ where $Q,K,V\in {ℝ}^{n\times d}$ are the query, key, and value matrices and $n$ is the number of spatial tokens. The crucial difference from standard softmax attention is the order of operations: by computing $K^{\top }V$ first (an $d\times d$ matrix), the cost drops from $O(n^{2}d)$ to $O(n{d}^2)$-a dramatic saving when $n$ (the number of spatial points) is large and $d$ (the embedding dimension) is moderate.

On Poisson-type benchmarks (including variable-coefficient variants on L-shaped and circular domains), Transformer operators achieve competitive or superior accuracy to FNO, with the added benefit of handling non-rectangular domains naturally.

The Heat Equation and Diffusion-Reaction Systems

We now add the dimension that elliptic equations lack: time. Parabolic PDEs describe processes that evolve toward equilibrium-the diffusion of heat, the spreading of a chemical concentration, the relaxation of a strained material. The simplest and most canonical parabolic PDE is the heat equation, and it occupies a special place in this book because the entire theory of diffusion generative models is built on a mathematical analogy with heat diffusion.

The heat equation: simplest parabolic PDE

The heat equation has several properties that make it both physically important and mathematically tractable:

Diffusion-reaction equations

The heat equation describes pure diffusion. Adding a nonlinear reaction term $R(u)$ creates a far richer class of equations:

The interplay between diffusion (which smooths and spreads) and reaction (which amplifies, saturates, or oscillates) produces some of the most visually striking phenomena in all of PDE theory.

Fisher–KPP equation

The simplest nontrivial diffusion-reaction equation is the Fisher–KPP equation (Fisher, 1937; Kolmogorov, Petrovskii, and Piskunov, 1937): (Fisher)$$\begin{array}{ccc}& \frac{\partial u}{\partial t}=D{\nabla }^{2}u+ru(1-u),& \text{<span class="tml-eqn"></span>}\end{array}$$ where $r>0$ is the growth rate. This equation models the spread of an advantageous gene through a population: $u$ represents the local frequency of the gene, diffusion models spatial migration, and the logistic term $ru(1-u)$ models selection. The Fisher–KPP equation admits travelling wave solutions of the form $u(𝐱,t)=\varphi (x-ct)$ that propagate at a minimum speed $c^{\ast }=2\sqrt{Dr}$.

Gray–Scott system

For truly spectacular pattern formation, we turn to the Gray–Scott system, a two-component reaction-diffusion model: (Grayscott)$$\begin{array}{ccc}& \begin{array}{cc}\frac{\partial u}{\partial t}& =D_u{\nabla }^{2}u-u{v}^2+F(1-u),\\ \frac{\partial v}{\partial t}& =D_v{\nabla }^{2}v+u{v}^2-(F+k)v,\end{array}& \text{<span class="tml-eqn"></span>}\end{array}$$ where $u$ and $v$ represent concentrations of two chemical species, $D_u$ and $D_v$ are their diffusivities, $F$ is a feed rate, and $k$ is a kill rate. Depending on the values of $F$ and $k$, this system produces a zoo of Turing-type patterns: spots, stripes, spirals, labyrinthine structures, and pulsating domains.

PDEBench diffusion-reaction benchmarks

The PDEBench benchmark suite (Takamoto et al., 2022) provides standardised datasets and evaluation protocols for neural PDE solvers. For diffusion-reaction systems, PDEBench includes:

• 1D diffusion-reaction: $\partial u/\partial t=D{\partial }^{2}u/\partial x^2+R(u)$ with $R(u)=\rho u(1-u)$ on $[\mathrm{0,1}]\times [\mathrm{0,1}]$. Dataset: 10000 trajectories at resolution $1024\times 101$ (space $\times$ time).

• 2D diffusion-reaction (2-species): The system $\partial u/\partial t=D_u{\nabla }^{2}u+R_u(u,v)$, $\partial v/\partial t=D_v{\nabla }^{2}v+R_v(u,v)$ on ${[\mathrm{0,1}]}^2\times [\mathrm{0,5}]$. Dataset: 1000 trajectories at $128\times 128\times 101$.

On these benchmarks, the relative performance of different neural architectures reveals interesting patterns:

PINNs for the heat equation

The heat equation is often the first PDE attempted with PINNs, and for smooth initial conditions, it works beautifully. Consider the 1D heat equation on $[\mathrm{0,1}]\times [0,T]$: (Heat1d)$$\begin{array}{ccc}& \frac{\partial u}{\partial t}=D\frac{{\partial }^{2}u}{\partial x^2},u(x,0)=u_0(x),u(0,t)=u(1,t)=0.& \text{<span class="tml-eqn"></span>}\end{array}$$ With $u_0(x)=\sin (\pi x)$, the exact solution is $u(x,t)=e^{-D{\pi }^{2}t}\sin (\pi x)$-a smooth, exponentially decaying function that PINNs approximate easily.

Neural operators for parametric heat problems

For the parametric heat equation-where we want to solve for many different initial conditions or diffusivities-neural operators offer a compelling alternative to PINNs. The task is to learn the solution operator: (Operator)$$\begin{array}{ccc}& 𝒢:u_0(\cdot )⟼u(\cdot ,T),& \text{<span class="tml-eqn"></span>}\end{array}$$ mapping the initial temperature distribution to the temperature field at a fixed future time $T$ (or, more ambitiously, to the entire trajectory $u(\cdot ,t)$ for $t\in [0,T]$).

By Theorem Theorem 17, this operator is given by convolution with the heat kernel-a linear, smoothing operator. In Fourier space: (Fourier)$$\begin{array}{ccc}& \hat{u}(𝐤,T)=e^{-D|𝐤{|}^{2}T}{\hat{u}}_0(𝐤).& \text{<span class="tml-eqn"></span>}\end{array}$$ This is a pointwise multiplication in Fourier space by an exponentially decaying filter-precisely the kind of operation that FNO is designed to learn.

This is a rare case where we can prove a neural architecture is perfectly matched to the PDE structure. Of course, real applications involve variable diffusivity, nonlinear reactions, and complex domains where the Fourier structure breaks down.

Score Matching via Differentiable Physics (SMDP)

Holl et al. introduced Score Matching via Differentiable Physics (SMDP), a framework that combines score-based generative modelling with differentiable PDE solvers for inverse problems. The idea is elegant: use a differentiable numerical solver (implemented in a framework like JAX or PhiFlow) as a “physics layer” inside the generative pipeline.

For the heat equation inverse problem-recovering the initial temperature $u_0$ from a noisy observation of $u(\cdot ,T)$-SMDP proceeds as follows:

1. Train a diffusion model on a dataset of plausible initial conditions $\{u_0^{(i)}\}$. This gives a learned score ${\nabla }_{u_0}\log p_{\sigma }(u_0)$ at noise level $\sigma$.

2. Compute the likelihood score by differentiating through the forward solver: (Likelihood)$$\begin{array}{ccc}& {\nabla }_{u_0}\log p(𝐲|u_0)=-\frac{1}{{\sigma }_{\text{obs}}^2}{\frac{\partial 𝒮}{\partial u_0}}^{\top }(𝒮(u_0)-𝐲),& \text{<span class="tml-eqn"></span>}\end{array}$$ where $𝒮(u_0)$ denotes the forward solver mapping $u_0$ to the observation $\mathscr{H}(u(\cdot ,T))$, and the Jacobian $\partial 𝒮/\partial u_0$ is computed via backpropagation through the solver.

3. Sample from the posterior using the combined score in a reverse diffusion process.

Poisson Flow Generative Models (PFGM)

We close our discussion of parabolic PDEs with a beautiful example of traffic in the opposite direction: rather than using generative models to solve PDEs, the Poisson Flow Generative Model (PFGM) of Xu et al. (2022) uses a PDE to build a better generative model.

The key observation is this: consider a data distribution $p_{\text{data}}$ supported on ${ℝ}^d$. Embed the data in a higher-dimensional space ${ℝ}^{d+1}$ by placing it on the hyperplane $z=0$. Now consider the electric field generated by the data charges: by Gauss's law, this field satisfies the Poisson equation $\nabla \cdot 𝐄=\rho$ in the augmented $(d+1)$-dimensional space.

The follow-up work PFGM++ (Xu et al., 2023) generalises this to arbitrary dimensions $D>d$ for the augmented space, interpolating between PFGM (when $D=d+1$) and standard diffusion models (in the limit $D\to \infty$). The optimal dimension $D^{\ast }$ depends on the dataset and provides a new hyperparameter that can be tuned for improved sample quality.

This result reveals a deep connection: diffusion models are the infinite-dimensional limit of electrostatic generative models. The Poisson equation (elliptic) and the heat equation (parabolic) are thus unified in the PFGM framework as two ends of a continuous spectrum.

Burgers' Equation and Conservation Laws

If the heat equation is the gentle giant of PDE theory-smoothing everything it touches, dissipating energy, approaching equilibrium-then Burgers' equation is its mischievous cousin. Burgers' equation adds a single ingredient to diffusion-nonlinear advection-and the consequences are dramatic: steep gradients, shock waves, and a rich interplay between smoothing and steepening that prefigures the full complexity of fluid dynamics.

For neural PDE solvers, Burgers' equation occupies a uniquely important position. It is simple enough to admit an exact analytical solution (via the Cole–Hopf transformation), yet complex enough to expose the fundamental limitations of PINNs and neural operators when confronted with discontinuities. It is, in a very real sense, the Drosophila of computational physics: a model organism for studying how neural architectures handle shocks.

The viscous Burgers equation

The equation can also be written in conservation form: (Conservation)$$\begin{array}{ccc}& \frac{\partial u}{\partial t}+\frac{\partial }{\partial x}\left(\frac{u^2}{2}\right)=\nu \frac{{\partial }^{2}u}{\partial x^2},& \text{<span class="tml-eqn"></span>}\end{array}$$ which reveals its structure as a conservation law with flux function $f(u)=u^2/2$ and viscous regularisation.

Shock formation and the inviscid limit

The physics of Burgers' equation is governed by a competition between two effects:

1. Nonlinear steepening: The advection term $u\partial u/\partial x$ causes the wave to steepen. Points where $u$ is large move faster than points where $u$ is small, so the wave profile tilts forward until it threatens to become multi-valued.

2. Viscous diffusion: The term $\nu {\partial }^{2}u/\partial x^2$ smooths out gradients, counteracting the steepening.

When viscosity is large, diffusion dominates and the solution remains smooth. When viscosity is small, steepening dominates and the solution develops a shock-a region where $u$ changes rapidly over a width of order $O(\nu )$.

In the inviscid limit $\nu \to 0$, the Burgers equation becomes the inviscid Burgers equation $\partial u/\partial t+u\partial u/\partial x=0$. For generic smooth initial data with $u_0^{\prime }(x)<0$ somewhere, the solution develops a gradient catastrophe (the derivative blows up) at a finite breaking time: (Breaktime)$$\begin{array}{ccc}& t^{\ast }=\frac{-1}{{\min }_x{u}_0^{\prime }(x)}.& \text{<span class="tml-eqn"></span>}\end{array}$$ After $t^{\ast }$, the classical solution ceases to exist, and we must pass to weak solutions satisfying the Rankine–Hugoniot jump conditions and an entropy condition.

The Cole–Hopf transformation

The viscous Burgers equation can be exactly solved by a remarkable change of variables:

Since the heat equation is exactly solvable (via convolution with the heat kernel), the Cole–Hopf transformation gives an exact solution of the viscous Burgers equation for any initial data.

Why PINNs fail on shocks

The Burgers shock is a litmus test for neural PDE solvers: any method that cannot resolve it is fundamentally limited. Standard PINNs fail on this problem, and understanding why they fail illuminates the path toward better architectures.

Spectral bias analysis

The shock in the viscous Burgers solution has width $O(\nu )$. When $\nu =0.01/\pi \approx 0.003$, the shock occupies roughly 0.3% of the spatial domain. Resolving this feature requires high-frequency Fourier modes up to $k\sim 1/\nu \approx 300$. But the spectral bias of MLPs with smooth activations (tanh, sigmoid) means that the Neural Tangent Kernel eigenvalues ${\lambda }_k$ decay exponentially for $k\gg 1$, making modes near $k=300$ essentially unlearnable within a reasonable number of gradient steps.

This proposition makes the problem precise: PINNs do not just struggle with shocks-they require exponentially more training time as the shock becomes thinner.

PINNsFormer: Transformer + wavelet activation

Zhao et al. (2024) proposed PINNsFormer, a Transformer-based architecture that addresses both spectral bias and temporal dependencies. The key innovations are:

1. Wavelet activation functions: Replace tanh/sigmoid with the Mexican hat wavelet $\psi (x)=(1-x^2)e^{-x^2/2}$. Wavelets have compact support in both physical and frequency space, providing multi-scale representation without spectral bias.

2. Temporal attention: Instead of treating time as just another input coordinate, PINNsFormer creates a sequence of pseudo-time-steps $\{t_1,\dots ,t_K\}$ and applies self-attention across them, allowing the network to capture temporal causality and long-range time dependencies.

3. Multi-head physics attention: Different attention heads specialise in different spatial scales, enabling simultaneous resolution of the smooth background and the sharp shock.

On the standard Burgers benchmark ($\nu =0.01/\pi$), PINNsFormer reduces the relative $L^2$ error from $\sim$11% (standard PINN) to $\sim$1.2%, a nearly 10$\times$ improvement. The wavelet activation alone accounts for approximately half of this gain, with temporal attention providing the other half.

FNO and operator learning benchmarks on Burgers

Unlike PINNs (which solve a single instance), neural operators learn the family of Burgers solutions parameterised by initial conditions. The standard benchmark task is: (Operator)$$\begin{array}{ccc}& 𝒢:u_0(\cdot )⟼u(\cdot ,T),& \text{<span class="tml-eqn"></span>}\end{array}$$ mapping the initial condition $u_0$ to the solution at a fixed future time $T=1$, with $\nu =0.1$ or $\nu =0.01$.

The results depend dramatically on the viscosity:

SafeDiffCon: safe diffusion for Burgers control

Wei et al. (2025) introduced SafeDiffCon (Safe Diffusion-based Controller), a framework that uses diffusion models for optimal control of PDE systems subject to safety constraints.

Consider the control problem: steer the Burgers equation from an initial state $u_0$ to a target state $u_{\text{target}}$ by applying a forcing term $g(x,t)$: (Safediffcon)$$\begin{array}{ccc}& \frac{\partial u}{\partial t}+u\frac{\partial u}{\partial x}=\nu \frac{{\partial }^{2}u}{\partial x^2}+g(x,t),& \text{<span class="tml-eqn"></span>}\end{array}$$ subject to the constraint that $u(x,t)$ must remain within safety bounds $[u_{\min },u_{\max }]$ for all $(x,t)$.

SafeDiffCon trains a conditional diffusion model on a dataset of optimal control trajectories (computed offline using adjoint methods). At inference, the diffusion model generates candidate control sequences, and a safety filter-based on control barrier functions-projects these onto the feasible set in real time.

On the Burgers control benchmark, SafeDiffCon achieves 37% lower tracking error than model predictive control (MPC) baselines while maintaining zero safety violations, compared to a 12% violation rate for unconstrained diffusion controllers.

Riemann problems and conservation law structure

The inviscid Burgers equation is the simplest example of a scalar conservation law: (Conslaw)$$\begin{array}{ccc}& \frac{\partial u}{\partial t}+\frac{\partial f(u)}{\partial x}=0,& \text{<span class="tml-eqn"></span>}\end{array}$$ with flux $f(u)=u^2/2$. The Riemann problem-the initial-value problem with piecewise-constant initial data (Riemann)$$\begin{array}{ccc}& u(x,0)=\{\begin{array}{cc}{u}_L& \text{if }x<0,\\ u_R& \text{if }x>0,\end{array}& \text{<span class="tml-eqn"></span>}\end{array}$$ -is the fundamental building block of conservation law theory.

For Burgers' equation, the Riemann problem has exactly two cases:

1. Shock wave ($u_L>u_R$): The solution is a discontinuity propagating at speed $s=(u_L+u_R)/2$.

2. Rarefaction wave ($u_L<u_R$): The solution is a continuous, self-similar expansion fan $u(x,t)=x/t$ connecting $u_L$ to $u_R$.

Shallow Water Equations

We conclude this first survey of PDE families with a system that is far more complex than anything we have encountered so far: the shallow water equations (SWE). These equations govern the motion of a thin layer of fluid under gravity-think of ocean waves, river flows, tidal bores, or the large-scale circulation of the atmosphere. The SWE are a system of hyperbolic conservation laws that couple mass and momentum, admit both smooth waves and shocks, and-when posed on the rotating sphere-incorporate the Coriolis force that dominates large-scale geophysical dynamics.

For neural PDE solvers, the shallow water equations represent a major step up in difficulty. They are a system (not a scalar equation), they are nonlinear (with wave-wave interactions), they live on non-trivial geometries (the sphere), and they exhibit multi-scale dynamics (from small ripples to planetary-scale Rossby waves). Success on the SWE is a prerequisite for tackling the full atmospheric equations that drive weather prediction.

The SWE system

Applications and physical significance

The shallow water equations arise whenever a fluid layer is much thinner than its horizontal extent:

• Tsunami propagation: Ocean waves generated by earthquakes are well-modelled by the SWE, as the wavelength ($\sim$100–500 km) vastly exceeds the ocean depth ($\sim$4 km). Real-time SWE simulations are critical for tsunami early warning systems.

• Atmospheric dynamics: The “shallow” atmosphere of Earth (thickness $\sim$10 km vs. radius $\sim$6400 km) supports large-scale dynamics that are well-approximated by the SWE on the rotating sphere-Rossby waves, Kelvin waves, and geostrophic adjustment.

• River and coastal flooding: The SWE with friction terms model the propagation of flood waves through river systems and coastal inundation during storm surges.

• Ocean circulation: The large-scale, depth-averaged ocean circulation (gyres, western boundary currents) is governed by the SWE with wind forcing and bottom friction.

Spherical Fourier Neural Operator (SFNO)

Standard FNO uses the FFT on rectangular grids, but the Earth is a sphere. For geophysical applications, we need an operator that respects spherical geometry. The Spherical Fourier Neural Operator (SFNO) of Bonev et al. (2023) generalises FNO by replacing the FFT with the spherical harmonic transform (SHT).

The SFNO architecture naturally handles the geometry of the sphere: there are no pole singularities (which plague latitude-longitude grids), and the spherical harmonics form a complete, orthonormal basis on ${𝕊}^2$. The spectral truncation at ${\ell }_{\max }$ provides a natural scale separation, analogous to the wavenumber truncation in standard FNO\@.

On the SWE benchmark (rotating shallow water on the sphere with realistic Coriolis forces), SFNO achieves significantly lower errors than standard FNO applied to equirectangular projections, because the latter introduces artificial distortions near the poles.

ECHO: million-point generative transformer

Lippe et al. (2024) introduced ECHO (Efficient Conditional Hierarchical Operator), a generative Transformer that scales to over one million spatial points-the resolution required for realistic geophysical simulations.

The key challenge is computational: standard self-attention over $N$ points has $O(N^2)$ cost, which is prohibitive for $N={10}^6$. ECHO addresses this through a hierarchical compression strategy:

1. Patch embedding: The spatial field is divided into patches (e.g., $16\times 16$), and each patch is embedded into a single token. This reduces $N={10}^6$ points to $\sim$4000 tokens.

2. Multi-scale attention: Attention is applied at multiple resolutions using a U-shaped architecture: first at the coarsest level (global context), then progressively finer levels (local detail).

3. Conditional generation: The current state $𝐔(t_k)$ conditions the model to generate the next state $𝐔(t_{k+1})$, enabling autoregressive rollout over long time horizons.

On the SWE benchmark, ECHO produces physically coherent solution fields that maintain conservation of mass and geostrophic balance over multi-day rollouts, even when generating at resolutions unseen during training.

PROSE-FD: multimodal PDE foundation model

Sun et al. (2024) proposed PROSE-FD (PDE Representation and Operator Solving via Finite Differences), a multimodal foundation model that can solve different PDEs with a single set of weights. The key insight is to encode PDE information as a symbolic token sequence (a textual description of the equation) alongside the numerical field data.

For the shallow water equations, PROSE-FD receives:

1. Symbolic tokens: A tokenised representation of the SWE, including the equation structure, parameter values ($g$, $f$), and boundary conditions.

2. Numerical tokens: The initial state $(h,v_1,v_2)$ discretised on a grid and flattened into a token sequence.

The Transformer processes both modalities jointly, using cross-attention between symbolic and numerical tokens, and outputs the predicted state at the next time step.

DiTTO: diffusion-inspired temporal transformer

Hemmasian and Barati Farimani (2024) introduced DiTTO (Diffusion-inspired Temporal Transformer Operator), which combines the temporal modelling strengths of diffusion processes with the spatial modelling strengths of Transformers.

The core idea is to treat the temporal evolution of a PDE solution as a “denoising” process:

1. Temporal diffusion: The difference between the solution at time $t_k$ and $t_{k+1}$ is modelled as a “noise” that must be predicted and removed.

2. Iterative refinement: Instead of predicting $u(\cdot ,t_{k+1})$ in a single forward pass, DiTTO iteratively refines the prediction over multiple “denoising” steps, each conditioned on $u(\cdot ,t_k)$.

3. Adaptive compute: For time steps where the dynamics are slow (smooth evolution), fewer refinement steps suffice; for time steps with rapid changes (shock formation, wave breaking), more steps are needed.

On the PDEBench shallow water benchmark, DiTTO achieves 15–25% lower error than single-pass neural operators (FNO, U-Net) while using only $\sim$3$\times$ more compute, a favourable trade-off.

PDEBench shallow water benchmarks

PDEBench provides a standardised 2D shallow water dataset:

• Domain: ${[\mathrm{0,1}]}^2$ with periodic boundary conditions.

• Resolution: $128\times 128$ spatial, 21 time steps.

• Initial conditions: Random height perturbations on a flat bottom ($b=0$), no Coriolis force ($f=0$).

• Dataset: 1000 trajectories for training, 200 for testing.

Navier–Stokes: The Crown Jewel of Fluid Dynamics

“Turbulence is the most important unsolved problem of classical physics.” - Richard Feynman

We have now assembled the full toolkit: physics-informed neural networks, neural operators, Transformers, and diffusion models. It is time to unleash these tools on the equation that matters most in applied mathematics and engineering-the Navier–Stokes (NS) equations. The NS equations govern the motion of viscous fluids and appear in an astonishing range of applications: aircraft design, weather prediction, cardiovascular modelling, turbine optimisation, ocean circulation, and the combustion dynamics inside every car engine. They are also the subject of one of the seven Clay Millennium Prize Problems, which asks whether smooth solutions in three dimensions exist for all time-a question that remains open after more than a century.

In this section we show how Transformer and diffusion-based generative models are transforming computational fluid dynamics (CFD), replacing hours of simulation with seconds of inference while preserving the multi-scale physics that makes these equations so challenging.

The incompressible Navier–Stokes equations

For an incompressible, Newtonian fluid with density $\rho$ and dynamic viscosity $\mu$, the velocity field $𝐮(𝐱,t)\in {ℝ}^d$ and pressure field $p(𝐱,t)$ satisfy:

The Reynolds number (Reynolds)$$\begin{array}{ccc}& \mathrm{Re}=\frac{\rho UL}{\mu },& \text{<span class="tml-eqn"></span>}\end{array}$$ where $U$ and $L$ are characteristic velocity and length scales, quantifies the ratio of inertial to viscous forces. At low $\mathrm{Re}$ (think honey pouring from a spoon), viscosity dominates and flows are smooth (laminar). As $\mathrm{Re}$ increases (think a Boeing 747 at cruising altitude, $\mathrm{Re}\sim {10}^7$), flows become turbulent-chaotic, multi-scale, and computationally catastrophic.

Compressible Navier–Stokes equations

When the fluid density $\rho (𝐱,t)$ varies (e.g. at high Mach numbers, in combustion, or in astrophysical flows), we must solve the full compressible system: (Compns MASS)$$\begin{array}{cccc}& \frac{\partial \rho }{\partial t}+\nabla \cdot (\rho 𝐮)& =0,& \text{<span class="tml-eqn"></span>}\\ & \frac{\partial (\rho 𝐮)}{\partial t}+\nabla \cdot (\rho 𝐮\otimes 𝐮)& =-\nabla p+\nabla \cdot 𝝉+\rho 𝐟,& \text{<span class="tml-eqn"></span>}\\ & \frac{\partial (\rho E)}{\partial t}+\nabla \cdot [(\rho E+p)𝐮]& =\nabla \cdot (𝝉\cdot 𝐮)+\nabla \cdot (\kappa \nabla T)+\rho 𝐟\cdot 𝐮,& \text{<span class="tml-eqn"></span>}\end{array}$$ where $E=e+\frac{1}{2}|𝐮{|}^2$ is the total specific energy, $e$ is the internal energy, $T$ is the temperature, $\kappa$ is the thermal conductivity, and $𝝉$ is the viscous stress tensor. An equation of state (e.g. the ideal gas law $p=\rho RT$) closes the system.

Transformer-based approaches for Navier–Stokes

Transformers bring two distinctive strengths to NS modelling: global receptive fields (via self-attention) that capture long-range pressure-velocity coupling, and geometry flexibility that handles unstructured meshes arising in industrial CFD\@. We survey the most impactful Transformer-based methods.

Transolver: Physics-aware Transformers for industrial CFD

Real-world CFD meshes are not regular grids. The mesh around an airfoil, for instance, is highly refined near the leading edge, coarse in the far field, and may contain millions of nodes. Standard Transformers struggle with this because naive self-attention scales as $O(N^2)$ in the number of mesh nodes $N$.

Transolver (Wu et al., 2024) addresses this by introducing physics-aware slicing: the computational domain is decomposed into a small number of overlapping “slices” that respect the underlying physics (e.g. boundary layers, wake regions, far field). Attention is computed within and across slices, reducing complexity from $O(N^2)$ to $O(N\cdot S+S^2)$, where $S\ll N$ is the number of slices.

HAMLET: Fluid-structure interaction

In many engineering problems, the fluid interacts with a deformable or moving solid-think of blood flow through a beating heart or aeroelastic flutter of a wing. The HAMLET framework (Graph Mesh Learner with Efficient Transformers; Bryutkin et al., 2024) tackles these fluid-structure interaction (FSI) problems by combining graph neural networks (for local mesh interactions) with efficient Transformer attention (for global coupling between the fluid and structural domains).

HAMLET uses a two-level architecture:

1. Local GNN encoder: A message-passing GNN operates on the mesh graph, capturing local interactions between neighbouring nodes (stress propagation, convection).

2. Global Transformer: The GNN node embeddings are pooled into a moderate number of “super-nodes” representing patches of the mesh, and a Transformer computes attention among these super-nodes. This captures the global fluid-structure coupling (e.g. the pressure wave from the fluid reaching the solid boundary and causing deformation, which in turn alters the flow).

GAOT: General geometry operator Transformer

The Geometry-Aware Operator Transformer (GAOT; Li et al., 2024) extends the operator-learning paradigm to arbitrary computational geometries. The key innovation is a geometry-conditioned attention mechanism: the attention weights depend not only on the function values at query/key positions but also on the local geometric features of the mesh (curvature, normal vectors, element aspect ratios).

Foundation models: MPP and CoDA-NO

A striking recent trend is the development of foundation models for PDEs-large Transformer-based architectures pretrained on diverse PDE datasets and fine-tuned for specific applications.

Multiple Physics Pretraining (MPP) (McCabe et al., 2023) pretrains a Vision Transformer on a large corpus of PDE simulation data spanning Navier–Stokes, shallow water, diffusion–reaction, compressible Euler, and Maxwell's equations. The key insight is that diverse PDE fields (velocity, pressure, temperature, electromagnetic fields) can be tokenised into patches and processed by a single Transformer backbone, much as language tokens from different languages can be processed by a multilingual LLM\@.

CoDA-NO (Cooperation of Domain Adaptation with Neural Operators; Subramanian et al., 2024) takes a complementary approach: instead of pretraining on simulation data, it uses domain adaptation techniques to transfer a neural operator trained on one PDE family to another. The method introduces learnable domain tokens that encode the “physics signature” of each PDE type, allowing a single model to serve as a multi-physics solver.

Diffusion-based approaches for Navier–Stokes

Diffusion models bring a complementary set of strengths: distributional modelling (generating diverse plausible solutions rather than a single point estimate), flexible conditioning (sparse observations, physics constraints), and high-frequency fidelity (the iterative denoising process can recover fine-scale structures that regression-based methods blur away).

DiffusionPDE: Sparse reconstruction

Imagine you have a sparse set of pressure sensors scattered across an airfoil surface and wish to reconstruct the full pressure and velocity fields. This is an inverse problem-ill-posed, since many flow configurations are consistent with a handful of sensor readings.

DiffusionPDE (Huang et al., 2024) formulates this as conditional generation. A diffusion model is trained on a dataset of full NS solution fields (velocity $𝐮$, pressure $p$). At inference, the sparse observations $𝐲=ℳ({𝐮}^{\ast },p^{\ast })$ (where $ℳ$ is the observation operator selecting sensor locations) are incorporated via a guidance mechanism.

Conditional diffusion for forecasting and assimilation

A landmark paper at NeurIPS 2024 (Rozet et al., 2024) demonstrated that conditional diffusion models can simultaneously perform forecasting (predicting the future state given the current state) and data assimilation (incorporating new observations as they arrive) for fluid dynamics.

The key idea is to train a score network $s_{\varphi }({𝐮}_t,t|{𝐮}_0^{\text{prev}})$ that is conditioned on the previous time step's solution field. This score network defines a conditional distribution over the next time step: (COND DIFF Forecast)$$\begin{array}{ccc}& p_{\varphi }({𝐮}^{(n+1)}|{𝐮}^{(n)})\propto \int p_{\varphi }({𝐮}^{(n+1)}|{𝐮}_t^{(n)})q({𝐮}_t^{(n)}|{𝐮}^{(n)})\mathrm{d}{𝐮}_t^{(n)}.& \text{<span class="tml-eqn"></span>}\end{array}$$ For data assimilation, observations ${𝐲}^{(n+1)}$ at the new time step are incorporated via Bayes' rule: (COND DIFF Assim)$$\begin{array}{ccc}& p({𝐮}^{(n+1)}|{𝐮}^{(n)},{𝐲}^{(n+1)})\propto p({𝐲}^{(n+1)}|{𝐮}^{(n+1)})p_{\varphi }({𝐮}^{(n+1)}|{𝐮}^{(n)}).& \text{<span class="tml-eqn"></span>}\end{array}$$

PDE-Diffusion: Temporal coherence

A common pitfall of applying image diffusion models frame-by-frame to PDE trajectories is temporal incoherence: each frame is individually plausible, but consecutive frames may exhibit flickering, non-physical jumps, or violation of conservation laws.

PDE-Diffusion (Shu et al., 2024) addresses this by conditioning the denoising network on a temporal context window of $K$ previous frames and adding a physics regulariser that penalises violations of the PDE residual in the generated trajectory: (PDE Diffusion)$$\begin{array}{ccc}& {ℒ}_{\text{PDE-Diff}}={ℒ}_{\text{DDPM}}+{\lambda }_{\text{phys}}𝔼\left[{\Vert 𝒩[{\widehat{𝐮}}_{0|t}]\Vert }^2\right],& \text{<span class="tml-eqn"></span>}\end{array}$$ where $𝒩$ denotes the PDE residual operator and ${\widehat{𝐮}}_{0|t}$ is the denoised prediction at diffusion time $t$. The physics loss guides the model toward temporally coherent, physically consistent trajectories.

SimDiffPDE: Image-to-image translation

A refreshingly direct approach is taken by SimDiffPDE (Wang et al., 2024), which treats PDE solving as a conditional image-to-image translation problem. The input “image” encodes the PDE parameters (initial condition, boundary geometry, forcing term) as a multi-channel image, and the output “image” is the solution field. A standard conditional diffusion model (essentially a Palette-style architecture) is trained on pairs $(c_i,{𝐮}_i)$ of parameter images and solution fields.

Key results: Error metrics and speedups

To contextualise the methods we have discussed, we present a comparison of representative results on standard NS benchmarks. The metrics reported are the relative $L^2$ error (lower is better) and wall-clock speedup over the reference solver (higher is better).

Generative Modelling for Turbulence

“Big whirls have little whirls that feed on their velocity, and little whirls have lesser whirls and so on to viscosity.” - Lewis Fry Richardson (1922)

Turbulence is, by many accounts, the last great unsolved problem of classical physics. It is ubiquitous-from the mixing of cream in coffee to the formation of galaxies-yet its full mathematical characterisation has eluded physicists and mathematicians for over a century. In this section, we explain why turbulence is the ultimate test case for generative models of physical systems, and survey the remarkable recent progress in generating turbulent flows using GANs, diffusion models, and hybrid approaches.

Why turbulence is THE challenge

Turbulent flows are characterised by:

• Chaotic sensitivity: Infinitesimal perturbations to initial conditions lead to exponentially diverging trajectories (positive Lyapunov exponents).

• Multi-scale structure: Energy is injected at large scales and cascades down to the Kolmogorov microscale $\eta$, where it is dissipated by viscosity. The ratio of the largest to smallest scales is $L/\eta \sim {\mathrm{Re}}^{3/4}$.

• Statistical stationarity: While individual realisations are unpredictable, the statistics (mean velocity profiles, energy spectra, structure functions) are reproducible and governed by universal scaling laws.

This presents a fundamental mismatch with deterministic neural surrogates. A single “best guess” prediction of a turbulent flow cannot be physically meaningful beyond the predictability time (typically a few eddy turnover times). What is meaningful are the statistics: the energy spectrum, probability density functions of velocity increments, and higher-order moments.

Energy cascade and the Kolmogorov spectrum

In 1941, Andrei Kolmogorov proposed one of the most celebrated results in theoretical physics: in the inertial range (scales much smaller than the energy injection scale $L$ and much larger than the dissipation scale $\eta$), the energy spectrum takes a universal form.

The energy spectrum is the single most important diagnostic for evaluating generative turbulence models: any model that fails to reproduce the $-5/3$ scaling in the inertial range is physically inadequate, regardless of its pixel-wise accuracy.

GAN-era approaches to turbulence

Before diffusion models dominated the field, GANs were the primary generative approach for turbulence modelling. Several key ideas from this era remain influential.

Physics-constrained GANs

The naive application of a GAN (e.g. DCGAN) to turbulence data produces visually plausible velocity fields but with incorrect statistics: the energy spectrum may have the wrong slope, or the two-point correlations may be inaccurate.

To address this, several groups introduced physics-constrained discriminator losses:

1. Covariance matching (Wu et al., 2020): Add a penalty to the generator loss that measures the discrepancy between the spatial covariance of generated and real velocity fields: (GAN COV)$$\begin{array}{ccc}& {ℒ}_{\text{cov}}={\Vert R_{ij}^{\text{gen}}(𝐫)-R_{ij}^{\text{data}}(𝐫)\Vert }_{L^2(𝐫)}^2,& \text{<span class="tml-eqn"></span>}\end{array}$$ where $R_{ij}(𝐫)=⟨u_i^{\prime }(𝐱)u_j^{\prime }(𝐱+𝐫)⟩$ is the two-point velocity correlation tensor.

2. Energy spectrum constraint (Kim & Lee, 2020): Compute the energy spectrum of generated fields via FFT and penalise deviations from the target spectrum: (GAN Spectrum)$$\begin{array}{ccc}& {ℒ}_{\text{spec}}={\int }_0^{k_{\max }}|E^{\text{gen}}(k)-E^{\text{data}}(k){|}^2\mathrm{d}k.& \text{<span class="tml-eqn"></span>}\end{array}$$

3. Deterministic constraint (Bode et al., 2021): Enforce the divergence-free condition $\nabla \cdot 𝐮=0$ in the generator architecture via a stream-function or vector-potential parameterisation.

Diffusion models for turbulence

The limitations of GANs-mode collapse, training instability, and difficulty enforcing complex statistical constraints-motivated the shift to diffusion models.

“From Zero to Turbulence” (ICLR 2024)

A conceptually striking result is the ability to generate turbulent velocity fields from scratch, without requiring an initial condition from a previous simulation. Shu et al. (ICLR 2024) demonstrated this with a score-based diffusion model trained on a dataset of statistically stationary turbulence snapshots.

The model learns the equilibrium distribution of the turbulent flow-the stationary measure of the Navier–Stokes dynamics. Sampling from this distribution produces velocity fields that are statistically indistinguishable from DNS snapshots: the energy spectrum, velocity PDFs, and structure functions all match to within statistical uncertainty.

CoNFiLD: Conditional neural-field latent diffusion

A key limitation of direct diffusion on high-resolution turbulence fields is computational cost: 3D turbulence on a ${256}^3$ grid requires denoising in a space of $\sim 50$ million dimensions.

CoNFiLD (Conditional Neural Field Latent Diffusion; Du et al., Nature Communications 2024) addresses this by combining three ideas:

1. Neural field representation: The turbulence field $𝐮(𝐱)$ is represented as a continuous neural field ${𝐮}_{𝜽}(𝐱)=f_{𝜽}(𝐱)$, parameterised by a compact latent vector $𝜽\in {ℝ}^d$.

2. Latent diffusion: Instead of denoising the high-resolution field directly, CoNFiLD performs diffusion in the low-dimensional latent space ${ℝ}^d$ ($d\sim 256$–$512$ in practice), reducing the computational cost by orders of magnitude.

3. Conditional generation: The model is conditioned on physical parameters (Reynolds number, forcing type, boundary conditions) via cross-attention, enabling generation of turbulent fields across a range of flow regimes.

GenCFD: Score-based diffusion for CFD ensembles

GenCFD (Dreler et al., 2024) focuses on generating ensembles of turbulent flow realisations for uncertainty quantification. The key insight is that in many engineering applications, one does not need a single accurate prediction but rather a distribution of outcomes (e.g. the 95th percentile of drag on a structure exposed to turbulent wind).

GenCFD trains a score-based diffusion model on flow snapshots and uses it to generate large ensembles efficiently. The ensemble statistics (mean, variance, higher-order moments) are then used for risk assessment.

Learning turbulent flows: Operator-generative hybrids

A 2026 paper in Nature Communications (Lienen et al., 2026) proposes a powerful hybrid: a deterministic neural operator provides the conditional mean prediction, and a diffusion model generates the stochastic residual (the difference between the true field and the conditional mean).

This decomposition leverages the strengths of both paradigms: the neural operator captures the large-scale structure efficiently (one forward pass), while the diffusion model fills in the physically correct small-scale fluctuations. The result is both fast (the neural operator handles the bulk of the prediction) and statistically accurate (the diffusion model ensures the correct energy spectrum at high wavenumbers).

Super-resolution of turbulent flows

A particularly impactful application of generative models is super-resolution: taking a coarse-grid (Large Eddy Simulation, LES) flow field and adding physically consistent small-scale detail to produce a field at DNS resolution. This is exactly the problem that diffusion models excel at-adding high-frequency detail that is statistically correct but unpredictable in its exact realisation.

Why $L^2$ loss fails and diffusion succeeds

This subsection addresses a subtle but crucial point: why minimising the pixel-wise $L^2$ error is fundamentally wrong for turbulent flows, and why diffusion models naturally avoid this pitfall.

Helmholtz Equation and Wave Phenomena

We now turn from fluid dynamics to wave physics. The Helmholtz equation (Helmholtz)$$\begin{array}{ccc}& {\nabla }^{2}u(𝐱)+k^{2}u(𝐱)=f(𝐱),𝐱\in \mathrm{\Omega }\subset {ℝ}^d,& \text{<span class="tml-eqn"></span>}\end{array}$$ governs time-harmonic waves: it arises when one assumes a time-periodic solution $U(𝐱,t)=u(𝐱)e^{-i\omega t}$ in the wave equation ${\partial }_t^{2}U=c^2{\nabla }^{2}U+F$. Here $k=\omega /c$ is the wavenumber (frequency divided by wave speed), and $f$ is a source term.

The Helmholtz equation appears throughout physics and engineering:

• Acoustics: Sound wave propagation in rooms, concert halls, and noise barriers.

• Seismology: Seismic wave propagation through the Earth's subsurface.

• Electromagnetics: Scattering and diffraction of electromagnetic waves (radar, antennas).

• Quantum mechanics: The time-independent Schrödinger equation is a Helmholtz-type equation.

Why high-frequency waves are hard

The Helmholtz equation becomes fiendishly difficult as the wavenumber $k$ increases. There are two fundamental reasons:

1. Oscillatory solutions: The solution oscillates on a length scale $\lambda =2\pi /k$. To resolve these oscillations, any discretisation (finite elements, neural network collocation points) must have a spacing $h\lesssim \lambda /10$, leading to $O(k^d)$ degrees of freedom in $d$ dimensions. For seismic applications at $k\sim 100$, this means $\sim {10}^6$ unknowns in 2D and $\sim {10}^9$ in 3D\@.

2. Pollution effect: Even with sufficient mesh resolution, standard finite element methods accumulate phase errors that grow with $k$, requiring $h=O(k^{-(d+1)/d})$ for a fixed error tolerance-worse than the Nyquist limit.

3. Spectral bias of neural networks: PINNs with smooth activations have eigenvalues that decay rapidly with frequency (as discussed in Section Spectral bias), making them catastrophically slow for high-$k$ Helmholtz problems.

NSNO: Neumann Series Neural Operator

The Neumann Series Neural Operator (NSNO; Agrawal et al., 2024) is a Transformer-based architecture specifically designed for the Helmholtz equation. The key idea is to incorporate the mathematical structure of the Neumann series (the iterative expansion of the Green's function) into the neural operator architecture.

NSNO approximates each term in the Neumann series by a Transformer layer, with the attention mechanism playing the role of the integral operator $G_0$. Each layer effectively computes one term in the series, and the total network depth corresponds to the truncation order.

Diffusion for probabilistic Helmholtz solutions

A 2026 paper at ICLR (Baldassari et al., 2026) introduces a diffusion-based approach to the Helmholtz equation that provides calibrated uncertainty quantification. The key motivation is that in many applications (e.g. subsurface imaging), the wave speed $c(𝐱)$ is only approximately known, leading to uncertainty in the wavenumber field $k(𝐱)=\omega /c(𝐱)$.

The diffusion model is conditioned on:

• The source term $f(𝐱)$,

• A (possibly uncertain) wavenumber field $k(𝐱)$,

• Sparse boundary observations of $u$.

The model generates an ensemble of solution fields, each consistent with the observations and the PDE (within the uncertainty on $k$). The ensemble spread provides a principled measure of solution uncertainty.

Spectrum-preserving wavefield surrogates

Just as turbulence models must preserve the energy spectrum, wavefield surrogates must preserve the spatial frequency content of the Helmholtz solution. This is particularly challenging because the solution contains sharp features (caustics, shadow boundaries, diffraction patterns) that span a wide range of spatial frequencies.

Diffusion models typically achieve spectral fidelity scores of $𝒮>0.95$ on Helmholtz benchmarks, compared to $𝒮\approx 0.80$–$0.85$ for FNO and $𝒮\approx 0.70$ for PINNs with smooth activations (due to spectral bias).

Applications: Acoustics, seismic, and material design

Room acoustics

In architectural acoustics, one must solve the Helmholtz equation at frequencies up to $20\text{kHz}$ ($k\sim 370{\text{m}}^{-1}$ in air) inside complex geometries (concert halls, lecture theatres). Neural surrogates trained on parameterised room geometries can evaluate the acoustic field in milliseconds, enabling real-time auralization in virtual reality applications.

Seismic full waveform inversion

The Helmholtz equation is the frequency-domain formulation of the wave equation used in seismic exploration. Generative models serve as fast forward solvers within iterative inversion loops, replacing each expensive Helmholtz solve (which dominates the cost of full waveform inversion) with a neural-network evaluation.

RadioDiff-k: Helmholtz-informed diffusion for radio maps

An innovative application by Lee et al. (2025) combines Helmholtz physics with diffusion models for wireless communication. RadioDiff-k generates radio propagation maps-the received signal strength as a function of position in an urban environment-by conditioning a diffusion model on the environment geometry and transmitter locations, with the Helmholtz equation providing a physics prior.

The model incorporates the wavenumber $k$ (determined by the carrier frequency) as a conditioning variable, allowing a single model to generate radio maps across a range of frequencies (700,MHz to 6,GHz).

Maxwell's Equations and Computational Electromagnetics

“From a long view of the history of mankind, seen from, say, ten thousand years from now, there can be little doubt that the most significant event of the 19th century will be judged as Maxwell's discovery of the laws of electrodynamics.” - Richard Feynman

Maxwell's equations unify electricity, magnetism, and optics into a single theoretical framework. Their numerical solution-known as computational electromagnetics (CEM)-underpins the design of antennas, optical devices, integrated circuits, MRI scanners, and the wireless networks through which you are likely accessing this textbook.

Maxwell's equations: A brief recap

In their differential form, Maxwell's equations in a linear, isotropic medium are:

For time-harmonic fields ($𝐄(𝐱,t)=𝐄(𝐱)e^{-i\omega t}$), Maxwell's equations reduce to the vector Helmholtz equation: (Vector Helmholtz)$$\begin{array}{ccc}& \nabla \times ({\mu }^{-1}\nabla \times 𝐄)-{\omega }^2\epsilon 𝐄=i\omega 𝐉,& \text{<span class="tml-eqn"></span>}\end{array}$$ connecting the electromagnetic problem to the Helmholtz theory of the previous section.

GNOT for Maxwell's equations

The General Neural Operator Transformer (GNOT; Hao et al., 2023) was evaluated on several electromagnetic benchmarks, including the Inductor2D problem: predicting the magnetic field distribution around a 2D cross-section of an inductor coil, given the coil geometry and current distribution.

Diffusion for electromagnetic inverse design

One of the most exciting applications of generative models in electromagnetics is inverse design: given a desired electromagnetic response (e.g. a specific scattering pattern, a prescribed transmission spectrum), generate the physical structure (geometry and material distribution) that achieves it.

This is a classic ill-posed inverse problem: many structures may produce similar electromagnetic responses, and the mapping from response to structure is highly nonlinear. Diffusion models are ideally suited because they naturally handle multi-modality (generating diverse valid designs) and can incorporate physics constraints.

A conditional diffusion model $p_{\varphi }(𝐬|{𝐫}^{\ast })$ is trained on paired (design, response) data. At inference, given a desired response ${𝐫}^{\ast }$, the model generates multiple candidate designs, which are then verified with a single forward simulation each (much cheaper than iterative optimisation).

Physics-aware Maxwell-guided diffusion

A 2026 advance incorporates Maxwell's equations directly into the diffusion process. Rather than training a diffusion model on simulation data and hoping it learns the physics, this approach uses the Maxwell residual as an explicit guidance signal during denoising.

This approach has two advantages: (i) it improves physical consistency without requiring additional training data, and (ii) it enables zero-shot generalisation to new material configurations (different $\epsilon$, $\mu$) not seen during training, since the physics guidance automatically adapts.

Applications in photonics, antenna design, and metamaterials

Photonic crystal design

Photonic crystals are periodic dielectric structures that control light propagation. Designing a photonic crystal with a desired band gap requires solving Maxwell's equations for many candidate geometries-a perfect task for generative surrogates. Diffusion models conditioned on the target band structure can generate diverse photonic crystal unit cells in milliseconds.

Antenna design

Modern antenna design (5G MIMO arrays, satellite antennas) involves optimising complex 3D geometries to achieve specified radiation patterns, impedance matching, and bandwidth. Generative models can explore the vast design space more efficiently than gradient-based topology optimisation, especially when multiple competing objectives must be balanced.

Metamaterials

Metamaterials are engineered structures with electromagnetic properties not found in nature (negative refractive index, cloaking, perfect absorption). The design space is combinatorially vast (each meta-atom can have a different geometry), and diffusion models provide a natural way to sample from the space of designs that achieve a target effective property.

Stefan Problem and Moving Boundaries

When ice melts in a glass of water, the boundary between solid and liquid is not fixed-it moves, and its position is itself an unknown of the problem. This deceptively simple observation underlies one of the oldest and most challenging classes of PDE problems: the Stefan problem and its many generalisations.

The classic Stefan problem

Consider a one-dimensional slab of material occupying $[0,L]$ that undergoes a phase transition (e.g., ice melting into water). Let $s(t)$ denote the position of the phase boundary at time $t$. In the liquid region $0<x<s(t)$ the temperature $u(x,t)$ satisfies the heat equation, and in the solid region $s(t)<x<L$ the temperature $v(x,t)$ satisfies a (possibly different) heat equation: (Liquid)$$\begin{array}{cccc}& \frac{\partial u}{\partial t}& ={\alpha }_{\ell }\frac{{\partial }^{2}u}{\partial x^2},0<x<s(t),& \text{<span class="tml-eqn"></span>}\\ & \frac{\partial v}{\partial t}& ={\alpha }_s\frac{{\partial }^{2}v}{\partial x^2},s(t)<x<L,& \text{<span class="tml-eqn"></span>}\end{array}$$ where ${\alpha }_{\ell }$ and ${\alpha }_s$ are the thermal diffusivities of the liquid and solid phases, respectively.

At the interface $x=s(t)$, the temperature equals the melting point $T_m$: (Melting)$$\begin{array}{ccc}& u(s(t),t)=v(s(t),t)=T_m.& \text{<span class="tml-eqn"></span>}\end{array}$$ The crucial ingredient is the Stefan condition, which links the interface velocity to the jump in heat flux:

The Stefan condition says that any imbalance in heat flux at the interface must go into moving the boundary-absorbing latent heat as ice melts, or releasing it as water freezes. This is an elegant expression of energy conservation, but it makes the mathematical problem nonlinear: the domain on which we solve the PDE is itself determined by the solution.

Why moving boundaries are hard for neural methods

Moving boundary problems present a triple challenge for neural PDE solvers:

1. Domain geometry is solution-dependent. In standard PINN or neural operator approaches, the domain $\mathrm{\Omega }$ is fixed. For a Stefan problem, the domain changes at every time step, and the domain boundary is itself an unknown. A neural network that takes $(x,t)$ as input has no built-in mechanism to “know” that $x$ lies in the liquid phase for $x<s(t)$ and in the solid phase for $x>s(t)$.

2. Discontinuous gradients at the interface. The temperature is continuous across the interface (it equals $T_m$ on both sides), but its gradient may be discontinuous-indeed, the difference in gradients is what drives the interface motion. Standard MLPs with smooth activations produce smooth outputs and struggle to represent gradient discontinuities sharply.

3. Implicit coupling. The Stefan condition couples the PDE solution to the interface evolution in a nonlinear, implicit fashion. Enforcing this condition as a soft penalty in a loss function leads to a delicate balance between PDE residual accuracy and interface tracking accuracy.

Physics-Informed Distillation of Diffusion Models (PIDDM)

A fundamental tension arises when we attempt to enforce PDE constraints within a diffusion model. Recall that a trained diffusion model learns the conditional expectation $𝔼[{𝐱}_0|{𝐱}_t]$ through its denoiser. But for a nonlinear PDE operator $𝒩$, enforcing the physics on this conditional mean is not the same as enforcing it on individual samples: (Jensen)$$\begin{array}{ccc}& 𝒩[𝔼[{𝐱}_0|{𝐱}_t]]\ne 𝔼[𝒩[{𝐱}_0]|{𝐱}_t],& \text{<span class="tml-eqn"></span>}\end{array}$$ by Jensen's inequality (applied to the convex/concave parts of $𝒩$). This discrepancy is called the Jensen's gap, and it means that a physics-informed diffusion model may produce samples whose mean satisfies the PDE while individual samples do not.

Post-hoc distillation to close the gap

The Physics-Informed Distillation of Diffusion Models (PIDDM) approach addresses the Jensen's gap through a two-stage strategy:

1. Stage 1: Train a standard diffusion model. A conditional diffusion model is trained on paired data $({\theta }_i,u_i)$ where ${\theta }_i$ are PDE parameters (initial conditions, coefficients) and $u_i$ are solution fields obtained from a numerical solver. This produces a teacher model that can generate diverse, data-consistent solution fields.

2. Stage 2: Distil into a single-step student. A deterministic student network $G_{\psi }(\theta ,𝐳)$ (where $𝐳\sim 𝒩(0,I)$) is trained to match the teacher's output distribution while simultaneously satisfying the PDE constraints on individual samples: (LOSS)$$\begin{array}{ccc}& {ℒ}_{\text{PIDDM}}(\psi )=\underset{\text{match teacher}}{\underbrace{{ℒ}_{\text{distill}}(\psi )}}+{\lambda }_{\text{phys}}\underset{\text{physics on samples}}{\underbrace{{ℒ}_{\text{PDE}}(\psi )}},& \text{<span class="tml-eqn"></span>}\end{array}$$ where the physics loss is evaluated on the student's actual outputs $\hat{u}=G_{\psi }(\theta ,𝐳)$: (Physics)$$\begin{array}{ccc}& {ℒ}_{\text{PDE}}(\psi )={𝔼}_{\theta ,𝐳}[\Vert 𝒩[G_{\psi }(\theta ,𝐳)]{\Vert }^2].& \text{<span class="tml-eqn"></span>}\end{array}$$

Because the student generates samples in a single forward pass (no iterative denoising), we can compute PDE residuals directly on each generated sample, completely avoiding the Jensen's gap.

Application to the Generalised Porous Medium Equation

The Generalised Porous Medium Equation (GPME) provides an excellent test bed for PIDDM because it exhibits moving fronts (compactly supported solutions with sharp interfaces) that require accurate tracking: (GPME)$$\begin{array}{ccc}& \frac{\partial u}{\partial t}=\nabla \cdot (m{u}^{m-1}\nabla u),m>1.& \text{<span class="tml-eqn"></span>}\end{array}$$ For $m=1$ this reduces to the heat equation; for $m>1$ the diffusion is degenerate (vanishing where $u=0$), producing solutions with compact support whose boundaries move at finite speed. The case $m=2$ arises in modelling gas flow through porous media, while $m=3$ appears in thin-film equations.

Cahn–Hilliard and Allen–Cahn: Phase-Field Dynamics

The Stefan problem of the previous section tracks a sharp interface between phases. An alternative approach, enormously influential in materials science and increasingly in machine learning, replaces the sharp interface with a diffuse interface-a thin transition layer of width $ϵ$ across which the order parameter varies smoothly from one phase to the other. This is the realm of phase-field models.

The Allen–Cahn equation

The Allen–Cahn equation describes the evolution of a non-conserved order parameter $u(𝐱,t)\in [-1,+1]$, where $u=+1$ represents one phase and $u=-1$ the other: (AC)$$\begin{array}{ccc}& \frac{\partial u}{\partial t}={ϵ}^2{\nabla }^{2}u-F^{\prime }(u),& \text{<span class="tml-eqn"></span>}\end{array}$$ where $ϵ>0$ is the interfacial width parameter and $F(u)=\frac{1}{4}{(u^2-1)}^2$ is the standard double-well potential with minima at $u=\pm 1$.

The two terms in the free energy compete: the gradient term $\frac{{ϵ}^2}{2}|\nabla u{|}^2$ penalises rapid spatial variation (favouring smooth, broad interfaces), while the double-well potential $F(u)$ penalises values away from $\pm 1$ (favouring sharp phase separation). The balance between these two tendencies produces interfaces of characteristic width $\sim ϵ$.

The Cahn–Hilliard equation

When the total amount of each phase must be conserved (as in a binary alloy or a polymer blend), the Allen–Cahn equation is inadequate because it does not preserve ${\int }_{\mathrm{\Omega }}ud𝐱$. The Cahn–Hilliard equation enforces conservation by making the dynamics a conserved gradient flow: (CH)$$\begin{array}{ccc}& \frac{\partial u}{\partial t}=\nabla \cdot (M\nabla (F^{\prime }(u)-{ϵ}^2{\nabla }^{2}u\left)\right),& \text{<span class="tml-eqn"></span>}\end{array}$$ where $M>0$ is the mobility. The quantity in parentheses, (Chempot)$$\begin{array}{ccc}& \mu =F^{\prime }(u)-{ϵ}^2{\nabla }^{2}u=\frac{\delta ℱ}{\delta u},& \text{<span class="tml-eqn"></span>}\end{array}$$ is the chemical potential, and the Cahn–Hilliard equation can be written compactly as ${\partial }_{t}u=\nabla \cdot (M\nabla \mu )$.

Spinodal decomposition

The Cahn–Hilliard equation produces one of the most visually striking phenomena in materials science: spinodal decomposition. Starting from a nearly homogeneous mixture (e.g., $u(𝐱,0)=\bar{u}+\eta (𝐱)$ where $\eta$ is small random noise and $|\bar{u}|<1/\sqrt{3}$ lies in the spinodal region where $F^{\prime \prime }(\bar{u})<0$), the system is linearly unstable: small perturbations grow exponentially.

The result is a cascade of phase separation: the mixture rapidly demixes into a labyrinthine pattern of domains rich in one component or the other, connected by thin interfaces. Over long times, these domains coarsen via the Lifshitz–Slyozov mechanism (larger domains grow at the expense of smaller ones), with the characteristic domain size growing as $\ell (t)\sim t^{1/3}$.

Why fourth-order derivatives kill PINNs

The Cahn–Hilliard equation is fourth-order in space: expanding the divergence term reveals ${\nabla }^{4}u$ (the biharmonic operator). This creates severe difficulties for physics-informed approaches:

1. Gradient accumulation. Computing ${\nabla }^4{u}_{\varphi }$ via automatic differentiation requires four nested applications of the chain rule. Each pass through the network multiplies the computational graph, and the memory cost scales roughly as $𝒪(4\cdot \text{depth}\cdot \text{width})$ for the backward pass. For a network with $L$ layers of width $W$, this translates to a per-collocation-point cost of $𝒪(L^{4}W)$ in the worst case.

2. Numerical ill-conditioning. Fourth-order derivatives amplify high-frequency noise in the network output. Small oscillations of amplitude $\delta$ at wavenumber $k$ produce fourth-derivative residuals of order $\delta k^4$, leading to loss-function values dominated by numerical artefacts rather than genuine PDE violation.

3. Loss landscape pathology. The ratio between the PDE residual scale (involving ${\nabla }^4$) and the boundary/initial condition scale creates extreme ill-conditioning in the loss landscape, making gradient descent ineffective without careful normalisation.

PINTO: Physics-Informed Neural Transformer Operator

Standard neural operators such as FNO and DeepONet learn the solution operator from data alone, without explicit PDE constraints. For phase-field equations, where the physics imposes strict conservation laws and energy dissipation, purely data-driven operators can produce solutions that violate these fundamental properties.

The Physics-Informed Neural Transformer Operator (PINTO) addresses this by combining the expressiveness of Transformers with physics-informed training.

Architecture

PINTO uses an encoder-decoder architecture where the key innovation is the use of cross-attention integral kernels to approximate the integral transforms that appear in operator learning:

Generalisation across initial conditions

A key advantage of PINTO over single-instance PINNs is amortised inference: once trained, the model maps any initial condition to the solution at a future time without re-training. For the Cahn–Hilliard equation, this means a single PINTO model can predict spinodal decomposition patterns starting from different compositions $\bar{u}$, different noise realisations, and even different domain sizes (if position encoding is appropriately designed).

Cahn–Hilliard for image inpainting

In a beautiful example of physics inspiring computer vision, the Cahn–Hilliard equation has been applied to image inpainting: the task of filling in missing or corrupted regions of an image.

The key insight is that image pixels, when viewed as a field $u(𝐱)$, can undergo “phase separation” that propagates structure from known regions into unknown ones. A wavelet decomposition separates the image into scale-specific components, and Cahn–Hilliard dynamics applied to each wavelet sub-band propagates texture information from the boundary of the missing region inward, respecting the natural scale hierarchy of the image.

Schrödinger Equation and Quantum PDEs

We now venture into perhaps the most consequential application of generative models to PDEs: solving the quantum many-body Schrödinger equation. The stakes are enormous-an accurate, scalable quantum solver would revolutionise chemistry, materials science, and drug design. The difficulty is equally enormous: the curse of dimensionality is at its most severe in quantum mechanics, where the wavefunction of $N$ electrons lives in a $3N$-dimensional space.

The time-independent Schrödinger equation

For a system of $N$ electrons in a fixed external potential (the Born–Oppenheimer approximation), the time-independent electronic Schrödinger equation reads: (TISE)$$\begin{array}{ccc}& \hat{H}\mathrm{\Psi }({𝐫}_1,\dots ,{𝐫}_N)=E\mathrm{\Psi }({𝐫}_1,\dots ,{𝐫}_N),& \text{<span class="tml-eqn"></span>}\end{array}$$ where ${𝐫}_i\in {ℝ}^3$ is the position of the $i$-th electron (we suppress spin for now), $E$ is the total electronic energy, and the Hamiltonian is: (Hamiltonian)$$\begin{array}{ccc}& \hat{H}=-\frac{1}{2}\sum _{i=1}^N{\nabla }_i^2+\sum _{i=1}^N{V}_{\text{ext}}({𝐫}_i)+\sum _{i<j}\frac{1}{|{𝐫}_i-{𝐫}_j|}.& \text{<span class="tml-eqn"></span>}\end{array}$$ The three terms represent kinetic energy, electron-nuclear attraction, and electron-electron repulsion, respectively. We use atomic units ($\hslash =m_e=e=1$).

The curse of dimensionality in Hilbert space

The wavefunction $\mathrm{\Psi }({𝐫}_1,\dots ,{𝐫}_N)$ is a function of $3N$ continuous variables. For a modestly sized molecule with $N=50$ electrons, this is a function in 150 dimensions. If we discretise each dimension with $G$ grid points, storing the wavefunction requires $G^{3N}$ numbers-a quantity that exceeds the number of atoms in the observable universe for $G=10$ and $N\ge 30$.

Antisymmetry and Slater determinants

Electrons are fermions, and the Pauli exclusion principle demands that the wavefunction be antisymmetric under exchange of any two electrons (including their spin): (Antisym)$$\begin{array}{ccc}& \mathrm{\Psi }(\dots ,{𝐫}_i,{\sigma }_i,\dots ,{𝐫}_j,{\sigma }_j,\dots )=-\mathrm{\Psi }(\dots ,{𝐫}_j,{\sigma }_j,\dots ,{𝐫}_i,{\sigma }_i,\dots ),& \text{<span class="tml-eqn"></span>}\end{array}$$ where ${\sigma }_i\in \{\uparrow ,\downarrow \}$ denotes spin. This constraint is non-negotiable: any ansatz for $\mathrm{\Psi }$ that does not satisfy will produce unphysical results.

The simplest antisymmetric ansatz is the Slater determinant: (Slater)$$\begin{array}{ccc}& {\mathrm{\Psi }}_{\text{SD}}({𝐫}_1,\dots ,{𝐫}_N)=\frac{1}{\sqrt{N!}}\left|\begin{array}{cccc}{\varphi }_1({𝐫}_1)& {\varphi }_1({𝐫}_2)& \cdots & {\varphi }_1({𝐫}_N)\\ {\varphi }_2({𝐫}_1)& {\varphi }_2({𝐫}_2)& \cdots & {\varphi }_2({𝐫}_N)\\ ⋮& ⋮& \ddots & ⋮\\ {\varphi }_N({𝐫}_1)& {\varphi }_N({𝐫}_2)& \cdots & {\varphi }_N({𝐫}_N)\end{array}\right|,& \text{<span class="tml-eqn"></span>}\end{array}$$ where $\{{\varphi }_i\}$ are single-particle orbitals. The determinant structure automatically ensures antisymmetry (swapping two columns changes the sign), but a single Slater determinant captures only mean-field (Hartree–Fock) physics, missing the crucial electron correlation that determines chemical bonding, reaction barriers, and materials properties.

FermiNet and PauliNet: pioneering neural wavefunctions

Two landmark architectures showed that neural networks could achieve near-exact accuracy for the electronic Schrödinger equation:

FermiNet (Pfau et al., 2020) parametrises $\mathrm{\Psi }$ as a sum of $D$ generalised Slater determinants: (Ferminet)$$\begin{array}{ccc}& {\mathrm{\Psi }}_{\theta }({𝐫}_1,\dots ,{𝐫}_N)=\sum _{d=1}^D{w}_d\det [{\varphi }_i^{(d)}({𝐫}_j;\{{𝐫}_k{\}}_{k\ne j}){]}_{i,j=1}^N,& \text{<span class="tml-eqn"></span>}\end{array}$$ where each orbital ${\varphi }_i^{(d)}$ depends not only on the position ${𝐫}_j$ of its assigned electron but also on the positions of all other electrons ${\{{𝐫}_k\}}_{k\ne j}$. This “backflow” dependence is the key innovation: it allows each orbital to adapt to the instantaneous configuration of the other electrons, capturing correlation effects far beyond Hartree–Fock.

PauliNet (Hermann, Schätzle, and Noé, 2020) takes a complementary approach, building on top of traditional quantum chemistry by starting from Hartree–Fock orbitals and adding neural correlation factors: $${\mathrm{\Psi }}_{\theta }={\mathrm{\Psi }}_{\text{HF}}\cdot e^{J_{\theta }}\cdot \det [{\varphi }_i^{\text{HF}}({𝐫}_j)+f_{ij}^{\theta }({𝐫}_j,\{{𝐫}_k\})],$$ where $J_{\theta }$ is a neural Jastrow factor and $f_{ij}^{\theta }$ are learned backflow corrections. By starting from the Hartree–Fock solution, PauliNet incorporates physical knowledge (cusp conditions, asymptotic decay) from the outset.

QiankunNet: Transformers for many-electron systems

While FermiNet and PauliNet demonstrated the viability of neural wavefunctions, they relied on hand-designed equivariant architectures with limited scalability. QiankunNet (Zhao et al., Nature Communications, 2025) introduces a Transformer-based architecture that scales to much larger systems, achieving chemical accuracy for challenging transition-metal complexes.

Architecture overview

The name “Qiankun” (meaning “universe” in Chinese) reflects the ambition to capture universal quantum correlations through attention mechanisms. The architecture consists of three key components:

1. Configuration-dependent orbitals via attention. Each electron $i$ attends to all other electrons and all nuclei through multi-head self-attention: (ATTN)$$\begin{array}{ccc}& h_i^{(\ell +1)}=h_i^{(\ell )}+\text{MHA}(h_i^{(\ell )},\{h_j^{(\ell )}{\}}_{j=1}^N,\{R_A{\}}_{A=1}^M),& \text{<span class="tml-eqn"></span>}\end{array}$$ where $\{R_A\}$ are nuclear positions and MHA denotes multi-head attention. The attention heads implement backflow correlations: each electron's representation is modulated by the positions of all other electrons, analogous to backflow transformations in quantum Monte Carlo.

2. Batched autoregressive sampling. To evaluate the variational energy by Monte Carlo integration, one needs to sample electron configurations $({𝐫}_1,\dots ,{𝐫}_N)\sim |{\mathrm{\Psi }}_{\theta }{|}^2$. QiankunNet uses a batched autoregressive approach: electrons are sampled sequentially, with each electron's proposal distribution conditioned on the positions of all previously sampled electrons. Batching across multiple configurations enables efficient GPU utilisation.

3. Generalised determinantal structure. The final wavefunction takes the form: (PSI)$$\begin{array}{ccc}& {\mathrm{\Psi }}_{\theta }({𝐫}_1,\dots ,{𝐫}_N)=\sum _{d=1}^D{w}_d(\{𝐫\})\det [{\tilde{\varphi }}_i^{(d)}({𝐫}_j;h_j^{(L)}){]}_{i,j},& \text{<span class="tml-eqn"></span>}\end{array}$$ where ${\tilde{\varphi }}_i^{(d)}$ are orbitals constructed from the final-layer Transformer embeddings $h_j^{(L)}$, and even the determinant weights $w_d$ are configuration-dependent.

Achieving chemical accuracy

Chemical accuracy-defined as an error of at most 1 kcal/mol ($\approx 1.6\times {10}^{-3}$ Hartree) relative to experiment-is the gold standard in quantum chemistry. QiankunNet achieves or approaches this threshold on systems that are extremely challenging for traditional methods:

Hellmann–Feynman forces without Pulay terms

A remarkable practical advantage of neural wavefunctions is that molecular forces (derivatives of the energy with respect to nuclear positions) simplify dramatically. The Hellmann–Feynman theorem states:

For approximate wavefunctions (such as those from Hartree–Fock or truncated CI), the theorem does not hold exactly, and additional “Pulay force” terms arising from the basis set dependence on nuclear positions must be computed. For neural wavefunctions, however, the basis set is the neural network itself, which depends on nuclear positions only through the input features. When the variational energy is well-converged, the Pulay terms become negligible, and forces can be computed cheaply from the Hellmann–Feynman expression alone.

Flow matching for ground-state sampling

An emerging approach replaces Variational Monte Carlo (VMC) with flow matching for sampling the ground-state distribution $|{\mathrm{\Psi }}_0{|}^2$. The idea is elegant: learn a continuous normalising flow that transforms a simple base distribution (e.g., independent Gaussians centred on nuclei) into the target distribution $|{\mathrm{\Psi }}_{\theta }{|}^2$: (FLOW)$$\begin{array}{ccc}& \frac{d𝐫}{dt}=v_{\varphi }(𝐫,t),𝐫(0)\sim p_{\text{base}},𝐫(1)\sim |{\mathrm{\Psi }}_{\theta }{|}^2.& \text{<span class="tml-eqn"></span>}\end{array}$$

The flow matching loss is: (Flowloss)$$\begin{array}{ccc}& {ℒ}_{\text{FM}}(\varphi )={𝔼}_{t,{𝐫}_0,{𝐫}_1}[\Vert v_{\varphi }({𝐫}_t,t)-({𝐫}_1-{𝐫}_0){\Vert }^2],& \text{<span class="tml-eqn"></span>}\end{array}$$ where ${𝐫}_t=(1-t){𝐫}_0+t{𝐫}_1$ is the linear interpolation. Importantly, the flow must be equivariant under electron permutations to respect the fermionic symmetry.

Coupled and Multiphysics PDE Systems

Real-world engineering systems rarely involve a single PDE\@. An aircraft wing experiences aerodynamic forces (Navier–Stokes), which deform the structure (elasticity), which changes the aerodynamic forces, which further deform the structure-a fully coupled fluid–structure interaction problem. A nuclear reactor core involves neutron transport, heat conduction, coolant flow, and structural mechanics, all coupled through temperature-dependent cross sections and thermal expansion. A battery cell involves electrochemistry, ion transport, heat generation, and mechanical stress from lithium intercalation.

The coupling challenge

Multiphysics systems are difficult for three fundamental reasons:

1. Scale separation. Different physics operate on vastly different spatial and temporal scales. In a turbine blade, the fluid boundary layer has thickness $\sim 0.1$ mm and fluctuates on microsecond timescales, while the structural vibration has a wavelength of $\sim 10$ cm and a period of $\sim 1$ ms. A single mesh or time step that resolves both is prohibitively expensive.

2. Stiffness. The coupling between physics often introduces stiff dynamics: some modes decay much faster than others. In thermoelasticity, thermal diffusion ($\sim {10}^{-5}$ m${}^2$/s) is much slower than elastic wave propagation ($\sim {10}^3$ m/s), leading to a stiffness ratio of $\sim {10}^8$ in the governing system of ODEs.

3. Interface conditions. At the boundary between different physics domains (e.g., the fluid–solid interface in an FSI problem), compatibility conditions must be satisfied: continuity of displacement, traction balance, heat flux continuity. These conditions couple the sub-problems and must be enforced to high accuracy; even small violations can lead to instabilities or unphysical energy generation.

MultiSimDiff: Compositional Generative Multiphysics

MultiSimDiff (Wei et al., ICML 2025) introduces a beautifully modular approach to multiphysics simulation using the framework of compositional diffusion models.

Core idea: composing energy functions

The key insight is that in diffusion models, the score function ${\nabla }_{𝐱}\log p(𝐱)$ can be decomposed as a sum of contributions from independent factors:

MultiSimDiff exploits this by training a separate diffusion model (equivalently, a separate score function or energy function) for each physics discipline, and composing them at inference time to generate solutions to the coupled system.

Training on decoupled data

The training procedure is strikingly simple:

1. For each physics $k$ (e.g., neutron transport, thermal hydraulics), independently generate training data by running the corresponding single-physics solver.

2. Train a conditional diffusion model ${ϵ}_{{\theta }_k}$ for each physics, conditioned on the relevant boundary conditions and parameters.

3. At inference, compose the score functions: (Composed)$$\begin{array}{ccc}& {ϵ}_{\text{coupled}}({𝐱}_t,t)=\sum _{k=1}^K{ϵ}_{{\theta }_k}({𝐱}_t,t).& \text{<span class="tml-eqn"></span>}\end{array}$$

The remarkable aspect is that no coupled simulation data is needed during training. Each model is trained independently on single-physics data, and the coupling emerges purely from the compositional structure of diffusion models at inference time.

Nuclear thermal coupling: results

MultiSimDiff was demonstrated on a nuclear reactor thermal coupling problem involving neutron diffusion coupled with heat conduction in a prismatic fuel element with up to 64 material components.

Transformer operators for multiphysics

Beyond compositional diffusion, several Transformer-based architectures have been proposed for directly learning multiphysics operators:

GNOT: multi-input operator Transformers

The General Neural Operator Transformer (GNOT) handles heterogeneous inputs naturally by treating different physical fields, boundary conditions, and PDE coefficients as distinct token types. Each token type has its own embedding, and cross-attention between types allows the model to learn input-output relationships across physics disciplines: (GNOT)$$\begin{array}{ccc}& h_{\text{out}}^{(k)}=\text{CrossAttn}(Q=h_{\text{query}}^{(k)},K=[h_{\text{in}}^{(1)},\dots ,h_{\text{in}}^{(P)}],V=[h_{\text{in}}^{(1)},\dots ,h_{\text{in}}^{(P)}]),& \text{<span class="tml-eqn"></span>}\end{array}$$ where $P$ is the number of input physics fields.

HAMLET: modular encoders for multiphysics

HAMLET (Heterogeneous Architecture with Multiple Learned Encoders for Transformers) uses physics-specific encoders that project each field into a shared latent space: (Hamlet)$$\begin{array}{ccc}& z^{(k)}={\text{Encoder}}_k(u^{(k)}),{\hat{u}}^{(k)}={\text{Decoder}}_k(\text{Transformer}([z^{(1)},\dots ,z^{(P)}])).& \text{<span class="tml-eqn"></span>}\end{array}$$ The modular encoder-decoder structure allows each physics module to be designed, trained, and tested independently, with the central Transformer learning the cross-physics interactions.

CoDA-NO: cross-system pretraining

The Codomain Attention Neural Operator (CoDA-NO) takes the transferability idea further by pretraining on a diverse collection of PDE systems (heat equation, advection, Burgers, Navier–Stokes) and fine-tuning on the target multiphysics system. The key architectural choice is codomain attention: instead of attending over spatial positions (as in standard Transformers), CoDA-NO attends over the output channels (solution components), which allows it to learn cross-physics correlations:

Summary of multiphysics approaches

Table Table 3 summarises the key characteristics of the multiphysics methods discussed.

Method	Coupled data	Modular	Generative
MultiSimDiff	Not needed	Yes	Yes (diffusion)
GNOT	Required	Partially	No
HAMLET	Required	Yes	No
CoDA-NO	Fine-tuning only	Yes	No

Function-Space Diffusion: Beyond Grids

Every diffusion model we have encountered so far operates on finite-dimensional objects: a vorticity field discretised on a $64\times 64$ grid, a velocity snapshot stored as a $128\times 128\times 2$ tensor, a temperature profile sampled at $N$ sensor locations. The forward and reverse SDEs live in ${ℝ}^n$ for some large but fixed $n$, and the score network is a finite-dimensional function $s_{\varphi }:{ℝ}^n\times [0,T]\to {ℝ}^n$.

This is fine-until you change the grid.

The remedy is to formulate diffusion directly on the function space to which the PDE solution belongs-typically a separable Hilbert space $\mathscr{H}$ such as $L^2(\mathrm{\Omega })$ or a Sobolev space $H^s(\mathrm{\Omega })$. In this formulation, each “data point” is an entire function $u\in \mathscr{H}$, the noise process is an $\mathscr{H}$-valued stochastic process, and the score is a map $s:\mathscr{H}\times [0,T]\to \mathscr{H}$. Any finite-dimensional grid merely provides a projection of this infinite-dimensional object, and a model trained in function space can be evaluated on any discretisation at test time.

Mathematical framework: diffusion on Hilbert spaces

Let $\mathscr{H}$ be a separable Hilbert space with inner product ${⟨\cdot ,\cdot ⟩}_{\mathscr{H}}$ and norm $\Vert \cdot {\Vert }_{\mathscr{H}}$. The prototypical choice is $\mathscr{H}=L^2(\mathrm{\Omega })$ for a bounded domain $\mathrm{\Omega }\subset {ℝ}^d$. We denote by ${\{e_k\}}_{k=1}^{\infty }$ a complete orthonormal basis of $\mathscr{H}$.

In practice, we work with a trace-class noise process by introducing a covariance operator $𝒞:\mathscr{H}\to \mathscr{H}$ that is positive, self-adjoint, and trace-class ($\mathrm{Tr}(𝒞)<\infty$). The $𝒞$-Wiener process is defined by (Cwiener)$$\begin{array}{ccc}& W^{𝒞}(t)=\sum _{k=1}^{\infty }\sqrt{{\lambda }_k}{\beta }_k(t)e_k,& \text{<span class="tml-eqn"></span>}\end{array}$$ where ${\lambda }_k$ are the eigenvalues of $𝒞$ and the eigenvectors coincide with the basis $\{e_k\}$ (which we may always arrange by diagonalising $𝒞$). The trace-class condition ${\sum }_k{\lambda }_k<\infty$ ensures that $W^{𝒞}(t)\in \mathscr{H}$ almost surely.

The marginal distribution at time $t$ has the form (Funcmarginal)$$\begin{array}{ccc}& u_t=\alpha (t)u_0+\sigma (t)\xi ,\xi \sim 𝒩(0,𝒞),& \text{<span class="tml-eqn"></span>}\end{array}$$ where $\alpha (t)$ and $\sigma (t)$ are scalar functions determined by $\beta (\cdot )$ and $𝒜$, and $𝒩(0,𝒞)$ denotes a Gaussian measure on $\mathscr{H}$ with covariance $𝒞$. The reverse-time SDE takes the form (Funcreverse)$$\begin{array}{ccc}& \mathrm{d}{u}_t=[-\frac{1}{2}\beta (t)𝒜u_t-\beta (t)𝒞{\nabla }_{\mathscr{H}}\log p_t(u_t)]\mathrm{d}t+\sqrt{\beta (t)}\mathrm{d}{\bar{W}}_t^{𝒞},& \text{<span class="tml-eqn"></span>}\end{array}$$ where ${\nabla }_{\mathscr{H}}\log p_t$ is the function-space score-the Fréchet derivative of the log-density with respect to the $\mathscr{H}$-inner product-and ${\bar{W}}^{𝒞}$ is a backward $𝒞$-Wiener process.

FunDPS: guided diffusion on function spaces

A major attraction of diffusion models is their ability to solve inverse problems via posterior sampling: given an observation $y=ℱ(u)+\eta$ of a PDE field $u$ through a (possibly nonlinear) forward operator $ℱ$ with noise $\eta$, we wish to sample from $p(u|y)$. In finite dimensions, this is achieved by Diffusion Posterior Sampling (DPS), which approximates the posterior score as (DPS Finite)$$\begin{array}{ccc}& {\nabla }_{𝐱}\log p({𝐱}_t|y)\approx {\nabla }_{𝐱}\log p_t({𝐱}_t)-{\rho }_t{\nabla }_{𝐱}\Vert y-ℱ({\widehat{𝐱}}_0){\Vert }^2,& \text{<span class="tml-eqn"></span>}\end{array}$$ where ${\widehat{𝐱}}_0$ is the Tweedie posterior mean.

FunDPS (Cardoso et al., NeurIPS 2025) extends this to function spaces. The key insight is that the Tweedie formula generalises naturally:

FunDPS uses this to construct a function-space guidance term: (Fundps)$$\begin{array}{ccc}& {\nabla }_{\mathscr{H}}\log p_t(u_t|y)\approx s_{\varphi }(u_t,t)-{\rho }_t{\nabla }_{\mathscr{H}}\Vert y-ℱ({\hat{u}}_0(u_t)){\Vert }_{𝒴}^2,& \text{<span class="tml-eqn"></span>}\end{array}$$ where $𝒴$ is the observation space and ${\nabla }_{\mathscr{H}}$ denotes the Fréchet derivative with respect to $u_t$.

The method is discretisation-agnostic: the unconditional model and the guidance gradient can be evaluated on different grids, and the posterior samples converge to the same continuous function regardless of how each component is discretised, up to discretisation error.

FunDiff: diffusion with function autoencoders

While FunDPS works directly in the ambient function space, an alternative strategy is to first compress the function into a latent representation via a learned encoder, perform diffusion in this latent space, and then decode back to function space. FunDiff (Lim et al., 2024) takes this approach with a carefully designed function autoencoder.

Function autoencoder architecture

The encoder maps a function $u\in \mathscr{H}$ to a finite-dimensional latent code $z\in {ℝ}^{d_z}$, and the decoder maps back: (Funcae)$$\begin{array}{ccc}& z={\mathrm{Enc}}_{\psi }(u),\tilde{u}={\mathrm{Dec}}_{\psi }(z).& \text{<span class="tml-eqn"></span>}\end{array}$$ Crucially, both the encoder and decoder are neural operators: the encoder takes function evaluations at an arbitrary set of input points and produces a fixed-dimensional code via cross-attention or integral transforms, and the decoder produces function evaluations at any queried set of output points. This makes the autoencoder itself discretisation-invariant.

Physics-informed priors

FunDiff incorporates physics via a regularisation term in the autoencoder loss: (Fundiff LOSS)$$\begin{array}{ccc}& {ℒ}_{\mathrm{AE}}(\psi )={𝔼}_{u\sim p_{\mathrm{data}}}[\Vert u-{\mathrm{Dec}}_{\psi }({\mathrm{Enc}}_{\psi }(u)){\Vert }_{\mathscr{H}}^2+\lambda \Vert 𝒩[{\mathrm{Dec}}_{\psi }({\mathrm{Enc}}_{\psi }(u))]{\Vert }^2],& \text{<span class="tml-eqn"></span>}\end{array}$$ where $𝒩[u]=0$ is the PDE constraint. This encourages the decoded functions to satisfy the governing equations, even when the latent diffusion introduces small artefacts.

This is a remarkable result: the convergence rate depends on the regularity $s$ of the target functions and the intrinsic dimensionality $d$ of the domain, not on the number of grid points used for discretisation.

Score-based diffusion in function space: the JMLR framework

The most comprehensive theoretical treatment of function-space diffusion to date is the framework of Pidstrigach, Berner, Richter, and Schönlieb (JMLR 2025). We summarise the key contributions.

Infinite-dimensional SDE formulation

The authors work with the SPDE (stochastic partial differential equation) (Funcjmlr Forward)$$\begin{array}{ccc}& \mathrm{d}{u}_t=-\frac{1}{2}\beta (t){(-\mathrm{\Delta })}^{\alpha /2}{u}_t\mathrm{d}t+\sqrt{\beta (t)}\mathrm{d}{W}_t^{𝒞},& \text{<span class="tml-eqn"></span>}\end{array}$$ where ${(-\mathrm{\Delta })}^{\alpha /2}$ is the fractional Laplacian and $𝒞={(-\mathrm{\Delta })}^{-\alpha }$ for $\alpha >d/2$. This choice ensures that:

1. The noise $W_t^{𝒞}$ takes values in $\mathscr{H}$ (the trace-class condition is satisfied).

2. The drift operator and noise covariance are spectrally matched: both decay at rate $k^{-\alpha }$ in the Fourier basis, ensuring that no frequency band is over-smoothed or under-damped.

3. The stationary distribution is a centred Gaussian $𝒩(0,𝒞)$ on $\mathscr{H}$.

Well-posedness and convergence

Flow Matching for PDE Solvers

Diffusion models generate samples by simulating a stochastic differential equation, which requires many small steps and injects randomness at every stage. An increasingly popular alternative is flow matching: instead of a stochastic path from noise to data, we learn a deterministic velocity field $v_{\varphi }:\mathscr{H}\times [\mathrm{0,1}]\to \mathscr{H}$ that transports a simple base distribution (typically Gaussian) to the data distribution along smooth trajectories.

Rectified Flow-Based Operator (RFO)

A particularly elegant application of flow matching to PDEs is the Rectified Flow-based Operator (RFO), which unifies forward and inverse PDE solving in a single framework.

Unified forward and inverse solving

The key observation is that a trained flow $v_{\varphi }$ defines a bijection between the noise space and the data space. For PDE applications, we can choose the coupling more creatively: instead of coupling noise to data, we couple PDE inputs (initial/boundary conditions, parameters) to PDE outputs (solution fields).

The beauty of this formulation is the bidirectional nature:

• Forward PDE solving: Given input $a$, integrate $\mathrm{d}{w}_t/\mathrm{d}t=v_{\varphi }(w_t,t)$ from $t=0$ (where $w_0=a$) to $t=1$ to obtain the solution $u\approx w_1$.

• Inverse PDE solving: Given a target solution $u$, integrate backward from $t=1$ (where $w_1=u$) to $t=0$ to recover the input $a\approx w_0$.

Generative latent neural PDE solver

For PDEs on complex geometries, working directly in the ambient function space can be unwieldy. The generative latent neural PDE solver combines an autoencoder with flow matching in the latent space.

Architecture

The pipeline consists of three stages:

1. Encode: A geometry-aware encoder ${\mathrm{Enc}}_{\psi }$ maps the PDE field $u\in \mathscr{H}$ to a latent code $z\in {ℝ}^{d_z}$. The encoder uses graph neural network layers to handle unstructured meshes.

2. Flow match in latent space: A velocity network $v_{\varphi }:{ℝ}^{d_z}\times [\mathrm{0,1}]\to {ℝ}^{d_z}$ is trained via conditional flow matching on the latent codes.

3. Decode: The decoder ${\mathrm{Dec}}_{\psi }$ maps the generated latent code back to a function on the original mesh.

Flow-matching schedules for PDE trajectories

A subtle design choice is the flow-matching schedule for temporal sequences of PDE fields. Rather than treating each time step independently, the generative latent solver conditions on a window of past latent states and generates the next state autoregressively: (Genlpde)$$\begin{array}{ccc}& z_{n+1}={\mathrm{\Phi }}_{0\to 1}^{v_{\varphi }}(z_{\mathrm{noise}};\mathrm{context}=\{z_n,z_{n-1},\dots ,z_{n-K+1}\}),& \text{<span class="tml-eqn"></span>}\end{array}$$ where ${\mathrm{\Phi }}_{0\to 1}^{v_{\varphi }}$ denotes integration of the learned velocity field from $t=0$ to $t=1$, and the context window $\{z_{n-K+1},\dots ,z_n\}$ provides temporal information.

Fast Sampling and Distillation

The Achilles' heel of diffusion-based PDE solvers is sampling cost: generating a single PDE solution field requires integrating the reverse SDE or probability-flow ODE over many time steps, with each step invoking the score network-typically a large neural operator. For a practical PDE solver that must produce thousands of samples (e.g., for Monte Carlo uncertainty quantification), 100–1000 neural function evaluations (NFEs) per sample is prohibitive.

This section surveys three strategies for breaking the NFE bottleneck: truncated sampling, physics-informed refinement, and knowledge distillation.

Improved sampling for fluid dynamics

The first strategy, introduced at ICLR 2025 for fluid-dynamics applications, is based on the observation that not all denoising steps are equally valuable for PDE fields.

Truncation: skipping early denoising steps

In standard diffusion sampling, the reverse process begins at $t=T$ with pure noise and gradually denoises to $t=0$. The early steps ($t$ near $T$) remove only large-scale, low-frequency noise and contribute primarily to the overall “shape” of the field. For PDE applications where the large-scale structure is already known (e.g., from boundary conditions or a coarse solver), these steps are redundant.

The truncation strategy initialises the reverse process not from pure noise but from a partially noised version of a coarse estimate: (Truncation)$$\begin{array}{ccc}& u_{T-K}=\alpha (T-K)u_{\mathrm{coarse}}+\sigma (T-K)\xi ,\xi \sim 𝒩(0,𝒞),& \text{<span class="tml-eqn"></span>}\end{array}$$ and runs the reverse process only from $t=T-K$ to $t=0$, saving $K$ denoising steps.

Iterative refinement with physics guidance

After truncated denoising, the resulting field may violate physical constraints (e.g., divergence-free condition for incompressible flows). An iterative refinement step applies a physics-based projection: (Physrefine)$$\begin{array}{ccc}& u^{(k+1)}=u^{(k)}-\eta {\nabla }_u\Vert 𝒩[u^{(k)}]{\Vert }^2-\gamma {\nabla }_u\Vert u^{(k)}-{\hat{u}}_0^{\mathrm{Tweedie}}{\Vert }^2,& \text{<span class="tml-eqn"></span>}\end{array}$$ where the first term enforces the PDE constraint and the second anchors the solution to the Tweedie posterior mean estimate.

Phys-Instruct: ultra-fast PDE solving via distillation

While truncation reduces the number of steps by a constant factor, knowledge distillation offers a more dramatic reduction: compressing the entire multi-step diffusion trajectory into a student network that generates high-quality PDE solutions in one or a few steps.

Phys-Instruct applies this idea to PDE solvers with a physics-aware twist: the distillation objective is augmented with PDE residual penalties.

Architecture and training

The teacher is a pre-trained diffusion-based PDE solver that generates high-quality solutions in $N$ steps. The student is a neural operator $G_{\theta }$ that maps $(\mathrm{noise},\mathrm{conditions})\to \mathrm{solution}$ in a single forward pass. The training loss combines: (Physinstruct)$$\begin{array}{ccc}& {ℒ}_{\mathrm{PI}}(\theta )=\underset{\text{distillation loss}}{\underbrace{𝔼[\Vert G_{\theta }(z,c)-u_{\mathrm{teacher}}{\Vert }^2]}}+\lambda \underset{\text{PDE residual loss}}{\underbrace{𝔼[\Vert 𝒩[G_{\theta }(z,c)]{\Vert }^2]}},& \text{<span class="tml-eqn"></span>}\end{array}$$ where $z\sim 𝒩(0,I)$, $c$ encodes PDE parameters, and $u_{\mathrm{teacher}}$ is the teacher's multi-step output.

PDE-guided distillation objectives

The PDE residual term in serves a dual purpose:

1. It compensates for approximation errors in the teacher's output, ensuring that even when the student imperfectly mimics the teacher, the generated solutions still satisfy the governing equations.

2. It provides gradient signal in regions of solution space where the teacher has low sample density, effectively regularising the student.

PIDDM: physics-informed distillation revisited

A subtler challenge in distilling diffusion models for PDE applications was identified in the PIDDM framework (ICLR 2026). The issue is Jensen's gap: when physics constraints are applied to the expected output of a multi-step diffusion process, they do not guarantee that individual samples satisfy those constraints.

Jensen's Gap in physics-informed diffusion

Consider a diffusion model that generates PDE fields $u\sim p_{\varphi }$. We might hope that if $𝔼[\Vert 𝒩[u]{\Vert }^2]$ is small, then individual samples $u$ have small PDE residuals. However, by Jensen's inequality, (Jensen)$$\begin{array}{ccc}& \Vert 𝒩[𝔼[u]]{\Vert }^2\le 𝔼[\Vert 𝒩[u]{\Vert }^2],& \text{<span class="tml-eqn"></span>}\end{array}$$ but the reverse direction does not hold. The gap between these two quantities can be large when the generated distribution is multi-modal, as is typical for chaotic PDEs.

Post-hoc physical constraint enforcement

PIDDM resolves Jensen's gap through a two-stage approach:

1. Distillation: Train a few-step (or single-step) student model via standard consistency distillation.

2. Post-hoc projection: After each generated sample, apply a differentiable physics projection: (Piddm PROJ)$$\begin{array}{ccc}& u_{\mathrm{phys}}=\arg \underset{\tilde{u}}{\min }\Vert \tilde{u}-u_{\mathrm{student}}{\Vert }^2\text{s.t.}\Vert 𝒩[\tilde{u}]\Vert \le \delta ,& \text{<span class="tml-eqn"></span>}\end{array}$$ where $\delta$ is a tolerance parameter. In practice, this projection is implemented via a few steps of projected gradient descent.

Latent Space Methods for PDE Fields

The PDE fields we wish to generate-velocity, pressure, temperature, concentration-live in high-dimensional spaces: a turbulent flow snapshot on a $256\times 256\times 256$ grid with 5 field variables contains over $8\times {10}^7$ degrees of freedom. Performing diffusion or flow matching in this ambient space is computationally punishing. Latent space methods address this by first compressing the high-dimensional field to a compact latent representation, performing generative modelling in the low-dimensional latent space, and then decoding back to the physical field.

Why compress PDE fields?

Three practical reasons motivate latent space approaches:

1. Computational cost: Diffusion models require $O(N_{\mathrm{dof}})$ operations per denoising step. Reducing $N_{\mathrm{dof}}$ from ${10}^7$ to ${10}^3$ yields a ${10}^4\times$ reduction in per-step cost.

2. Structured priors: PDE fields are not arbitrary high-dimensional vectors; they lie on (or near) a low-dimensional manifold determined by the governing equations. The autoencoder learns this manifold, and the generative model only needs to learn a distribution on a low-dimensional space.

3. Multi-resolution compatibility: The encoder can accept fields at any resolution (using, e.g., point-cloud or graph-based architectures), while the latent space has fixed dimensionality. This provides resolution invariance without the full machinery of function-space diffusion.

CoNFiLD: conditional neural-field latent diffusion

CoNFiLD (Conditional Neural-Field Latent Diffusion) represents PDE fields as neural fields-implicit neural representations that map coordinates to field values-and performs diffusion in the space of neural-field parameters.

Neural field representation

Each PDE field $u$ is parameterised by a small neural network $f_{\omega }:\mathrm{\Omega }\to {ℝ}^m$ with weights $\omega$. The weights are obtained by fitting $f_{\omega }$ to the discrete field values: (Confild FIT)$$\begin{array}{ccc}& {\omega }^{\ast }(u)=\arg \underset{\omega }{\min }\sum _{i=1}^N\Vert f_{\omega }({𝐱}_i)-u({𝐱}_i){\Vert }^2.& \text{<span class="tml-eqn"></span>}\end{array}$$ The latent code is then $z={\omega }^{\ast }\in {ℝ}^{d_z}$, where $d_z$ is the number of parameters of the small network.

Conditional generation

CoNFiLD conditions the latent diffusion on PDE parameters $c$ (Reynolds number, boundary conditions, etc.) using classifier-free guidance: (Confild CFG)$$\begin{array}{ccc}& {\tilde{s}}_{\varphi }(z_t,t,c)=(1+w)s_{\varphi }(z_t,t,c)-w{s}_{\varphi }(z_t,t,⌀),& \text{<span class="tml-eqn"></span>}\end{array}$$ where $w$ is the guidance scale and $⌀$ denotes the null condition (dropout of $c$ during training).

Variational autoencoders for PDE fields

A natural choice for the encoder-decoder pair is a variational autoencoder (VAE), which provides a probabilistic latent space with a known prior (typically $𝒩(0,I)$). The ELBO for PDE fields is: (VAE)$$\begin{array}{ccc}& {ℒ}_{\mathrm{VAE}}(\psi ,\phi )={𝔼}_{u\sim p_{\mathrm{data}}}[-{𝔼}_{z\sim q_{\phi }(z|u)}[\log p_{\psi }(u|z)]+\mathrm{KL}(q_{\phi }(z|u)\Vert p(z)\left)\right].& \text{<span class="tml-eqn"></span>}\end{array}$$ The decoder likelihood $p_{\psi }(u|z)$ can be augmented with a PDE residual penalty: (VAE PHYS)$$\begin{array}{ccc}& \log p_{\psi }(u|z)\propto -\Vert u-{\mathrm{Dec}}_{\psi }(z){\Vert }^2-\lambda \Vert 𝒩[{\mathrm{Dec}}_{\psi }(z)]{\Vert }^2.& \text{<span class="tml-eqn"></span>}\end{array}$$

ECHO: extreme compression for spatiotemporal fields

For spatiotemporal PDE data (e.g., time series of 3D flow fields), the compression ratios achievable by standard autoencoders may be insufficient. ECHO (Efficient Compression for High-dimensional Output) pushes compression to the extreme: $100\times$ reduction in spatiotemporal dimensionality.

The key ideas are:

1. Temporal factorisation: Separate the spatial and temporal dimensions, compressing each independently before combining them in the latent space.

2. Physics-aware bottleneck: The latent space is structured to respect conservation laws-e.g., the latent code decomposes into a “conserved” subspace (encoding total mass, energy, etc.) and a “dynamic” subspace (encoding the evolving spatial structure).

3. Multi-scale encoding: A hierarchical encoder captures both large-scale structures and fine-grained details, with coarser levels receiving more aggressive compression.

Comparing latent and function-space approaches

Both function-space diffusion (Section Function-Space Diffusion: Beyond Grids) and latent diffusion address the resolution-dependence problem, but via different mechanisms. We summarise the tradeoffs:

Exercises

PDE-Constrained Inverse Problems

In all of the preceding sections, we have been concerned with the forward problem: given the PDE parameters $\theta$-initial conditions, boundary data, material coefficients, source terms-compute the solution field $u(𝐱,t)$. We now turn the arrow around. In the inverse problem, we observe partial, noisy measurements of the solution and wish to recover the unknown parameters $\theta$ that gave rise to it.

“The inverse problem is the question that nature answers every time we make a measurement, and the one that scientists struggle with every time they try to interpret one.” - Albert Tarantola

Inverse problems are ubiquitous in science and engineering:

• Medical imaging: given X-ray or ultrasound measurements, reconstruct the internal tissue density (a coefficient in the wave equation).

• Geophysics: given seismic waveforms recorded at the surface, infer the subsurface velocity model (a coefficient in the elastic wave equation).

• Astrophysics: given sparse telescope data, reconstruct black hole accretion disk images (governed by radiative transfer equations).

• Climate science: given satellite observations of sea surface temperature, infer subsurface ocean currents (governed by the primitive equations).

Mathematical formulation

Consider a PDE parameterised by $\theta \in \mathrm{\Theta }$: (PDE)$$\begin{array}{ccc}& {𝒩}_{\theta }[u]=0\text{in }\mathrm{\Omega }\times (0,T],& \text{<span class="tml-eqn"></span>}\end{array}$$ where ${𝒩}_{\theta }$ is a (possibly nonlinear) differential operator depending on the parameter $\theta$. We do not observe the full solution $u$; instead, we measure (OBS)$$\begin{array}{ccc}& 𝒚=\mathscr{H}(u)+𝝐,& \text{<span class="tml-eqn"></span>}\end{array}$$ where $\mathscr{H}:𝒰\to {\mathrm{ℝ}}^m$ is the observation operator (mapping the infinite-dimensional solution to finitely many sensor readings) and $𝝐\sim 𝒩(0,\mathrm{\Gamma })$ is additive Gaussian noise with covariance $\mathrm{\Gamma }\in {\mathrm{ℝ}}^{m\times m}$.

This is typically cast as an optimisation problem: (Optimisation)$$\begin{array}{ccc}& \hat{\theta }={\mathrm{arg\hspace{0.17em}min}}_{\theta \in \mathrm{\Theta }}\frac{1}{2}\Vert 𝒚-\mathscr{H}(u(\theta )){\Vert }_{{\mathrm{\Gamma }}^{-1}}^2+\lambda R(\theta ),& \text{<span class="tml-eqn"></span>}\end{array}$$ where $\Vert 𝒗{\Vert }_{{\mathrm{\Gamma }}^{-1}}^2={𝒗}^{\top }{\mathrm{\Gamma }}^{-1}𝒗$ is the Mahalanobis norm, $R(\theta )$ is a regularisation functional, and $\lambda >0$ controls the strength of regularisation. The PDE ${𝒩}_{\theta }[u]=0$ enters as an equality constraint-hence the name PDE-constrained optimisation.

Ill-posedness and the Hadamard conditions

Jacques Hadamard (1902) defined a problem as well-posed if it satisfies three conditions:

Most PDE inverse problems are severely ill-posed: the observation operator $\mathscr{H}$ typically loses information (it maps an infinite-dimensional field to finitely many measurements), so multiple parameters $\theta$ can produce observations consistent with the data. Uniqueness fails. Even when a unique solution exists, it may be extremely sensitive to noise-stability fails.

Bayesian formulation

The modern approach to ill-posed inverse problems is Bayesian: instead of seeking a single point estimate $\hat{\theta }$, we characterise the full posterior distribution (Posterior)$$\begin{array}{ccc}& p(\theta |𝒚)=\frac{p(𝒚|\theta )p(\theta )}{p(𝒚)}\propto p(𝒚|\theta )p(\theta ),& \text{<span class="tml-eqn"></span>}\end{array}$$ where:

• $p(\theta )$ is the prior, encoding our beliefs about $\theta$ before seeing the data;

• $p(𝒚|\theta )=𝒩(𝒚|\mathscr{H}(u(\theta )),\mathrm{\Gamma })$ is the likelihood, measuring the fit of the PDE prediction to the observations;

• $p(𝒚)=\int p(𝒚|\theta )p(\theta )\mathrm{d}\theta$ is the evidence (a normalising constant);

• $p(\theta |𝒚)$ is the posterior, our updated belief about $\theta$ given the observations.

The challenge, of course, is that computing the posterior requires solving the PDE (within the likelihood) for each candidate $\theta$. In Markov chain Monte Carlo (MCMC) methods, this may mean running the PDE solver millions of times-a prohibitive cost for any but the simplest equations. This is precisely where generative models enter the picture.

Classical approaches: adjoint methods and Tikhonov regularisation

Before diving into diffusion-based methods, let us briefly review the classical workhorses of PDE-constrained optimisation.

Tikhonov regularisation.

The simplest approach to is the quadratic penalty $R(\theta )=\Vert \theta -{\theta }_0{\Vert }^2$ for some reference value ${\theta }_0$ (often zero). The resulting estimator is (Tikhonov)$$\begin{array}{ccc}& {\hat{\theta }}_{\lambda }={\mathrm{arg\hspace{0.17em}min}}_{\theta }\frac{1}{2}\Vert 𝒚-\mathscr{H}(u(\theta )){\Vert }_{{\mathrm{\Gamma }}^{-1}}^2+\frac{\lambda }{2}\Vert \theta -{\theta }_0{\Vert }^2.& \text{<span class="tml-eqn"></span>}\end{array}$$ This is the MAP estimate under a Gaussian prior $\theta \sim 𝒩({\theta }_0,{\lambda }^{-1}I)$. The regularisation parameter $\lambda$ balances data fidelity against prior belief.

The adjoint method.

Computing the gradient ${\nabla }_{\theta }J(\theta )$ of the objective $J(\theta )=\frac{1}{2}\Vert 𝒚-\mathscr{H}(u(\theta )){\Vert }_{{\mathrm{\Gamma }}^{-1}}^2$ requires differentiating through the PDE solve. The adjoint method achieves this at a cost comparable to one additional PDE solve (the adjoint equation), regardless of the dimension of $\theta$.

Diffusion Models as Priors for Inverse Problems

The key insight connecting diffusion models to inverse problems is deceptively simple:

Posterior sampling via guided diffusion

Recall from the diffusion chapter that a score-based generative model learns the score function ${\nabla }_u\log p_t(u)$ at each noise level $t$. The reverse-time SDE that generates samples from $p(u)$ is (Reverse SDE)$$\begin{array}{ccc}& \mathrm{d}u=[f(u,t)-g(t{)}^2{\nabla }_u\log p_t(u)]\mathrm{d}t+g(t)\mathrm{d}\bar{W},& \text{<span class="tml-eqn"></span>}\end{array}$$ where $f$ and $g$ are the drift and diffusion coefficients.

To sample from the posterior $p(u|𝒚)$ instead of the prior $p(u)$, we apply Bayes' rule at each noise level: (Posterior Score)$$\begin{array}{ccc}& {\nabla }_u\log p_t(u|𝒚)=\underset{\text{prior score (learned)}}{\underbrace{{\nabla }_u\log p_t(u)}}+\underset{\text{likelihood score}}{\underbrace{{\nabla }_u\log p_t(𝒚|u)}}.& \text{<span class="tml-eqn"></span>}\end{array}$$ The prior score is provided by the pre-trained diffusion model. The likelihood score requires approximation, since $p_t(𝒚|u)$ at intermediate noise levels is generally intractable.

Approximation strategies.

Several approaches have been proposed for the likelihood score:

1. Diffusion Posterior Sampling (DPS): Approximate $p_t(𝒚|u_t)\approx p(𝒚|{\hat{u}}_0(u_t))$ where ${\hat{u}}_0$ is the denoised estimate from Tweedie's formula: ${\hat{u}}_0(u_t)=𝔼[u|u_t]$. The gradient is then ${\nabla }_{u_t}\log p_t(𝒚|u_t)\approx -{\nabla }_{u_t}\Vert 𝒚-\mathscr{H}({\hat{u}}_0(u_t)){\Vert }_{{\mathrm{\Gamma }}^{-1}}^2/2$.

2. Pseudoinverse guidance: Project the measurement residual back to the data space using the pseudoinverse of the linearised observation operator.

3. $\mathrm{\Pi }$GDM (Projected Gradient Descent in Measurement space): Alternate between unconditional denoising steps and projection steps that enforce measurement consistency.

The plug-and-play paradigm

A remarkable feature of the diffusion-prior approach is that it is plug-and-play: the diffusion model (prior) and the forward model (likelihood) are completely separate modules. One can train a single unconditional diffusion model on a corpus of PDE solutions and then re-use it for any downstream inverse problem-tomography, data assimilation, parameter estimation-simply by changing the observation operator $\mathscr{H}$ and the noise model. No retraining or fine-tuning of the diffusion model is required.

Property	Classical (e.g. Tikhonov/TV)	Diffusion Prior
Prior form	Hand-crafted penalty	Learned from data
Expressiveness	Low (smooth/sparse)	High (multimodal)
Retraining needed	No	No (plug-and-play)
Uncertainty quantification	Via MCMC (expensive)	Via sample diversity
Nonlinear structure	Limited	Fully captured
Computational cost	Low per iteration	Moderate (many NFEs)
Theoretical guarantees	Strong (regularisation theory)	Emerging

RED-DiffEq: Full Waveform Inversion

The geophysical inverse problem

Full Waveform Inversion (FWI) is one of the most challenging inverse problems in the geosciences. The goal is to reconstruct the subsurface velocity model $c(𝐱)$-which determines how fast seismic waves propagate through the Earth-from seismograms recorded at the surface.

The forward model is the acoustic wave equation: (WAVE)$$\begin{array}{ccc}& \frac{1}{c{(𝐱)}^2}\frac{{\partial }^{2}p}{\partial t^2}={\nabla }^{2}p+s(𝐱,t),& \text{<span class="tml-eqn"></span>}\end{array}$$ where $p(𝐱,t)$ is the pressure field and $s$ is the source term (an impulsive or swept-frequency excitation at a known location). Receivers at the surface record $p$ at discrete positions, yielding the observed seismograms $𝒚$.

The RED-DiffEq framework

RED-DiffEq (Sun et al., 2024) addresses FWI by combining a physics-driven inversion loop with a diffusion model acting as a Regularisation by Denoising (RED) prior. The core idea is elegant: at each iteration of the inversion, the current velocity estimate is “cleaned up” by a diffusion denoising step, which projects it back onto the manifold of geologically plausible velocity models.

Domain decomposition for large-scale problems

Real-world FWI problems involve domains of tens or hundreds of kilometres, discretised at sub-wavelength resolution. The velocity field $c(𝐱)$ may live on a grid of ${10}^6$–${10}^8$ unknowns. Training a diffusion model on the full domain at this resolution is impractical.

RED-DiffEq addresses this via domain decomposition: the large domain is partitioned into overlapping patches, each of moderate resolution (e.g., $128\times 128$). The diffusion model is trained on individual patches and applied independently during the denoising step. Consistency across patch boundaries is enforced by the physics-based gradient step, which operates on the full domain. In practice, this strategy has been shown to reconstruct velocity models on domains $8\times$ larger than the diffusion model's training resolution.

InverseBench: A Community Benchmark for PDE Inverse Problems

The rapid growth of diffusion-based methods for inverse problems has created an urgent need for standardised evaluation. InverseBench (Ye et al., ICLR 2025) addresses this by providing a unified benchmark spanning diverse PDE inverse problems with standardised datasets, metrics, and evaluation protocols.

Benchmark tasks

InverseBench covers five tasks drawn from different physical domains:

1. Optical diffuse tomography: Reconstruct the spatially varying absorption coefficient ${\mu }_a(𝐱)$ from boundary measurements of diffuse light. Governed by the diffusion approximation to the radiative transfer equation.

2. Electrical impedance tomography (EIT): Reconstruct the conductivity distribution $\sigma (𝐱)$ from boundary voltage measurements. Governed by Laplace's equation $\nabla \cdot (\sigma \nabla V)=0$.

3. Black hole imaging: Reconstruct the emission profile of a black hole accretion disk from sparse interferometric measurements (as in the Event Horizon Telescope). Governed by radiative transfer in curved spacetime.

4. Turbulent flow reconstruction: Recover the full velocity field of a turbulent flow from sparse point measurements. Governed by the incompressible Navier–Stokes equations.

5. Seismic wave inversion: A simplified FWI task: recover the velocity model from surface seismograms.

Evaluated methods

InverseBench evaluates several plug-and-play diffusion methods:

• DPS (Diffusion Posterior Sampling)

• DiffPIR (Diffusion-based Plug-and-Play Image Restoration)

• $\mathrm{\Pi }$GDM (Pseudoinverse-Guided Diffusion Models)

• DDRM (Denoising Diffusion Restoration Models)

• RED-Diff (Regularisation by Denoising with Diffusion)

Method	Opt. Tomo.	EIT	Black Hole	Turbulence	Seismic
Tikhonov	21.3	18.7	16.2	19.8	20.1
DPS	26.8	24.1	22.5	24.3	25.7
DiffPIR	27.2	23.8	23.1	25.0	26.3
$\mathrm{\Pi }$GDM	27.5	25.3	23.8	25.2	26.8
DDRM	25.1	22.6	21.9	23.1	24.5
RED-Diff	28.1	24.9	24.4	26.1	27.5

Score Matching via Differentiable Physics (SMDP)

While the plug-and-play approach separates the diffusion prior from the physics, an alternative paradigm integrates the physics simulator directly into the diffusion process. Score Matching via Differentiable Physics (SMDP, Holzschuh et al., 2023) embeds the physics simulator in the drift term of the SDE that defines the generative process.

The SMDP framework

Consider a parameterised PDE family with parameters $\theta$. SMDP constructs a diffusion process whose forward SDE gradually corrupts a pair $(\theta ,u(\theta ))$ of parameter and solution: (Forward)$$\begin{array}{ccc}& \mathrm{d}\left(\begin{array}{c}{\theta }_t\\ u_t\end{array}\right)=f(t)\left(\begin{array}{c}{\theta }_t\\ u_t\end{array}\right)\mathrm{d}t+g(t)\mathrm{d}{W}_t.& \text{<span class="tml-eqn"></span>}\end{array}$$ The reverse-time SDE then generates joint samples $(\theta ,u)\sim p(\theta ,u)$.

The key innovation is in the score network architecture: instead of predicting the full joint score ${\nabla }_{(\theta ,u)}\log p_t(\theta ,u)$, SMDP factorises it as (Score)$$\begin{array}{ccc}& {\nabla }_u\log p_t(u|\theta )\approx {\nabla }_u[-\frac{1}{2{\sigma }_t^2}\Vert u-S(\theta ){\Vert }^2]=\frac{S(\theta )-u}{{\sigma }_t^2},& \text{<span class="tml-eqn"></span>}\end{array}$$ where $S(\theta )$ is a differentiable PDE simulator (e.g., a neural PDE surrogate or a differentiable finite-difference solver). The parameter score ${\nabla }_{\theta }\log p_t(\theta )$ is learned by a separate network.

Uncertainty Quantification in Neural PDE Solvers

A neural PDE solver that produces a single prediction without any indication of confidence is of limited use in safety-critical applications. When a bridge designer uses a surrogate model to estimate stress distributions, or a weather forecaster relies on a neural operator for rapid prediction, the question is not just “what is the predicted solution?” but “how much should I trust this prediction?”

“It is better to be approximately right than precisely wrong.” - attributed to John Maynard Keynes

Why UQ matters

Epistemic vs. aleatoric uncertainty

It is essential to distinguish two fundamentally different sources of uncertainty:

For PDE surrogates, both types are relevant:

• Epistemic: The neural operator has been trained on a finite dataset of PDE solutions. In regions of parameter space far from the training distribution, its predictions are uncertain due to lack of exposure.

• Aleatoric: For chaotic systems (turbulence, weather), the solution at long times is inherently unpredictable: tiny perturbations in initial conditions lead to vastly different trajectories. No amount of data can eliminate this uncertainty.

Bayesian neural operators

The most principled approach to epistemic UQ is to place a prior $p(\varphi )$ over the neural operator's parameters $\varphi$ and compute the posterior $p(\varphi |𝒟)$ given training data $𝒟$. Predictions are then made by marginalising over parameters: (Bayesian PRED)$$\begin{array}{ccc}& p(u|\theta ,𝒟)=\int p(u|\theta ,\varphi )p(\varphi |𝒟)\mathrm{d}\varphi .& \text{<span class="tml-eqn"></span>}\end{array}$$ The predictive variance decomposes as (VAR Decomp)$$\begin{array}{ccc}& \underset{\text{total}}{\underbrace{\mathrm{𝖵ar}[u|\theta ,𝒟]}}=\underset{\text{aleatoric}}{\underbrace{{𝔼}_{\varphi }[\mathrm{𝖵ar}[u|\theta ,\varphi ]]}}+\underset{\text{epistemic}}{\underbrace{{\mathrm{𝖵ar}}_{\varphi }[𝔼[u|\theta ,\varphi ]]}}.& \text{<span class="tml-eqn"></span>}\end{array}$$

In practice, exact Bayesian inference over the millions of parameters of a neural operator is intractable. Common approximations include:

MC Dropout.

At test time, keep dropout active and run $M$ forward passes. The empirical variance of the $M$ predictions provides an estimate of epistemic uncertainty. This corresponds to an approximate variational posterior (Gal and Ghahramani, 2016).

Deep ensembles.

Train $M$ independent copies of the neural operator with different random seeds and/or data orderings. The ensemble mean approximates the Bayesian predictive mean, and the ensemble variance captures epistemic uncertainty. Despite its simplicity, deep ensembles consistently provide well-calibrated uncertainty estimates and are often hard to beat.

Laplace approximation.

Fit a Gaussian $𝒩(\hat{\varphi },H^{-1})$ to the posterior, where $\hat{\varphi }$ is the MAP estimate and $H$ is the Hessian of the negative log-posterior. Modern implementations use the Generalised Gauss–Newton approximation to make this scalable.

Diffusion models as natural UQ tools

One of the most compelling advantages of generative PDE solvers over deterministic surrogates is that they provide uncertainty quantification “for free.” A conditional diffusion model $p_{\psi }(u|\theta )$ can generate diverse samples $u^{(1)},u^{(2)},\dots ,u^{(M)}\sim p_{\psi }(u|\theta )$. The spread of these samples directly quantifies the model's uncertainty about the solution.

Calibration

Generating uncertainty estimates is useful only if those estimates are reliable. A model is said to be calibrated if its predicted confidence levels match empirical frequencies.

In practice, we assess calibration by the calibration plot: for each nominal level $\alpha$ (e.g., $50\%$, $80\%$, $90\%$, $95\%$), we compute the fraction of test points that fall within the predicted $\alpha$-credible interval. A well-calibrated model yields points that lie on the diagonal of this plot.

Conformal prediction for PDE control

Conformal prediction provides distribution-free, finite-sample coverage guarantees for uncertainty estimates, requiring only the assumption of exchangeability (a weaker condition than i.i.d.).

SafeDiffCon (Wei et al., 2024) applies this idea to PDE-based control: a diffusion model generates candidate control trajectories, conformal prediction wraps these in guaranteed coverage bands, and a safety filter ensures that the executed control action keeps the system state within the conformal uncertainty tube. This provides formal safety guarantees for neural PDE-based controllers, even when the underlying diffusion model is imperfect.

Benchmarks and Evaluation Practice

The scientific machine learning community has experienced explosive growth over the past five years. Hundreds of papers propose new neural PDE solvers each year, each reporting results on a different set of equations, resolutions, training protocols, and metrics. This heterogeneity makes it nearly impossible to compare methods fairly or to identify genuine progress.

The antidote is standardised benchmarking: community-maintained datasets, evaluation protocols, and metrics that allow apples-to-apples comparison. In this section, we survey the major PDE benchmarks, discuss evaluation metrics beyond the ubiquitous $L^2$ relative error, and distill principles for designing good benchmarks.

“Without data, you're just another person with an opinion.” - W. Edwards Deming

PDEBench

PDEBench (Takamoto et al., NeurIPS 2022) was the first large-scale, community benchmark for neural PDE solvers. It provides high-resolution simulation data for a diverse collection of PDEs, together with standardised training/test splits and evaluation code.

Datasets

PDEBench covers the following equations:

1. 1D Advection equation: ${\partial }_{t}u+\beta {\partial }_{x}u=0$ with varying advection speed $\beta$.

2. 1D/2D Burgers equation: ${\partial }_{t}u+u{\partial }_{x}u=\nu {\partial }_x^{2}u$ with varying viscosity $\nu$.

3. 1D/2D Diffusion-reaction equation: ${\partial }_{t}u=D{\nabla }^{2}u+R(u)$ with varying diffusion coefficient $D$ and reaction rate.

4. 1D/2D Diffusion-sorption equation: ${\partial }_t(u+s(u))=D{\nabla }^{2}u$ modelling contaminant transport with nonlinear sorption.

5. 2D Compressible Navier–Stokes: The Euler equations with a Mach number ranging from $0.1$ to $1.0$, including shock formation.

6. 2D Darcy flow: $-\nabla \cdot (a(𝐱)\nabla u)=f$ with log-normal random permeability fields.

7. 2D Shallow water equations: Modelling surface gravity waves with varying bathymetry.

Each dataset contains $\mathrm{10,000}$ simulation instances at resolutions ranging from $128$ to $1024$ grid points per spatial dimension, stored in HDF5 format with metadata.

Baseline methods

PDEBench evaluates three baseline architectures:

• U-Net: A standard encoder-decoder CNN with skip connections, operating on the solution field at a fixed resolution.

• FNO: The Fourier Neural Operator, learning mappings in Fourier space.

• PINN: The physics-informed approach, trained on PDE residuals without data.

Equation	U-Net	FNO	PINN
1D Advection ($\beta =0.4$)	0.12	0.08	2.31
1D Burgers ($\nu =0.01$)	1.84	0.93	8.72
1D Diff.-React.	0.45	0.31	4.18
2D Darcy Flow	1.21	0.56	6.93
2D Comp. Navier–Stokes	5.73	3.82	15.4
2D Shallow Water	3.41	2.17	12.8

Reproducibility infrastructure

PDEBench provides:

• Standardised data loaders in PyTorch with consistent normalisation, train/validation/test splits, and augmentation protocols.

• Reference implementations of all baseline models with hyperparameter configurations.

• Evaluation scripts computing all reported metrics.

• Leaderboard for community submissions.

This infrastructure has made PDEBench the de facto standard for evaluating new neural PDE solvers, with over 200 papers using it as of early 2025.

PINNacle

While PDEBench evaluates data-driven surrogates, PINNacle (Hao et al., NeurIPS 2024) focuses specifically on the evaluation of Physics-Informed Neural Networks.

Scope and design

PINNacle includes over 20 PDEs spanning:

• Fluid dynamics: Navier–Stokes (lid-driven cavity, cylinder wake), Stokes flow, Kovasznay flow.

• Solid mechanics: linear elasticity, plate bending.

• Heat transfer: steady and unsteady conduction, convection-diffusion.

• Electromagnetism: Poisson, Helmholtz.

• Nonlinear PDEs: Allen–Cahn, Cahn–Hilliard, Korteweg–de Vries (KdV), Burgers.

• High-dimensional PDEs: Black–Scholes in up to 10 dimensions, Hamilton–Jacobi–Bellman.

For each PDE, PINNacle provides:

• Reference solutions (high-resolution FEM/spectral).

• Multiple PINN variants: vanilla, causal, gradient-reweighted, domain-decomposition, and hard-constraint architectures.

• Systematic hyperparameter sweeps: learning rate, network width/depth, activation function, number of collocation points.

Key findings

1. Optimisation pathologies are pervasive: Across the benchmark, PINNs frequently fail to converge to accurate solutions, especially for problems with sharp gradients, multiscale features, or complex geometries. The failure rate (defined as $>10\%$ relative error) exceeds $30\%$ for challenging problems even with extensive hyperparameter tuning.

2. No universal PINN architecture: Different PINN variants excel on different problem classes. Causal PINNs help for time-dependent problems but not for steady-state. Hard-constraint architectures help for problems with important boundary conditions but introduce their own optimisation issues.

3. Scaling limits: Increasing the network size beyond $\sim$4 layers and $\sim$256 neurons per layer provides diminishing returns for most problems. The bottleneck is optimisation, not capacity.

4. Collocation point distribution matters more than quantity: Adaptive sampling strategies (e.g., residual-based refinement) consistently outperform uniform random sampling, often achieving the same accuracy with $5\times$ fewer points.

Evaluation Metrics Beyond $L^2$

The most commonly reported metric in the neural PDE literature is the relative $L^2$ error: (L2)$$\begin{array}{ccc}& e_{L^2}=\frac{\Vert u_{\text{pred}}-u_{\text{true}}{\Vert }_{L^2}}{\Vert u_{\text{true}}{\Vert }_{L^2}},& \text{<span class="tml-eqn"></span>}\end{array}$$ where the $L^2$ norm is computed over the spatial domain (and possibly time). While convenient and widely understood, the $L^2$ error has significant limitations.

Energy norms: $H^1$ and $H^2$

Many PDEs have natural energy norms associated with their variational formulations. For instance, the Poisson equation $-{\nabla }^{2}u=f$ has the energy norm (H1)$$\begin{array}{ccc}& \Vert u{\Vert }_{H^1}^2=\Vert u{\Vert }_{L^2}^2+\Vert \nabla u{\Vert }_{L^2}^2.& \text{<span class="tml-eqn"></span>}\end{array}$$ The $H^1$ error penalises errors in the gradient field, which is physically relevant because the gradient often corresponds to a physical flux (heat flux, stress, velocity gradient).

The $H^2$ seminorm $|\cdot {|}_{H^2}=\Vert {\nabla }^{2}u{\Vert }_{L^2}$ is relevant for fourth-order problems (plate bending, Cahn–Hilliard) where second derivatives carry physical meaning.

Spectral metrics

For turbulent flows, the energy spectrum $E(k)$-the distribution of kinetic energy across wavenumbers $k$-is a fundamental diagnostic. A surrogate that matches the solution pointwise in $L^2$ may nevertheless produce a completely wrong energy spectrum (too much energy at high wavenumbers, or too little).

For turbulence, the Kolmogorov $-5/3$ scaling law predicts $E(k)\sim k^{-5/3}$ in the inertial range. A good generative model should reproduce this scaling, even if the individual vortex positions differ from the reference solution.

Conservation residuals

Physical PDEs are derived from conservation laws: conservation of mass, momentum, and energy. A surrogate that violates these laws may produce plausible-looking but physically meaningless solutions.

Common conservation checks include:

• Mass conservation: Is ${\int }_{\mathrm{\Omega }}\rho (t)\mathrm{d}𝐱$ constant over time?

• Energy conservation/dissipation: Does the total energy ${\int }_{\mathrm{\Omega }}\frac{1}{2}|𝒖{|}^2\mathrm{d}𝐱$ decay at the correct rate for viscous flows?

• Momentum conservation: For inviscid flows, is the total momentum preserved?

• Divergence-free condition: For incompressible flows, is $\nabla \cdot 𝒖=0$ satisfied pointwise?

Physically relevant statistics

For stochastic or chaotic systems, pointwise comparison is meaningless at long times. Instead, we compare statistical properties:

• Vorticity PDF: The probability distribution of vorticity values, capturing the prevalence of extreme events.

• Structure functions: $S_p(r)=⟨|\delta u(r){|}^p⟩$ where $\delta u(r)=u(𝐱+𝐫)-u(𝐱)$ and $|𝐫|=r$. For Kolmogorov turbulence, $S_2(r)\sim r^{2/3}$ in the inertial range.

• Temporal autocorrelation: How rapidly decorrelation occurs-a key indicator of whether the model captures the correct dynamical timescales.

• Extreme value statistics: The distribution of the maximum vorticity, peak stress, or other extreme quantities-crucial for engineering applications where failure is governed by extremes.

Stability-informed metrics

For time-dependent problems, it is important to assess not just accuracy at a single time but stability over long rollouts. Key metrics include:

• Valid prediction time: The time $T_{\text{valid}}$ until the relative $L^2$ error first exceeds a threshold (e.g., $50\%$). For chaotic systems, this is analogous to the Lyapunov time.

• Error growth rate: The rate at which the $L^2$ error grows with time. Physically consistent models should exhibit error growth rates consistent with the system's Lyapunov exponents.

• Blow-up detection: Does the model produce divergent predictions (NaN, exponentially growing values) during long rollouts? A metric measuring the fraction of test trajectories that remain bounded provides a crude but important stability measure.

Category	Metric	What it captures
3*Pointwise	$L^2$ relative error	Overall solution accuracy
	$L^{\infty }$ error	Worst-case local error
	MSE at sensors	Prediction at observation points
2*Regularity	$H^1$ error	Gradient accuracy
	$H^2$ error	Curvature accuracy
2*Spectral	Energy spectrum error	Scale-by-scale fidelity
	Spectral rolloff	High-frequency content
3*Conservation	Mass residual	Mass conservation
	Energy residual	Energy conservation/dissipation
	Divergence error	Incompressibility constraint
3*Statistical	Vorticity PDF error	Distributional accuracy
	Structure functions	Multi-scale correlations
	Temporal autocorrelation	Dynamical timescales
3*Stability	Valid prediction time	Temporal stability
	Error growth rate	Lyapunov-like behaviour
	Blow-up fraction	Numerical stability

What Makes a Good Benchmark?

Having surveyed existing benchmarks and metrics, we now distill principles for designing effective PDE benchmarks.

Feature	PDEBench	PINNacle	InverseBench
Focus	Data-driven surrogates	PINNs	Inverse problems
Number of PDEs	7 families	20+	5
Spatial dims	1D, 2D	1D–3D	2D
Metrics	$L^2$, MSE	$L^2$, $L^{\infty }$	PSNR, SSIM, LPIPS
Conservation checks	No	No	No
Spectral metrics	No	No	No
Open-source	Yes	Yes	Yes
Year	2022	2024	2025

Open Problems and Future Challenges

“In so far as the laws of mathematics refer to reality, they are not certain; and in so far as they are certain, they do not refer to reality.” - Albert Einstein, 1921

We have now traversed a long path-from the classical PDE solvers of the twentieth century through physics-informed neural networks, neural operators, Transformer-based architectures, and diffusion models on function spaces. The tools are powerful, the progress remarkable, and the pace of discovery exhilarating. Yet the honest truth is that we are still near the beginning. The most profound questions-about turbulence, about guarantees, about the relationship between data-driven and physics-driven understanding-remain wide open.

This section surveys the frontier. We revisit the most famous unsolved problem in mathematical physics, chart the challenges that must be overcome for generative PDE solvers to fulfil their promise, and close with philosophical reflections that we hope will inspire the next generation of researchers.

The Navier–Stokes Millennium Prize Revisited

In Section A First Look at PDEs and Why They Matter we encountered the Navier–Stokes equations and mentioned that the question of their regularity in three dimensions is one of the seven Clay Millennium Prize Problems, carrying a prize of one million US dollars. Now that we have developed the full machinery of generative PDE solvers, we can ask a more pointed question: can machine learning help crack this problem?

The formal statement

Let us state the problem precisely.

At first glance, this problem may seem remote from machine learning. It is a question about the mathematical existence and regularity of solutions, not about their numerical computation. But as we shall argue, there are surprising and potentially transformative connections.

Why 3D is fundamentally different from 2D

The key to understanding why the 3D problem is so hard lies in a single mechanism: vortex stretching. Define the vorticity $𝝎=\nabla \times 𝐮$. Taking the curl of , we obtain the vorticity equation: (Vorticity)$$\begin{array}{ccc}& \frac{\partial 𝝎}{\partial t}+(𝐮\cdot \nabla )𝝎=\underset{\text{vortex stretching}}{\underbrace{(𝝎\cdot \nabla )𝐮}}+\nu \mathrm{\Delta }𝝎.& \text{<span class="tml-eqn"></span>}\end{array}$$

This criterion tells us exactly what to look for: scenarios where the pointwise maximum of vorticity grows without bound.

Partial results: from Leray to the present

Can ML help? Probing near-singular behaviour

Here is where generative models enter the picture. Even if they cannot prove theorems directly, they can serve as extraordinarily powerful experimental tools for probing the landscape of possible Navier–Stokes solutions.

Several concrete research directions are emerging:

1. Diffusion models for extreme vorticity sampling. Train a score-based generative model on the distribution of vorticity fields from turbulent simulations. Condition the model on high values of $\Vert 𝝎{\Vert }_{\infty }$ and sample configurations that push the envelope of vorticity concentration. This is an instance of rare-event sampling using generative models.

2. AI-assisted conjecture generation. Feed a large language model the corpus of partial regularity results and ask it to propose new conditional regularity criteria-new functional norms or geometric conditions under which smoothness can be guaranteed. The model can suggest candidates; formal theorem provers (see the mathematics chapter) can attempt to verify them.

3. Learned Lyapunov functionals. A classical approach to proving regularity is to construct a Lyapunov functional (an energy-like quantity that decreases along solutions). Neural networks can be trained to parameterise candidate Lyapunov functionals, with the loss function penalising violations of the monotonicity condition. If a valid Lyapunov functional is found, it constitutes a machine-generated proof of regularity.

4. Symmetry-aware search. The Navier–Stokes equations have rich symmetry groups (spatial translations, rotations, scaling). Generative models that respect these symmetries (equivariant architectures) can focus the search on genuinely distinct configurations, avoiding redundant exploration of symmetry-related states.

Discretisation-Agnostic Generation

Every method we have studied in this chapter operates, at some level, on a discrete representation: grid points, mesh nodes, spectral coefficients, or tokens. This discretisation is both a practical necessity and a fundamental limitation.

The grid-dependence problem

Consider training a diffusion model on fluid-flow snapshots at resolution $64\times 64$. The model learns to generate fields at precisely that resolution. If we need a $256\times 256$ field-perhaps because a fine-scale feature demands higher resolution-the trained model cannot help. We must retrain on higher-resolution data, which is expensive to generate and requires more GPU memory.

Function-space methods as the solution

The key idea is to work directly in function space rather than in ${ℝ}^N$ for a fixed $N$. We encountered this perspective in the diffusion-on-function-space formulation: the forward process adds noise in a Hilbert space $H$, and the reverse process denoises back to a clean function. The resulting model is, in principle, resolution-independent.

Open questions abound:

1. Can we define well-posed diffusion processes in infinite-dimensional Hilbert spaces with convergence guarantees for the score-matching objective?

2. How do we efficiently parameterise score functions on function spaces? Current approaches use discretised approximations-can we avoid this entirely?

3. What is the correct notion of “distribution” over solutions of a nonlinear PDE? For linear PDEs, Gaussian measures on Sobolev spaces provide a natural framework. For nonlinear PDEs, the situation is far less clear.

4. Can neural implicit representations (NeRF-like coordinate networks for PDE solutions) achieve true resolution independence, or do they inherit the spectral bias of MLPs?

Real-Time Turbulence and Interactive Simulation

“The goal is not to replace DNS, but to make its insights available in real time.” - Karthik Duraisamy, 2023

For many applications-flight simulators, digital twins of industrial processes, real-time weather nowcasting, interactive surgical planning-the speed of a PDE solver matters as much as its accuracy. Classical high-fidelity solvers deliver accuracy but require hours or days of computation. Generative PDE solvers promise a qualitative leap: from offline batch processing to interactive, real-time simulation.

The 1000$\times$ speedup challenge

Current state-of-the-art neural surrogates achieve speedups of $100$–$1000\times$ over classical solvers for moderate-complexity problems. But for industrially relevant turbulent flows at high Reynolds numbers ($\mathrm{Re}\gtrsim {10}^6$), we still fall short of real-time requirements by one to two orders of magnitude.

Latent-space approaches

The most promising path to real-time performance is to operate entirely in a compressed latent space:

1. Encode the high-dimensional PDE state $u(𝐱,t)$ into a low-dimensional latent vector $𝐳(t)\in {ℝ}^d$ using a trained encoder.

2. Evolve the latent state using a lightweight temporal model: $𝐳(t+\mathrm{\Delta }t)=f_{\theta }(𝐳(t))$. This step is fast because $d\ll N$ (the number of grid points).

3. Decode back to the physical state: $\hat{u}(𝐱,t+\mathrm{\Delta }t)=g_{\psi }(𝐳(t+\mathrm{\Delta }t))$.

The challenge is ensuring that the latent dynamics respect the physics: conservation laws, stability, and long-time accuracy.

Hardware co-design: neural PDE accelerators

Generative PDE Foundation Models

“Can we build one model that understands all of physics?” - Animashree Anandkumar, 2024

The success of foundation models in natural language (GPT, Claude, Gemini) and vision (DALL-E, Stable Diffusion) naturally raises the question: can we build a foundation model for PDEs? A single, large, pre-trained model that, given a PDE specification (equations, domain geometry, initial/boundary conditions, material parameters), generates accurate solution fields-without any task-specific fine-tuning.

The vision: one model for all PDEs

Scaling laws: does more data help?

In language modelling, the “scaling laws” of Kaplan et al. (2020) and Hoffmann et al. (2022) showed that test loss decreases as a power law in model size, dataset size, and compute budget. Do similar scaling laws hold for PDE foundation models?

Current state of the art

Several groups have taken first steps towards PDE foundation models:

• MPP (Multiple Physics Pretraining, McCabe et al., 2023): Pre-trained on 14 PDE families, demonstrated zero-shot transfer to unseen PDEs, achieving competitive accuracy with specialised models.

• Unisolver (Unified PDE Solver with PDE-Conditioned Transformers, 2024): Conditions the Transformer attention mechanism directly on the PDE coefficients via “physics cross-attention,” enabling a single model to handle diverse PDE families.

• PROSE-FD (Pre-trained Neural Operator for PDE Families, 2024): Learns symbolic representations of PDE operators and conditions the neural solver on these representations.

• PDEformer (2024): Uses a tokenised representation of PDE specifications (akin to a “PDE language”) to condition a Transformer-based solver.

What's missing

Despite these promising beginnings, significant gaps remain:

Physical Guarantees and Trustworthiness

If we are to deploy neural PDE solvers in applications where lives or infrastructure are at stake-nuclear reactor safety, aircraft design, climate policy-we need guarantees, not just empirical performance metrics.

Conservation laws

Classical numerical schemes-finite volumes, symplectic integrators, structure-preserving discretisations-are carefully designed to respect these conservation laws exactly (at the discrete level). Neural surrogates, trained by minimising a data-driven loss, generally do not conserve these quantities. The mass, energy, or momentum of a neural solution may drift over long time horizons, leading to physically meaningless predictions.

Several approaches are being explored:

1. Hard constraints: Design the neural architecture so that conservation laws are satisfied by construction. For divergence-free fields, use curl-based parameterisations $𝐮=\nabla \times 𝝍$; for energy-preserving dynamics, use Hamiltonian neural networks.

2. Soft constraints: Add conservation-law penalties to the loss function. This is simpler but provides no guarantees: the constraints are satisfied approximately, not exactly.

3. Projection methods: After each neural prediction, project the output onto the constraint manifold (e.g., onto the space of divergence-free fields). This guarantees satisfaction of the constraint at the cost of additional computation.

4. Hybrid solvers: Use the neural model for fast prediction and a classical solver for correction, in a predict-correct cycle.

A posteriori error estimation

The key idea is to exploit the residual: even without knowing the true solution $u_{\text{true}}$, we can evaluate the PDE residual $r(𝐱,t)=𝒩[{\hat{u}}_{\theta }](𝐱,t)$, which should be zero for a perfect solution. Classical a posteriori error estimation theory relates the norm of the residual to the norm of the error, provided stability estimates for the PDE are available.

Out-of-distribution detection

High-Dimensional PDEs and the Curse of Dimensionality

Most of this chapter has focused on PDEs in two or three spatial dimensions-the “natural” setting for fluid dynamics, electromagnetism, and elasticity. But some of the most important PDEs in science and engineering live in much higher dimensions.

Where high-dimensional PDEs arise

• Quantum mechanics: The Schrödinger equation for $N$ particles lives in $3N$ spatial dimensions. For a molecule with 100 electrons, this is a PDE in 300 dimensions.

• Optimal control and finance: The Hamilton–Jacobi–Bellman (HJB) equation for an optimal control problem with $d$ state variables is a PDE in $d$ dimensions. Realistic portfolio optimisation problems can have $d\sim 100$.

• Mean-field games: Models of large populations of interacting agents (traffic flow, epidemic spreading, economic markets) lead to coupled systems of high-dimensional PDEs.

• Kinetic theory: The Boltzmann equation for a gas lives in 6-dimensional phase space (3 positions + 3 velocities), and its numerical solution is already at the edge of feasibility.

Classical grid-based methods are hopeless in high dimensions: a grid with $M$ points per dimension requires $M^d$ total points, which exceeds the number of atoms in the observable universe for $d>15$ and $M=100$. This is the infamous curse of dimensionality.

Deep BSDE methods

One of the most striking success stories of deep learning for PDEs is the deep BSDE method of E, Han, and Jentzen (2017), which exploits a beautiful connection between parabolic PDEs and stochastic processes.

The deep BSDE method parameterises the gradient $\nabla u$ at each time step by a neural network and optimises the terminal condition $u(𝐱,T)\approx g(𝐱)$ along sampled stochastic trajectories. Crucially, the method's computational cost scales polynomially in dimension $d$-breaking the curse of dimensionality.

The potential of generative models

Can generative models go further than the deep BSDE method? Several tantalising directions exist:

1. Score-based solvers for HJB equations: The Hamilton–Jacobi–Bellman equation can be viewed as defining a score function (the gradient of the value function) on a high-dimensional state space. Score-based generative models are designed precisely to learn such gradients-can they be repurposed as HJB solvers?

2. Mean-field generative models: For mean-field games, the coupled PDE system involves a forward Fokker–Planck equation and a backward HJB equation. A generative model could learn the joint distribution of trajectories, from which both the density and the value function can be extracted.

3. Quantum many-body wavefunctions: The Schrödinger equation for many particles is a high-dimensional PDE. Recent work (PauliNet, FermiNet, DiffusionQMC) uses neural networks to represent many-body wavefunctions. Generative models could sample from the Born probability distribution $|\psi (𝐱){|}^2$ directly.

Climate and Weather: The Ultimate Test

“Prediction is very difficult, especially if it's about the future.” - Niels Bohr (attributed)

If there is one domain where the promise and peril of generative PDE solvers converge with maximum intensity, it is weather and climate prediction. The atmosphere is governed by the Navier–Stokes equations coupled with thermodynamics, radiation transfer, moisture transport, and ocean-atmosphere interaction. The system is chaotic, multi-scale, and of existential importance to humanity.

Weather prediction as a PDE problem

Numerical Weather Prediction (NWP) has been one of the great scientific and computational achievements of the twentieth century. The basic approach-discretise the atmospheric equations on a global grid and time-step them forward-has been refined over 70 years at operational centres like ECMWF, NOAA, and the Met Office. The current state of the art, ECMWF's IFS (Integrated Forecasting System), uses a spectral transform method at roughly 9,km horizontal resolution with 137 vertical levels, producing 51-member ensemble forecasts that take approximately one hour of wall-clock time on a dedicated supercomputer.

The AI weather revolution

In 2022–2024, a series of landmark results showed that AI models can match or exceed the accuracy of IFS for medium-range weather forecasting:

• Pangu-Weather (Huawei, 2023): A 3D Swin Transformer trained on 39 years of ERA5 reanalysis data, producing 5-day global forecasts in seconds.

• GraphCast (DeepMind, 2023): A graph neural network operating on an icosahedral mesh, achieving state-of-the-art accuracy on 90% of 1,380 verification targets.

• GenCast (DeepMind, 2024): A conditional diffusion model that generates probabilistic ensemble forecasts, providing not just a single prediction but a distribution over possible weather outcomes. GenCast outperformed ECMWF's operational ensemble on probabilistic skill metrics.

Climate modelling: the century-scale challenge

Weather prediction operates on a time horizon of days to weeks. Climate modelling asks a much harder question: what will the statistical distribution of weather look like decades or centuries from now, under different greenhouse gas emission scenarios?

This is a problem of a fundamentally different character. Weather prediction is an initial-value problem (given today's atmospheric state, predict tomorrow's). Climate modelling is a boundary-value problem in the space of forcing scenarios: given a trajectory of CO${}_2$ concentrations, predict the long-term statistical equilibrium of the climate system.

A Grand Challenge: Twenty Open Questions

“The formulation of a problem is often more essential than its solution, which may be merely a matter of mathematical or experimental skill.” - Albert Einstein

We now present twenty concrete open research questions. They range from deep theoretical problems that may take decades to resolve, to practical engineering challenges that a determined team could tackle in a year. We have marked each question with a difficulty rating: -1pt[defblue] (0,0) circle (3pt); for problems accessible to motivated undergraduates, -1pt[exorange] (0,0) circle (3pt);[exorange] (8pt,0) circle (3pt); for PhD-level research, and -1pt[exrcoral] (0,0) circle (3pt);[exrcoral] (8pt,0) circle (3pt);[exrcoral] (16pt,0) circle (3pt); for “Millennium Prize”-level challenges.

I. Theoretical Foundations

1. Convergence of function-space diffusion. -1pt[exorange] (0,0) circle (3pt);[exorange] (8pt,0) circle (3pt); Prove that score-matching in infinite-dimensional Hilbert spaces converges to the true score as the number of training samples and network capacity grow. What are the rates?

2. Approximation theory for neural operators. -1pt[exorange] (0,0) circle (3pt);[exorange] (8pt,0) circle (3pt); Determine tight upper and lower bounds on the number of parameters needed for a neural operator to approximate a solution map $𝒢:𝒜\to 𝒰$ to within $ϵ$ error. How do these bounds depend on the smoothness and nonlinearity of the PDE?

3. Navier–Stokes regularity via learned Lyapunov functionals. -1pt[exrcoral] (0,0) circle (3pt);[exrcoral] (8pt,0) circle (3pt);[exrcoral] (16pt,0) circle (3pt); Can a neural network discover a Lyapunov-type energy functional that proves global regularity (or yields a finite-time blow-up scenario) for the 3D Navier–Stokes equations?

4. Generative models and ergodic theory. -1pt[exorange] (0,0) circle (3pt);[exorange] (8pt,0) circle (3pt); For chaotic PDEs (e.g., turbulent Navier–Stokes), the long-time statistics are governed by an invariant measure. Can a diffusion model learn this invariant measure, and if so, does sampling from the model produce trajectories with the correct ergodic properties?

5. Complexity of PDE solving with neural networks. -1pt[exrcoral] (0,0) circle (3pt);[exrcoral] (8pt,0) circle (3pt);[exrcoral] (16pt,0) circle (3pt); Is there a complexity-theoretic separation between neural PDE solvers and classical solvers? Can neural networks solve certain PDE families in polynomial time (in dimension) where classical grid methods require exponential time?

II. Architectures and Methods

1. Equivariant diffusion on function spaces. -1pt[exorange] (0,0) circle (3pt);[exorange] (8pt,0) circle (3pt); Design a diffusion model on function space that is equivariant to the full symmetry group of the PDE (spatial translations, rotations, Galilean boosts, scaling). Does symmetry equivariance improve sample quality and data efficiency?

2. Autoregressive vs. diffusion for temporal rollouts. -1pt[defblue] (0,0) circle (3pt); Systematically compare autoregressive (one-step-at-a-time) and diffusion-based (full-trajectory) generation for time-dependent PDEs. Under what conditions does each approach dominate?

3. Latent-space dynamics with physical constraints. -1pt[exorange] (0,0) circle (3pt);[exorange] (8pt,0) circle (3pt); Develop an encoder-decoder architecture for PDE states where the latent dynamics provably conserve energy, momentum, and mass. Can symplectic or Hamiltonian structures be imposed in the latent space?

4. Tokenisation of PDE specifications. -1pt[defblue] (0,0) circle (3pt); What is the optimal way to represent a PDE specification (equations, domain, boundary conditions) as a sequence of tokens for a Transformer? Can symbolic tokenisation (parsing the PDE into an abstract syntax tree) outperform naive string tokenisation?

5. Multi-fidelity generative training. -1pt[defblue] (0,0) circle (3pt); Train a generative PDE model on a mix of cheap low-resolution and expensive high-resolution simulation data. Can multi-fidelity training achieve the accuracy of high-fidelity training at a fraction of the data-generation cost?

III. Applications and Deployment

1. Certified neural surrogates for nuclear safety. -1pt[exorange] (0,0) circle (3pt);[exorange] (8pt,0) circle (3pt); Develop a neural PDE surrogate for reactor thermal-hydraulics that comes with provable error bounds sufficient for regulatory certification.

2. Real-time turbulence for flight simulation. -1pt[exorange] (0,0) circle (3pt);[exorange] (8pt,0) circle (3pt); Build a generative model that produces realistic turbulent flow fields around an aircraft at interactive frame rates ($\ge 30$,fps), with statistical properties matching high-fidelity LES or DNS within specified tolerances.

3. Generative climate emulators. -1pt[exorange] (0,0) circle (3pt);[exorange] (8pt,0) circle (3pt); Train a generative model that produces physically consistent century-scale climate trajectories conditioned on emission scenarios, matching the statistical behaviour of CMIP6 models but running $\mathrm{10,000}\times$ faster.

4. Inverse design with generative PDE solvers. -1pt[defblue] (0,0) circle (3pt); Use a conditional diffusion model to solve inverse design problems: given a desired flow property (e.g., minimum drag), generate the geometry that achieves it. Compare with classical adjoint-based optimisation.

5. Patient-specific surgical planning. -1pt[exorange] (0,0) circle (3pt);[exorange] (8pt,0) circle (3pt); Build a generative cardiovascular flow solver that takes a patient's CT/MRI-derived vascular geometry, produces an ensemble of possible haemodynamic scenarios, and runs fast enough for intraoperative use.

IV. Benchmarks and Evaluation

1. A unified benchmark for generative PDE solvers. -1pt[defblue] (0,0) circle (3pt); Create a comprehensive benchmark suite (analogous to ImageNet for vision or GLUE for NLP) covering diverse PDE families, geometries, and evaluation metrics, with standardised train/test splits and evaluation protocols.

2. Perceptual metrics for PDE solutions. -1pt[defblue] (0,0) circle (3pt); Pointwise $L^2$ error fails to capture the perceptual quality of turbulent flow fields (analogous to how MSE fails for images). Develop physically motivated perceptual metrics that assess the statistical, spectral, and structural fidelity of generated PDE solutions.

3. Stress-testing out-of-distribution. -1pt[defblue] (0,0) circle (3pt); Systematically characterise the failure modes of neural PDE solvers under distribution shift: changes in Reynolds number, domain geometry, boundary conditions, and PDE coefficients. Build a “breaking point” dataset.

V. Connections and Unifications

1. PDE solving as language modelling. -1pt[exorange] (0,0) circle (3pt);[exorange] (8pt,0) circle (3pt); Is there a meaningful sense in which PDE solutions can be “narrated” as sequences, and autoregressive language models can serve as PDE solvers? What are the theoretical limitations of this analogy?

2. The grand unification: one model for science. -1pt[exrcoral] (0,0) circle (3pt);[exrcoral] (8pt,0) circle (3pt);[exrcoral] (16pt,0) circle (3pt); Can we build a single foundation model that simultaneously handles fluid dynamics, solid mechanics, electromagnetism, quantum mechanics, and their couplings? What are the fundamental obstacles-architectural, data-related, or information-theoretic-to such a model?

Philosophical Reflections

“The miracle of the appropriateness of the language of mathematics for the formulation of the laws of physics is a wonderful gift which we neither understand nor deserve.” - Eugene Wigner, 1960

We close this chapter-and this part of the book-with some reflections that are more philosophical than technical. They are offered not as definitive answers but as invitations to thought.

Does a Neural PDE Solver “Understand” Physics?

When a Fourier Neural Operator produces a flow field that matches DNS to within 2% relative error, does it “understand” fluid mechanics? The question may seem idle, but it has practical implications.

Consider two scenarios:

1. A student who memorises the answers to every problem in a fluid mechanics textbook without understanding the derivations. Ask them a problem they have seen before, and they answer correctly. Ask them a genuinely novel problem, and they fail.

2. A student who understands the Navier–Stokes equations deeply-the role of each term, the physical mechanisms, the dimensional analysis. They may not have memorised any answers, but they can derive solutions to problems they have never seen.

Current neural PDE solvers are closer to the first student. They excel at interpolation (problems “between” training examples) but struggle at extrapolation (problems genuinely outside the training distribution). A model trained on flows at Reynolds numbers ${10}^3$ to ${10}^5$ may fail catastrophically at ${10}^6$.

Data-Driven vs. Physics-Driven: The Central Tension

The field of scientific machine learning is animated by a fundamental tension between two philosophies:

• The data-driven philosophy: “Let the data speak. Given enough data and a sufficiently flexible model, the physics will emerge automatically. We do not need to encode physical laws; the model will discover them.”

• The physics-driven philosophy: “Physical laws are hard-won distillations of centuries of human understanding. Ignoring them is wasteful and dangerous. Encode them into the model architecture and loss function; use data only to fill in what physics alone cannot determine.”

The truth, as usual, lies somewhere in between-but the optimal balance is far from obvious and likely depends on the application.

Symbolic AI vs. Connectionist AI for Scientific Discovery

The tension between physics-driven and data-driven approaches is a specific instance of a much older debate in artificial intelligence: symbolic vs. connectionist AI.

Symbolic AI-exemplified by expert systems, theorem provers, and symbolic regression-represents knowledge as explicit rules and manipulates them logically. Connectionist AI-exemplified by neural networks-represents knowledge as distributed patterns of activation in high-dimensional weight spaces.

For scientific discovery, the ideal would be a system that combines:

1. The generalisation and pattern recognition abilities of connectionist models (finding structure in vast datasets that no human could survey);

2. The interpretability and compositionality of symbolic models (producing results that humans can understand, verify, and build upon);

3. The rigour of formal mathematical proof (ensuring that discovered relationships are not just correlations but provable truths).

Two Kinds of Unreasonable Effectiveness

In 1960, the physicist Eugene Wigner published his celebrated essay The Unreasonable Effectiveness of Mathematics in the Natural Sciences. He marvelled at the fact that abstract mathematical structures-developed by pure mathematicians with no physical application in mind-turn out, again and again, to be precisely the right language for describing physical reality. Riemannian geometry, conceived in the abstract, turned out to be the language of general relativity. Group theory, developed for its own beauty, became the foundation of particle physics.

Six decades later, we face an analogous puzzle: the unreasonable effectiveness of deep learning. Neural networks-universal function approximators with no built-in physical knowledge-turn out to be extraordinarily effective at learning the solutions of PDEs. A generic architecture (Transformer, U-Net) trained on examples can match or exceed decades of domain-specific numerical method development.

The Role of Human Intuition in an AI-Augmented World

Perhaps the most important question of all is not technical but human: in a world where AI can solve PDEs faster and more accurately than humans, what is the role of human intuition?

We believe the answer is: more important than ever.

The history of science suggests that the greatest discoveries occur not when a tool replaces human thought but when it extends it. The telescope did not replace the astronomer; it gave the astronomer new eyes. The computer did not replace the mathematician; it gave the mathematician new hands. Generative PDE solvers will not replace the physicist or engineer; they will give them new powers of exploration.

“We stand at the beginning of a new relationship between humans and the mathematical laws that govern the physical world. For three centuries, we have used our minds to derive solutions; for seventy years, we have used our computers to approximate them. Now, for the first time, we can learn them. The question is not whether AI will transform computational physics-it already has. The question is whether we will be wise enough to wield this power with the care it demands, and bold enough to ask the questions it makes possible.”

We close with a grand challenge and a conjecture, offered in the spirit of Hilbert's 1900 address and the Clay Millennium Problems, as a beacon for the next generation.

“The important thing is not to stop questioning. Curiosity has its own reason for existence.” - Albert Einstein

This is not the end. It is not even the beginning of the end. But it is, perhaps, the end of the beginning.

Pen-and-Paper Exercises

The exercises in this section are designed to solidify the theoretical foundations developed throughout the chapter. They range from straightforward calculations (3570END) to demanding proofs that require combining ideas from multiple sections (3571END). We encourage the reader to attempt each exercise before consulting any hints, and to write complete, rigorous solutions.

Coding Projects

The following projects are designed to be completed in two to four weeks each, assuming familiarity with PyTorch (or JAX) and basic numerical PDE methods. Each project specifies the expected deliverables, suggested architecture, and evaluation criteria. Projects marked with $†$ require GPU access; all others can be completed on a laptop CPU (though GPU acceleration is always welcome).

Challenge Problems

The following challenge problems are deliberately open-ended and sit at the frontier of current research. There is no single “correct” answer; instead, the goal is to develop creative solutions, identify fundamental limitations, and-ideally-push the state of the art. Each problem could form the basis of a research paper or a master's thesis.