Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Structured Simulators and Neural Operators › Operator Learning for Fields and PDEs

Michael BrenndoerferJune 22, 202610 min read

World Models HandbookAll chapters

Part 4 of 12Space, Physics, and Causality

Chapter 22Structured Simulators and Neural Operators › Operator Learning for Fields and PDEs

Reading section 6 of 11 · Operator Learning for Fields and PDEs · 10 minSection 6 of 11 · 10 min

Section 6 of 11 · 10 min read

Structured Simulators and Neural Operators › Operator Learning for Fields and PDEsOne section · Sections in this chapter6 / 11 ↓

Sections in this chapter11

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

Everything so far concerns systems represented with finite-dimensional state. Continuum fields, such as fluid velocity, material temperature, or ocean pressure, are functions governed by partial differential equations. A numerical discretization with 10 million unknowns is a large finite-dimensional ODE or DAE approximation of that continuum problem. A fixed-grid architecture may couple its learned representation to one resolution or geometry, while convolutional, graph, interpolation, multigrid, and operator constructions can share parameters across compatible discretizations. Whether they transfer depends on coordinates, normalization, geometry handling, numerical resolution, and the training distribution.

Operator learning reframes the learning target. For a PDE with solution $u(x, t)$ , the object of interest is often not one trajectory but the solution operator: the map

G: u_0(\cdot) \;\longmapsto\; u(\cdot, T),

which takes an entire initial condition to the entire solution at time $T$ . A neural operator is designed to approximate this function-space map from sampled representations. Appropriately constructed operators can share parameters across discretizations and sometimes evaluate zero-shot at new resolutions. That capability is not automatic: standard FFT-based FNOs assume regular grids, while new geometries need compatible coordinates, meshes, or geometry-aware variants and remain distribution-dependent.

Discretization invariance

The use of one learned parameterization across compatible discretizations of the same underlying function-space problem. A well-designed method should converge toward a consistent continuum operator as adequately resolved meshes are refined; it need not have equal accuracy on arbitrary grids, geometries, or resolutions.

Two foundational architectures illustrate different answers to the question "how does a network eat a function?"

DeepONet (Lu and colleagues, 2019) factorizes the output function evaluation as a learned inner product:

G(u)(y) \approx \sum_{k=1}^{p} b_k(u)\, t_k(y),

where:

$G(u)(y)$ : the operator output at the query location $y$ for the input function $u$
$b_k(u)$ : the $k$ -th coefficient the branch network produces from the entire input function $u$
$t_k(y)$ : the $k$ -th feature the trunk network produces at the query location $y$
$p$ : the number of coefficient-feature pairs, which sets the expressiveness of the factorization

A branch network encodes sampled values of the input function $u$ into coefficients $b_k$ , and a trunk network encodes each supported query location $y$ into features $t_k$ . Their dot product lets one trained model evaluate new query points within the represented domain, subject to the sensor layout and training distribution. The factorization resembles a learned basis expansion: the trunk provides basis-like spatial functions and the branch provides input-dependent coefficients. Universal approximation results apply under compactness and continuity assumptions; they are existence results, not finite-data generalization guarantees.

The Fourier Neural Operator (FNO, Li and colleagues, 2020) builds the nonlocality into a single layer. Its core is the spectral convolution:

(K v)(x) = \mathcal{F}^{-1}\!\Big( R \cdot \mathcal{F}\, v \Big)(x),

where:

$v$ : the input field (for example, the feature values on the grid at the current layer)
$\mathcal{F}$ : the Fourier transform, mapping the field to its frequency modes
$R$ : a learnable complex weight applied to each frequency mode, with the highest modes truncated
$\mathcal{F}^{-1}$ : the inverse Fourier transform, mapping the weighted modes back to the grid
$(K v)(x)$ : the output of the spectral convolution at grid location $x$

Here $\mathcal{F}$ is the Fourier transform, $R$ is a learnable complex weight applied to each frequency mode (with the highest modes truncated), and $\mathcal{F}^{-1}$ transforms back. Multiplication by a learned $R$ in frequency space is a convolution on the grid, and unlike a local CNN kernel it can couple every point to every other point in one operation. This global path is useful for solution operators with nonlocal dependencies, including positive-time diffusion on the periodic domain used below; not every PDE or forecast horizon requires all-to-all influence. A full FNO stacks several layers of the form

v_{l+1}(x) = \sigma\!\Big( W v_l(x) + (K v_l)(x) \Big),

where:

$v_l(x)$ : the layer- $l$ feature value at grid location $x$
$W$ : a local linear map applied pointwise at each location
$(K v_l)(x)$ : the global spectral path from the previous equation, coupling every location to every other
$\sigma$ : a pointwise nonlinearity applied after summing the local and global paths

The layer combines the global spectral path with a local linear path and a pointwise nonlinearity. Nonlinear PDEs couple modes, while a purely linear spectral map remains confined to linear transformations regardless of depth. The pointwise nonlinearity lets successive spectral layers compose nonlinear operators.

Why is the spectral view so natural for PDEs? Because many linear PDE solution operators are diagonal in the Fourier basis. The heat equation $u_t = \nu u_{xx}$ on a periodic domain multiplies each wavenumber- $k$ Fourier coefficient by $\exp(-\nu k^2 T)$ . A spectral architecture can represent this multiplier exactly on its retained modes; an architecture that looks at one grid point at a time cannot express the full map for arbitrary fields, because $u(x,T)$ depends on the entire initial profile. When you choose an architecture for a field-valued world, you decide which dependency structures are available to the model.

Weather forecasting provides large-scale examples. In their 2022--2023 evaluations, FourCastNet, based on an adaptive Fourier operator, matched then-current IFS performance on selected short-lead variables, while GraphCast, a multimesh graph model, beat then-current HRES on most specified verification targets. These were benchmark-, lead-time-, hardware-, and baseline-specific results rather than universal dominance. FourCastNet reported a week-long forecast in under two seconds on its stated accelerator setup; GraphCast reported a ten-day forecast in under 60 seconds on a TPU. Both timings describe inference and exclude training and operational pipeline costs.

Let us see the argument play out on the heat equation. We generate discrete spectral solutions, then train two models on the same data: a minimal spectral operator and a pointwise MLP that sees only $u_0(x)$ at each location. The spectral architecture can represent the heat multiplier on its retained modes; modes above its cutoff are present but strongly attenuated in the targets. Both models use the same data and number of training steps. The spectral model has 2,073 trainable real parameters and the pointwise MLP has 4,353, so this comparison isolates the dependency structure more clearly than it isolates parameter count.

In[21]:

Code

import numpy as np

rng = np.random.default_rng(5)

N_grid, nu_heat, T_final = 256, 0.1, 1.0
k_modes = np.arange(N_grid // 2 + 1)


def make_heat_sample():
    """Random smooth initial condition and its exact heat-equation solution at T."""
    c0 = rng.normal(size=len(k_modes)) + 1j * rng.normal(size=len(k_modes))
    c0 = c0 / (1.0 + k_modes**2) ** 1.2
    c0[0] = 0.0
    u0 = np.fft.irfft(c0, n=N_grid)
    uT = np.fft.irfft(c0 * np.exp(-nu_heat * k_modes**2 * T_final), n=N_grid)
    return u0, uT


n_train, n_test = 200, 50
train_u0 = np.zeros((n_train, N_grid))
train_uT = np.zeros((n_train, N_grid))
test_u0 = np.zeros((n_test, N_grid))
test_uT = np.zeros((n_test, N_grid))
for i in range(n_train):
    train_u0[i], train_uT[i] = make_heat_sample()
for i in range(n_test):
    test_u0[i], test_uT[i] = make_heat_sample()

The spectral model lifts each grid value to a small feature width, applies a learnable complex weight per retained Fourier mode, and projects back. The random initial conditions contain nonzero coefficients above mode 15, but heat evolution suppresses them strongly at $T=1$ . The model can represent the exact heat multiplier on the retained modes and approximate the full target closely; the cutoff remains a real modeling choice rather than an exact support boundary.

In[22]:

Code

import torch
import torch.nn as nn

torch.manual_seed(0)


class SpectralOperator(nn.Module):
    """Minimal FNO-style layer: lift, learned per-mode weights, project."""

    def __init__(self, modes=16, width=8):
        super().__init__()
        self.modes = modes
        self.lift = nn.Linear(1, width)
        self.W = nn.Parameter(0.05 * torch.randn(modes, width, width, 2))
        self.proj = nn.Linear(width, 1)

    def forward(self, u):  # u: (batch, N_grid)
        h = self.lift(u.unsqueeze(-1))  # (batch, N, width)
        h_ft = torch.fft.rfft(h, dim=1)  # (batch, N//2+1, width)
        out_ft = torch.zeros_like(h_ft)
        Wc = torch.complex(self.W[..., 0], self.W[..., 1])
        out_ft[:, : self.modes] = torch.einsum(
            "bkm,kmo->bko", h_ft[:, : self.modes], Wc
        )
        h = torch.fft.irfft(out_ft, n=u.shape[-1], dim=1)
        return self.proj(h).squeeze(-1)


train_u0_t = torch.tensor(train_u0, dtype=torch.float32)
train_uT_t = torch.tensor(train_uT, dtype=torch.float32)

spectral_model = SpectralOperator(modes=16, width=8)
opt = torch.optim.Adam(spectral_model.parameters(), lr=5e-3)
spectral_losses = []
for step in range(1000):
    pred = spectral_model(train_u0_t)
    loss = torch.mean((pred - train_uT_t) ** 2)
    opt.zero_grad()
    loss.backward()
    opt.step()
    spectral_losses.append(loss.item())

The pointwise baseline has more parameters but the wrong dependency structure: its output at $x$ can depend only on its input at $x$ . Watch what that does to the error.

In[23]:

Code

pointwise_model = nn.Sequential(
    nn.Linear(1, 64),
    nn.GELU(),
    nn.Linear(64, 64),
    nn.GELU(),
    nn.Linear(64, 1),
)
opt = torch.optim.Adam(pointwise_model.parameters(), lr=5e-3)
pointwise_losses = []
for step in range(1000):
    pred = pointwise_model(train_u0_t.unsqueeze(-1)).squeeze(-1)
    loss = torch.mean((pred - train_uT_t) ** 2)
    opt.zero_grad()
    loss.backward()
    opt.step()
    pointwise_losses.append(loss.item())

In[24]:

Code

def rel_l2(model, u0, uT, pointwise=False):
    u0_t = torch.tensor(u0, dtype=torch.float32)
    uT_t = torch.tensor(uT, dtype=torch.float32)
    with torch.no_grad():
        x_in = u0_t.unsqueeze(-1) if pointwise else u0_t
        pred = model(x_in).squeeze(-1)
    err = torch.mean((pred - uT_t) ** 2, dim=1)
    ref = torch.mean(uT_t**2, dim=1)
    return torch.sqrt(err / ref).mean().item()


err_spectral = rel_l2(spectral_model, test_u0, test_uT)
err_pointwise = rel_l2(pointwise_model, test_u0, test_uT, pointwise=True)

with torch.no_grad():
    spec_pred = spectral_model(
        torch.tensor(test_u0[:1], dtype=torch.float32)
    ).numpy()[0]
    pw_pred = (
        pointwise_model(
            torch.tensor(test_u0[:1], dtype=torch.float32).unsqueeze(-1)
        )
        .squeeze(-1)
        .numpy()[0]
    )
    spectra_u0 = np.abs(np.fft.rfft(test_u0[0]))
    spectra_uT = np.abs(np.fft.rfft(test_uT[0]))

Out[25]:

Console

Test relative L2 error, spectral operator: 0.0011
Test relative L2 error, pointwise MLP:     0.2292

Out[26]:

Visualization

Curves showing the spectral operator prediction tracking the true solution while the pointwise prediction deviates. — Discrete spectral heat-equation solution at time T for a held-out initial condition, compared with the spectral operator and pointwise MLP. The spectral prediction closely tracks the target, while the pointwise prediction deviates despite having more trainable parameters and the same data.

Log-log spectra with the solution spectrum decaying faster than the initial spectrum. — Fourier amplitude spectra of the initial condition and the solution at time T. Diffusion damps high wavenumbers, so the solution at T depends on the whole initial profile, not on pointwise values.

The contrast is stark. The spectral operator tracks the target closely, while the larger pointwise MLP does not. The failure is an expressiveness problem imposed by the dependency structure: no function of the single number $u_0(x)$ can equal a value that depends on the entire field for arbitrary initial conditions. Diffusion damps high wavenumbers, so by time $T$ the output at one point is a weighted average of values across the periodic domain. A pointwise map has no mechanism for that averaging. This fixed-grid experiment isolates nonlocal representability; it does not test transfer to another resolution or geometry.