Neural Operators and DeepONet

Post 4 built a network that steps the wave equation forward in time, and ended by naming resolution-independent neural-operator methods as something it would not go into. This post goes into it. It starts with the mathematical object these networks are trying to approximate, then states the one property that separates a neural operator from the network of Post 4, and then works through DeepONet as the first architecture built around that property.

Take the wave equation from Post 3 with an initial condition. Post 4 showed that displacement alone does not determine a wave's future, so the initial condition consists of two fields, displacement and velocity, given over the whole domain. Call that pair $a$.

Given $a$, the wave equation has exactly one solution. That is the fact the following definition rests on:

where $a$ denotes the initial condition as a whole function, $T$ denotes some later time, $u(\cdot ,T)$ denotes the whole solution field at that time (the dot marks the spatial argument as left open, so this is a function of space, not a single number), $X$ denotes the set of input functions we allow, $Y$ denotes the set of output functions, and $\mathcal{G}$ denotes the solution operator, the rule that takes any starting configuration to the field it turns into.

It is called an operator instead of a function because its input and its output are both entire functions rather than numbers. That is the whole meaning of the word here.

$T$ is not required to be far away. Setting $T=\mathrm{\Delta }t$ gives a solution operator that advances the state by one small step, where $\mathrm{\Delta }t$ denotes the step size, and applying it repeatedly advances the state further. Both readings are solution operators. This matters because a common first impression is that the point of a neural operator is to skip the rollout and land on a distant time in one shot. Some are trained that way and some are not. The Fourier Neural Operator, covered in the next post, is usually trained on a single small step and then rolled out exactly like the network from Post 4.

$\mathcal{G}$ is not something machine learning invented. It comes out of classical PDE analysis and exists whether or not anyone ever trains a network on it. The definition above says which function $\mathcal{G}(a)$ is, namely the one that satisfies the equation with that starting data. It does not say how to compute it, and for most equations nobody can write it down.

Three things have to be true before the operator is even well defined, and together they are called well-posedness: a solution has to exist for every input we allow, it has to be the only one, and a small change to the input has to cause only a small change to the output. The third one is stability, and it is the one that tends to fail for inverse problems, where the task is to recover a cause from measured effects. That is why inverse methods usually cannot get away with returning a single answer and have to report uncertainty as well.

The symplectic Euler scheme from Post 6 does not compute $\mathcal{G}$. It produces grid values $U_i^n$ that approximate it, with the error $e_i^n:=U_i^n-u_i^n$ that Posts 5 and 6 analysed, and it converges to the true solution as the grid gets finer, as long as the stability condition holds.

A network trained on PDE data is doing the same thing from a different direction. It approximates the same $\mathcal{G}$, fitted from examples instead of derived from a stencil, and those examples come from running the solver. One run per input function. Since the solver is the expensive part, the number of runs, not the speed of the network, is what limits these methods in practice.

Both the Post 4 network and DeepONet approximate a solution operator. That is not what distinguishes them, so something else has to.

It is not that one works on continuous data and the other does not. Every network that runs on a computer reads and writes finite arrays of numbers. Nothing about a neural operator stores a continuous function or evaluates an integral exactly. When the general definition of a neural operator layer is written with an integral over the domain, that integral is turned into a finite sum before anything executes. Working with functions describes the object being approximated, never the bytes going through the model.

The real difference is what the learned weights are attached to.

The input layer of the Post on the Neural PDE utilizing the Euler rollout has one slot per grid point per field. Slot 47 is wired to grid point $x_{47}$ and to nothing else, and every weight leaving that slot was trained on values read at that one location. The output layer is built the same way and emits one number per grid point.

Sample the same physical field at four times the resolution and the input vector no longer fits. The extra values cannot be padded or dropped, because they are not additional information about the old locations, they are information about new locations the network has no slots for. Retraining from scratch is the only option. The same limitation cuts the other way on the output side: there is no way to ask this network for the displacement halfway between two grid points, because a coordinate that is not one of the $N$ trained locations is not something the input format can express, where $N$ denotes the number of grid points used in training.

A neural operator attaches its weights to a fixed number of objects whose count does not depend on the grid. In DeepONet those objects are learned basis functions. In the Fourier Neural Operator they are frequency modes. Either way, changing the discretisation or resolution does not change how many weights there are or what each one means, so the trained parameters stay valid.

The simplest objects of this kind are whole shapes spread over the domain. For a string of length 1 m fixed at both ends, as in Post 3, a natural choice is the sine shapes ${\phi }_k(x)=\sqrt{2}\sin (k\pi x)$, where $k=1,2,3,\dots$ denotes the number of half-waves of the shape and the factor $\sqrt{2}$ normalises it, so that the area under ${\phi }_k^2$ is 1. One weight belongs to each shape, and the number it multiplies says how much of that shape the input contains:

Where $c_k$ denotes the amount of shape $k$ in the input, $N$ denotes the number of grid points, $x_i$ denotes grid point $i$, $a(x_i)$ denotes the input's value there, ${\phi }_k(x_i)$ denotes the shape's value there, and $\mathrm{\Delta }x$ denotes the grid spacing. For the wave equation, displacement and velocity are each measured this way and get their own coefficients.

Each term of the sum is a strip of width $\mathrm{\Delta }x$ under the product of input and shape, so the sum approximates the area under that product. A finer grid cuts the same area into more, thinner strips: $c_k$ approaches the same value at every resolution, and the weight that multiplies it keeps its meaning. The output is rebuilt from the same shapes, which can be evaluated at any point:

Where $\hat{\mathcal{G}}(a)(\xi )$ denotes the prediction at the query coordinate $\xi$, $w_k$ denotes the weight attached to shape $k$, and $M$ denotes the number of shapes. The grid size $N$ appears nowhere in this formula: the same $M$ weights apply whether the input arrived on 8 grid points or on 64.

The property this buys is called discretisation invariance: the same trained weights keep working as the discretisation changes, and the predictions get better rather than breaking as it gets finer. That is the actual definition, and the payoff is flexibility. Train once, evaluate at whatever resolution the application needs.

One thing to keep in mind before moving on: discretisation invariance is not a single switch that is either on or off. It can hold on the input side, on the output side, or on both, and most architectures do not have both. DeepONet, which the rest of this post covers, has it on the output side only, which resulted in many additional papers that tried to fix the input side.

Answering one question about a solution takes two pieces of information, and they are not alike.

The first is which initial condition the question is about. On a 100 by 100 grid with two fields that is $2\times \mathrm{10,000}=\mathrm{20,000}$ numbers, where the factor 2 denotes the displacement and velocity fields and 10,000 denotes the number of grid points, flattened into one vector under a fixed layout.

The second is where the answer is wanted: a coordinate $\xi =(x,y,t)$.

DeepONet gives each of them its own sub-network. The branch network takes the 20,000 numbers describing the initial condition. The trunk network takes the coordinate. Both are ordinary feedforward networks, meaning stacks of linear layers with a nonlinear activation in between.

Both sub-networks end in the same number of output neurons, and the prediction is their dot product:

where $\mathcal{G}(a)(\xi )$ denotes the true solution for initial condition $a$ at coordinate $\xi$, $b_k(a)$ denotes the $k$-th output of the branch network for that initial condition, $t_k(\xi )$ denotes the $k$-th output of the trunk network for that coordinate, and $p$ denotes the shared output width of the two sub-networks, usually in the low hundreds.

So let's unpack what this formula means. First of all, this is a linear combination of the trunk outputs, with the branch outputs as coefficients. The trunk $t_k$ produces a set of shapes, and the branch $b_k$ says how to mix them together for this particular input. Zooming in on the two halves. The trunk network takes in only a coordinate $\xi$ and forwards it through a stack of fully connected linear layers with $p$ output dimensions. Each output dimension $t_k$ represents one neuron each, is some function of the coordiante - some shape. If we think back about the concept of the Von Neumann Stability Analysis, we conducted a Fourier decomposition, as every continuos & arbitrarily complex signal can be reconstructed from a linear combination of sines and cosines. The trunk network is doing something similar, but instead of using a fixed basis of sines and cosines, it is learning a basis of shapes that are useful for the problem at hand. We call this the learned basis of shapes.

Now that we have the learned basis, we need to know how to combine them to get the final output. This is where the branch network comes in. The branch network takes in the initial condition $a$ and forwards it through another stack of fully connected linear layers with $p$ output dimensions. Each output dimension $b_k$ represents one neuron each, is some function of the initial condition - some coefficient. These coefficients tell us how much of each shape from the trunk to mix together for this particular input.

Assuming we $p$ equals 3 for simplicity, one spelled out example of the dot product would be:

Putting the initial condition and the coordinate into one 20,003-input network would also give something that accepts any query point, so the split needs a reason. There are two.

The first is cost. The branch runs once per initial condition, and its $p$ outputs are reused for every query about that same condition. Asking for a thousand points costs one branch pass, a thousand cheap trunk passes and a thousand dot products. A combined network would push all 20,000 initial-condition values through the model again for every single query, and that is by far the larger input.

The second is that the two inputs do different jobs. One says which function we are talking about, the other says where in the output we want to look. Splitting them puts that distinction into the architecture instead of hoping gradient descent works it out.

One training example is a triplet: an initial condition, a query coordinate, and the true value there.

where $L$ denotes the loss, $\theta$ denotes the parameters of branch and trunk together, $Q$ denotes the number of triplets, $q$ indexes them, $a^{(q)}$ denotes the initial condition of triplet $q$, ${\xi }^{(q)}$ denotes its query coordinate, ${\hat{\mathcal{G}}}_{\theta }$ denotes the network's prediction, and $\mathcal{G}$ denotes the true solution operator, whose values come from a classical solver. With two output fields the squared errors of both are summed. Branch and trunk are trained together by backpropagating through the dot product.

What makes this affordable is that one solver run yields many triplets. A single run produces the whole solution field over the whole domain and every saved time, and sampling that field at a few hundred scattered coordinates gives a few hundred training examples that share one expensive run and one branch input. Ten thousand training points can come from a hundred runs.

The trunk takes a coordinate as an actual number, so any point in the domain can be asked for, including points between grid points and points at a resolution never used in training. The output side is genuinely free of the grid.

The branch is not. It is an ordinary feedforward network with one input slot per sensor location, those locations are fixed when the dataset is built, and they have to be identical for every input function the model ever sees. Feeding it an initial condition sampled on a 200 by 200 grid fails for exactly the same reason the Post 4 network fails, and in exactly the same way.

So measured against the definition given earlier, DeepONet satisfies discretisation invariance on one side and not the other. This is a known limitation rather than a technicality, and it is the specific gap that later variants were built to close.

The trunk is made of linear layers and smooth activations, so nearby coordinates give nearby outputs. Across most of a solution field that is what you want, and it is why the trunk interpolates sensibly between training points.

At a wavefront it is the wrong behaviour. Just ahead of the front the displacement is near zero and just behind it is large, so two coordinates a fraction of a cell apart carry very different true values. The basis has to be steep exactly there and flat nearly everywhere else. Networks are bad at this, and the reason for that is spectral bias, which is the tendency of neural networks to fit low-frequency structure first and to struggle to produce sharp local variation at all.

The Fourier Neural Operator attaches its weights to a fixed number of frequency modes instead of to grid positions or to a learned basis. That gives discretisation invariance on both sides rather than one, and by the result above it handles sharp fronts far better.