From solving one PDE at a time to learning the map from any input function to its corresponding output, the architecture behind amortised scientific computation.
Every notebook and article in this series up to now has trained a neural network to represent one solution to one problem instance. The network approximates $u(x,t;\theta^*)$ for a specific set of initial conditions, boundary conditions, and PDE parameters, and does it remarkably well. But change any of those inputs, even slightly, and you discard the trained weights and start from scratch.
Consider what this costs in practice. Design exploration over 1,000 initial condition variations for a pharmaceutical reactor: $1{,}000 \times 45\text{s} \approx 12.5$ hours of training. Uncertainty quantification via Monte Carlo over 10,000 parameter samples: weeks of GPU time. Real-time digital twin control with millisecond update requirements: completely out of reach. The PINN paradigm, however powerful, is fundamentally a per-instance solver.
The root cause is how PINNs are parameterised. The network maps coordinates $(x,t)$ to a solution value, it is a function approximator in the classical sense. What we need instead is a machine that takes the problem specification (the initial condition function, the forcing function, the parameter values) and outputs the corresponding solution function. That is not a function approximator. It is an operator approximator.
Before describing the architecture, the distinction must be made precise. A function maps a number to a number (or a finite-dimensional vector to a vector):
An operator maps a function to a function, it is a map between infinite-dimensional spaces:
The input $u$ is an entire function; the output $s$ is an entire function. Three canonical examples build intuition:
The differential operator: $\mathcal{L}[u](x) = \frac{du}{dx}$, maps a differentiable function to its derivative.
The integral (antiderivative) operator:
Given any integrable function $u$, the operator returns its antiderivative $s$ satisfying $s(0)=0$. This is the benchmark problem in Lu et al. 2021 and in the notebook for this article.
The solution operator of a PDE: Given an initial condition $u_0(x)$, the solution operator $\mathcal{G}$ returns the full space-time field $s(x,t)$ satisfying the governing equation. This is the most valuable example, it maps any admissible IC to the corresponding solution, amortising the cost of solving the PDE across all possible inputs.
Figure 1. The distinction between a function (maps numbers to numbers) and an operator (maps functions to functions). DeepONet approximates the latter.
The theoretical foundation for DeepONet is the Chen & Chen 1995 universal approximation theorem for operators. In its original form (IEEE Trans. Neural Netw., 1995), the theorem states:
The theorem guarantees that a network of the form $\sum_k b_k(u) \cdot t_k(y)$ can approximate any continuous operator to arbitrary precision. The $b_k$ functions are evaluated on the input function $u$ sampled at $m$ fixed sensors; the $t_k$ functions depend on the query location $y$ alone. This dot-product structure is not an architectural choice made for computational convenience, it emerges directly from the theoretical guarantee.
The approximation is in the uniform norm over the entire function space: the network must be simultaneously accurate for all admissible input functions $u$ and all query locations $y$. This is a much stronger guarantee than pointwise approximation, and it is what makes DeepONet a principled tool rather than an ad hoc regression.
Lu et al. 2021 (Nature Machine Intelligence) operationalise the Chen & Chen theorem into a practical neural architecture with two sub-networks connected by a dot product. The architecture is elegantly minimal yet expressive enough to handle a wide class of PDE solution operators.
Figure 2. DeepONet architecture (Lu et al. 2021). The branch network encodes the input function evaluated at $m$ fixed sensors; the trunk network encodes the query location $y$. The output is a dot product of the two embeddings plus a bias, a weighted sum of $p$ learned basis functions.
The branch network takes the input function $u$ evaluated at $m$ fixed sensor locations $\{x_1, \ldots, x_m\}$ as a finite-dimensional vector $[u(x_1), \ldots, u(x_m)] \in \mathbb{R}^m$. It encodes this into a coefficient vector $\mathbf{b}(u) \in \mathbb{R}^p$. The branch can be any feed-forward architecture; a three-layer tanh MLP with hidden width 80 is standard for the antiderivative benchmark. The branch weights encode information about the entire input function, effectively, they are the "context encoder."
The trunk network takes a scalar query location $y \in \mathbb{R}$ and encodes it into $\mathbf{t}(y) \in \mathbb{R}^p$. Crucially, the trunk's final activation is a Tanh (not linear), ensuring that the basis functions $\{t_k(y)\}$ are bounded. The trunk learns a set of $p$ basis functions over the output domain, what amounts to a learned, PDE-adapted basis analogous to Fourier modes or POD eigenfunctions, but without any prior specification.
The output is:
The scalar $b_0$ is a learnable bias. This dot product is a weighted combination of the $p$ trunk basis functions, where the weights $\{b_k(u)\}$ are determined by the input function via the branch network. The analogy with a functional basis expansion is exact: if the trunk learned the Fourier basis $\{sin(k\pi y), cos(k\pi y)\}$, the branch would be predicting Fourier coefficients. But the basis is learned jointly with the coefficients, which is what makes DeepONet so flexible.
The canonical benchmark from Lu et al. 2021 is the antiderivative operator: $\mathcal{G}[u](x) = \int_0^x u(\tau)\,d\tau$, $s(0)=0$. Input functions are smooth random draws from a Gaussian process with squared-exponential kernel and length scale 0.2. This produces a rich diversity of smooth but highly variable functions, exactly the kind of distribution over which we want the operator to generalise.
The dataset generation uses a Gaussian process to sample random input functions, then computes their antiderivatives via cumulative trapezoid integration. The GP covariance structure ensures the generated functions are smooth (controlled by the length scale) but cover a wide range of shapes.
import torch, torch.nn as nn import numpy as np torch.manual_seed(0); np.random.seed(0) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") m = 100 # number of sensor locations x_grid = np.linspace(0, 1, m) def random_function(n=1, length_scale=0.2): # Sample smooth functions via a squared-exponential GP kernel X = x_grid.reshape(-1, 1) K = np.exp(-(X - X.T)**2 / (2*length_scale**2)) + 1e-8*np.eye(m) L = np.linalg.cholesky(K) return (L @ np.random.randn(m, n)).T # (n, m) def antiderivative(u): # Cumulative trapezoid, s(0) = 0 dx = x_grid[1] - x_grid[0] s = np.zeros_like(u) s[:, 1:] = np.cumsum((u[:, 1:] + u[:, :-1]) / 2 * dx, axis=1) return s def make_dataset(N): U = random_function(N) # (N, m) input functions S = antiderivative(U) # (N, m) target functions Q = 30 # query points per function idx = np.random.randint(0, m, size=(N, Q)) y = x_grid[idx] # (N, Q) s_y = np.take_along_axis(S, idx, axis=1) # (N, Q) # Flatten into (N*Q) training triples branch = np.repeat(U, Q, axis=0) # (N*Q, m) trunk = y.reshape(-1, 1) # (N*Q, 1) target = s_y.reshape(-1, 1) # (N*Q, 1) return (torch.tensor(branch, dtype=torch.float32, device=device), torch.tensor(trunk, dtype=torch.float32, device=device), torch.tensor(target, dtype=torch.float32, device=device)) br_tr, tr_tr, y_tr = make_dataset(1000) print("training triples:", br_tr.shape[0]) # → 30,000
The PyTorch implementation directly mirrors the mathematical structure. The branch and trunk are three-layer MLPs with tanh activations; the final trunk layer also applies tanh to keep the basis functions bounded. The forward pass is a single dot product plus bias:
class DeepONet(nn.Module): def __init__(self, m, p=40, width=80): super().__init__() self.branch = nn.Sequential( nn.Linear(m, width), nn.Tanh(), nn.Linear(width, width), nn.Tanh(), nn.Linear(width, p) ) self.trunk = nn.Sequential( nn.Linear(1, width), nn.Tanh(), nn.Linear(width, width), nn.Tanh(), nn.Linear(width, p), nn.Tanh() # final tanh: bounds basis functions ) self.b0 = nn.Parameter(torch.zeros(1)) def forward(self, u_sensors, y): b = self.branch(u_sensors) # (N, p) — coefficients t = self.trunk(y) # (N, p) — basis functions return torch.sum(b * t, dim=1, keepdim=True) + self.b0 # (N, 1) model = DeepONet(m).to(device) print("params:", sum(p.numel() for p in model.parameters())) # → 27,681 parameters
Training uses Adam with random mini-batches of 4,096 triples (out of 30,000). The loss is simple mean squared error between predicted and true antiderivative values. Over 8,000 epochs the loss falls three orders of magnitude:
opt = torch.optim.Adam(model.parameters(), lr=1e-3) EPOCHS, bs = 8000, 4096 N = br_tr.shape[0] for ep in range(EPOCHS): perm = torch.randperm(N, device=device)[:bs] opt.zero_grad() pred = model(br_tr[perm], tr_tr[perm]) loss = torch.mean((pred - y_tr[perm])**2) loss.backward(); opt.step() if ep % 1500 == 0: print(f"epoch {ep:5d} | loss {loss.item():.3e}") # epoch 0 | loss 2.276e-01 # epoch 1500 | loss 3.719e-04 # epoch 3000 | loss 1.387e-04 # epoch 4500 | loss 4.582e-05 # epoch 6000 | loss 4.642e-05 # epoch 7500 | loss 3.076e-05
The network has only 27,681 parameters, which is small relative to the complexity of the operator it learns. This is a key advantage: the parameterisation is efficient because the operator has a compact representation in the learned basis.
The output $\sum_k b_k(u) \cdot t_k(y)$ is a functional basis expansion, a weighted sum of $p$ functions $\{t_k(y)\}$ over the output domain, where the weights $\{b_k(u)\}$ depend on the input function. This structure has a long history in scientific computing:
| Method | Basis $\{t_k\}$ | Coefficients $\{b_k\}$ | Key property |
|---|---|---|---|
| Fourier Series | $\sin(k\pi y), \cos(k\pi y)$, fixed | Computed analytically from $u$ | Universal for periodic functions |
| POD / PCA | Eigenfunctions of empirical covariance, data-driven, fixed post-training | Projected coefficients | Optimal $L^2$ basis for the training ensemble |
| DeepONet | Trunk outputs, learned jointly | Branch outputs, learned jointly | Jointly optimal; adapts to the PDE structure |
The key advantage over POD is that the DeepONet basis is not constrained to be linear in the input. The branch network can represent highly nonlinear dependence of the coefficients on the input function. The trunk learns basis functions that need not be orthogonal or satisfy any a priori constraint, they adapt to whatever representation minimises the total loss across the training distribution of input functions.
The critical test is evaluation on input functions that were never seen during training. The following code evaluates the trained model on three freshly sampled GP functions and computes the relative $L^2$ error:
U_test = random_function(3) # 3 new, unseen functions S_test = antiderivative(U_test) y_full = torch.tensor(x_grid, dtype=torch.float32, device=device).view(-1, 1) for i in range(3): br = torch.tensor(np.repeat(U_test[i:i+1], m, axis=0), dtype=torch.float32, device=device) with torch.no_grad(): pred = model(br, y_full).cpu().numpy().ravel() err = np.linalg.norm(pred - S_test[i]) / np.linalg.norm(S_test[i]) print(f"Unseen function #{i+1}: relative L2 error = {err:.2e}") # Unseen function #1: relative L2 error ≈ 5.1e-03 # Unseen function #2: relative L2 error ≈ 3.8e-03 # Unseen function #3: relative L2 error ≈ 7.2e-03
Relative $L^2$ errors of $10^{-3}$ to $10^{-2}$ on completely unseen input functions represent genuine operator generalisation, the network has learned something structural about the antiderivative operation, not just memorised the training set. More importantly, each evaluation takes a single forward pass through the 27,681-parameter network, requiring less than 1 millisecond on CPU.
Compare this to training a PINN for each new input function from scratch: 30–60 seconds of gradient descent, every time. For a design sweep over 1,000 initial conditions, this difference is the distinction between a 1.5-second deployment and a 12-hour batch job.
| Workflow | Cost per query | Cost for N = 1000 design sweep | Notes |
|---|---|---|---|
| PINN, single solve | ~45 s | ~45,000 s (12.5 h) | Full retraining required for each new IC/BC/parameter |
| DeepONet, offline training | ~600 s (once) | ~600 s (once) | Paid once; amortised over all future queries |
| DeepONet, inference | <1 ms | ~1.5 s | Single forward pass; GPU parallelism makes it faster still |
| Classical FEM/FD solver | 1–10 s (simple) | 1,000–10,000 s | Accurate but scales linearly with N; no generalisation |
Figure 3. Logarithmic time comparison for N = 1,000 design sweep. DeepONet amortises the training cost across all inferences; total query time drops from 12.5 hours (PINN) to 1.5 seconds.
For parabolic and hyperbolic PDEs, the trunk query is extended to include both space and time: $y \to (x, t) \in \mathbb{R}^2$. The trunk becomes a 2-input MLP encoding spatial-temporal basis functions. The branch continues to encode the initial condition evaluated at $m$ spatial sensors. The antiderivative benchmark generalises: the operator maps $u_0(x)$ to the full space-time field $s(x,t)$. For the heat equation, for example, a single DeepONet handles all initial conditions simultaneously after training.
For systems with multiple output fields, for example, the velocity components $(u, v)$ and pressure $p$ in Navier-Stokes, the output layer is extended to $q$ scalar outputs, each with its own dot product: $\mathcal{G}_j(u)(y) = \mathbf{b}_j(u) \cdot \mathbf{t}(y) + b_{0,j}$ for $j = 1, \ldots, q$. The trunk can be shared (same basis functions for all outputs) or separate (different bases per output). Shared trunk reduces parameter count; separate trunks allow each field to develop its own optimal basis.
Wang et al. 2022 (arXiv:2204.13188) identified that the standard DeepONet initialization leads to gradient pathologies similar to those observed in PINNs. Their "Modified DeepONet" uses a gating mechanism between two encoders (analogous to the Modified MLP for PINNs), achieving more stable training and better accuracy with fewer epochs. The key modification introduces encoder outputs $U = \sigma(W_1 x + b_1)$ and $V = \sigma(W_2 x + b_2)$, then gates each hidden layer: $H^{(k)} = \sigma(W^{(k)} H^{(k-1)} + b^{(k)}) \odot U + (1 - \sigma(W^{(k)} H^{(k-1)} + b^{(k)})) \odot V$.
Cai et al. 2021 extended DeepONet to multi-physics problems with multiple coupled equations. The "DeepM&Mnet" architecture uses a network for each physical field, with cross-field coupling through shared branch encodings. This is directly applicable to pharmaceutical process models where heat transfer, mass transfer, and reaction kinetics must be solved simultaneously.
DeepONet is one of two dominant operator learning families that emerged around 2021. The Fourier Neural Operator (FNO, Li et al. 2021) takes a fundamentally different approach: it applies convolutions in Fourier space to learn the operator. Understanding the tradeoffs is essential for choosing the right tool.
| Property | DeepONet | FNO (Fourier Neural Operator) |
|---|---|---|
| Theoretical basis | Chen & Chen 1995 UAT for operators | Universal approximation via Fourier convolution |
| Inductive bias | Branch/trunk factorisation; basis expansion | Translation equivariance; spectral convolution |
| Discretisation dependence | Mesh-free (trunk is continuous in $y$) | Tied to the training grid resolution |
| Sensor placement | Fixed sensor locations for branch input | Regular grid assumed |
| Data requirement | Smaller datasets often sufficient | Benefits from large datasets; excels for turbulence-like problems |
| Computational cost | Low; small parameter count for simple operators | Higher; FFT overhead but very fast for large grids |
| Best suited for | Low-dimensional output domains, sensor-based inputs, time-continuous queries | High-resolution 2D/3D PDE fields; climate, turbulence, fluid dynamics |
| Physics-informed variant | PI-DeepONet (Wang et al. 2021) | PINO (Li et al. 2021) |
Beyond these two, the neural operator landscape includes the Graph Neural Operator (GNO, Li et al. 2020) which handles irregular geometries via message-passing graphs; the Spherical FNO (SFNO, Bonev et al. 2023) adapted for global climate modeling on the sphere; and the GNOT (Han et al. 2022) which uses attention mechanisms for general operator learning. Kovachki et al. 2023 (JMLR) provides a comprehensive theoretical unification of all these methods under the neural operator framework.