Source(s): Basically just https://arxiv.org/pdf/2108.08481 and https://zongyi-li.github.io/blog/2020/fourier-pde/

Introduction

The world around us is basically governed entirely by partial differential equations. The famous Navier-Stokes equation governs fluid flow, Laplace's equation shows up in practically every imaginable subfield of physics (electromagnetism, thermodynamics), Schrodinger's equation describes the most fundamental quantum mechanical interaction between particles. Solving PDE's is of great importance to society too -- numerical simulations of Navier-Stokes are used everywhere in mechanical engineering to gauge how efficient different constructions of vehicles are, and quantum simulations of Schrodinger's equation allow us to basically simulate reality itself (although really crudely), with applications in molecular engineering, drug design, etc.

The problem, though, is that PDEs are really hard to solve. The goal of solving a PDE is to find some solution "function" that satisfies your differential equation everywhere (so, your solution itself is a function over your space).

Take a simple differential equation as example:

Here, we just have a 1-dimensional function $f(x)$ we want to solve for, so our "space" is just the real number line. We're told that the derivative of our "goal function" $f(x)$ is equal to the value of that function everywhere. Our solution is a function that satisfies this equation for all $x$, which happens to be $f(x) = e^x$ (ignoring constants and all that shit).

Very few PDEs have completely solvable analytical answers, so researchers have to resort to approximate numerical techniques (most of which basically amount to chopping up your space into a bunch of little boxes and making some assumption about your function in each of those boxes -- that it's constant, linear, etc. This is what's known as the finite element method in engineering). The problem with this approach is not necessarily accuracy (with small enough boxes, it can get pretty accurate), but rather how computationally intensive it is. More precise solutions require more fine-grained chopping, which especially scales with the dimension of your space (this is a big reason why we need huge supercomputers to do any sort of meaningful chemistry simulations for drug design/related work).

The key idea behind the neural operators paper is to do away with this idea of finite element approximation completely. Instead, we see differential equations as an "operator" that takes in some input function (our solution), and outputs some target function. This sounds really abstract, but lets see an example:

Consider the equation

Scientifically inclined readers may recognize this as the inhomogeneous wave equation, but for our purposes it's just a ton of symbols -- the point of putting it there is to develop the abstraction necessary to understand neural operators. The point is, this is a partial differential equation, where we want to "solve" for the function $u(x, t)$ in the same manner that we did above. Our $f(x, t)$ (sometimes called the source function) is known and depends on the situation -- it could be $f(x, t) = 0, f(x, t) = e^x$, whatever, it doesn't really matter. We want to think of that messy left hand side (the $(\nabla^2 - \frac{1}{c^2}\frac{\partial^2}{\partial t^2})$) as "taking in a function" $u(x, t)$, doing some transformation on that function, and spitting out the $f(x, t)$. This is what we call an operator -- it takes in some object of some "type" (in this case, a function), transforms it in some way, and spits out a different object of the same type ($f(x, t)$ is different, but still a function of two variables). In this specific case we call it the "differential operator," which is denoted by the symbol $L_a$ (to be consistent with the paper, the $L$ means linear). Any partial differential equation, thus, can be expressed as the following:

A couple clarifications:

• Why the $a$ subscript for $L$? We want to somehow parameterize different differential operators, such as the wave equation above that performs the transformation $(\nabla^2 - \frac{1}{c^2}\frac{\partial^2}{\partial t^2})$ or, as a different example, the "elliptic operator" studied in the original paper that performs the transformation $-\nabla \cdot (a \nabla)$. $a \in \mathcal{A}$ represents our parameter space. In practice, this parameter space is typically some sort of function space, meaning our parameter $a$ itself is a function. Think of the above example as our differential operator having the overall structure $-\nabla \cdot (a \cdot \nabla)$ (whatever that means), but the specific choice of $a$ function that we plug into this operator will actually concretely define what transformation the operator is doing.
• The $D$ is our "domain," or the underlying space our functions operate over. In the above example $df/dx = f(x)$, our domain is just $D = \mathbb{R}$, the real numbers. But it could also be complex numbers, some higher dimensional $\mathbb{R}^n$ or some other weird space that isn't so easily described (suppose your functions take in points on a sphere or some other surface).
• What is the $u(x) = 0$ line? If you've taken a differential equations course, you may remember to define a concrete solution function, we need "initial conditions" which tell us what the solution should be equal to (for some intuition regarding this, in our previous example our solution could be any function of the form $f(x) = A e^x$ for some number $A$, but giving the initial condition $f(0) = 1$ forces that $A = 0$). Then, $\delta D$ represents the "boundary" of our domain $D$. Don't really worry about this, just think of it as some subset of points in which we define our initial conditions (so in the example $f(0) = 1$, our subset is just a single point $x = 0$).

With this, our goal is to use deep learning to directly model this operator. Specifically, we want a map that will take in our parameter $a \in \mathcal{A}$ for the differential operator (which itself is a function), and spit out some solution function $u$ which lies in our "function space" $\mathcal{U}$. How to actually do this is where the beauty of this paper lies, but unfortunately that beauty lies behind a boatload of functional analysis/theory. For those who are interested, I'll try to dissect it a little in the next few sections.

Neural Operator

Our goal with neural operators is to learn the mapping

where $\theta \in \mathbb{R}^p$ just represents the parameters of our model (neural network, smth else). We then fit this model to the true map $\mathcal{G}^\dagger \triangleq L_a^{-1}f : \mathcal{A} \xrightarrow{} \mathcal{U}$.

Immediately, there's the issue that these "function spaces" $\mathcal{A}$ and $\mathcal{U}$ are infinite-dimensional: how could we ever represent them as neural network input/outputs? To remedy this, each function is sampled at certain points from our domains $D$ and $D'$. In other words, we select some list of points $x_1, x_2, \cdots, x_L \in D$ at which we "sample" our function, which gives us a list of values $a(x_1), a(x_2), \cdots, a(x_L) \in \mathbb{R}$. Likewise, we do the same for $u$, sampling some points from $D'$.

This becomes our "dataset" for neural operator learning, where we are given $N$ datapoints. Each datapoint can be thought of as a pair of these lists, representing input parametr values $a^{(i)}$ and the corresponding "desired" solution values $u^{(i)}$. So our input and output to the neural operator model end up being vectors of numbers, but don't forget that what these vectors actually represent are discretizations of continuous-valued functions!

If our functions are represented as discrete vectors anyways, what's the whole point of this operator modelling approach? We can actually prove (using something called discrete refinements) that as we refine our selection of points over our domain $D$ the "distance" between our output solution when fed in the discretization versus the true output solution converges to zero. What this means is that adding more points will gradually have negligible effect, so we don't need to worry about overarching approximation errors coming from this step.

Now we have some input and some output that we can actually train our model with! The next step is figuring out what the model should actually be. The neural operators paper defines the following structure for $\mathcal{G}_\theta : \mathcal{A} \xrightarrow{}\mathcal{U}$

1. Lifting: Our input is some vector in $\mathbb{R}^{d_a}$ that represents our input function $a$ "sampled" at some point. This step just lifts this vector to a higher dimension, call it $\mathbb{R}^{d_{v_0}}$; intuitively, our function $a$ is a lot more complicated than the small amount of points we sample it from, so our internal representation of this function should be a lot higher dimensional than what actually comes into our model.
2. Iterative Kernel Integration: This is a fancy way to say the "hidden layers" of our network. Here, we map one hidden representation of our function which lies in some $\mathbb{R}^{d_{v_t}}$ (think of $d_{v_t}$ as the $t$th hidden layer's "dimension") to some next hidden layer's dimension $\mathbb{R}^{d_{v_{t+1}}}$. What this layer actually does is apply a linear operator (think: weight matrix), an integral kernel operator (no real analog to normal NNs), and a bias function (think: bias scalar), applying a nonlinearity at the end. We see the clear parallels with neural networks that act on numbers here.
3. Projection: The last hidden layer will output the representation of our function $u$, which we need to cast down to whatever output dimension $\mathbb{R}^{d_u}$ we defined for our model (remember this output dimension represents the size of the vector -- "how many points we sample the function $u$ from").

For those that enjoy syntax (ew), take the lifting operator to be $\mathcal{P}$, each iterative "layer" to be $\sigma_i(W_i+ \mathcal{K}_i + b_i)$, where $\sigma_i$ is the nonlinearity, $W_i$ is the linear operator, $\mathcal{K}_i$ is the kernel operator, and $b_i$ is the bias function. Lastly, take the projection operator to be $\mathcal{Q}$. Then, our total map is

For anyone in CS189, this should look pretty familar! Each layer is roughly like a $h(x) = \sigma(Wx + b)$ hidden layer, which we then stack on top of each other. The key difference here is that while all of these operations are technically done on vectors, but what these vectors actually represent are functions, and each layer is really a transformation between functions! So really, our model here is acting on functions themselves, which is a really cool idea.

Integral Kernel Operators

At this point, there's still a big question mark around what an "integral kernel operator" is, which seems to be the core of this model. Lets dive deeper into that:

In general, integral operators are the way we transform one function into another. Taking from Wikipedia, integral operators (denoted by $T$) are of the form

Here, $K(t, u)$ is what we call the "kernel," which is just some function that governs what the transform actually does (so, for example, the kernel $K(t, u) = \delta(t - u)$ does the operation of "picking out" the value of the function at $u$, which ends up actually just returning the same function. Fourier transforms are actually a type of integral transform with a different kernel!). Think of $T$ then as another map from a function to another function (so $Tf$ is the output function which we then apply on some input $u$, whose value is given by the integral on the RHS) (the best analogy I can think of here is like currying in 61A).

The integral kernel operators in neural operators are basically just this, except with a few extra details. For sake of completeness, the integral kernel operator for some layer $t$ with input function $v_t(x)$ that maps from some domain $D_t$ to $\mathbb{R}^{d_{v_t}}$ is given by

Overall, this integral operator should "feel" similar to the one I introduced above. To explain the syntax:

• $\kappa_t(x, y)$ is just our kernel function, which we define to be a continuous function of two variables $x$ and $y$ (similar to our $t$ and $u$ above)
• $v_t(y)$ in the integral directly maps to our $f(t)$ above, where we integrate over the product of $v_t$ and the kernel $\kappa^{(t)}$.
• Think of $d\nu_t(y)$ as some "generalization" of the integral differential element $dt$ (specifically, it's a measure on $\mathbb{R}^{d_t}$). This is just a mathematical technicality because our hidden layer function domains $D_t$ may not be as nice or straightforward as $\mathbb{R}^n$, so we need some more advanced techniques to be able to take integrals over them. Understanding it as a simple $dy$ is perfectly fine.
• Then, we integrate over the entire input domain which is $D_t$.

The paper also defines different kernel operator parameterization options, such as one that takes in the $a$ parameters at each layer (think of these like skip connections to keep the input "relevant" as we get deeper in the models), and even one that puts the $v_t$ function itself as an input to the kernel operator. Choosing a good kernel function, then, becomes the big design question of neural operators, among other hyperparameter considerations. Interestingly enough, this architecture can be shown to be a generalization of hyperparameters!

For clarity, now that we know what each $\mathcal{K}_t(v_t)$ looks like, lets explicitly write the update rule for going from $v_t$ to $v_{t+1}$:

That big $\Pi_t(x)$ can just be thought of as some extra transformations on our input $x$. So, our entire model can be thought of as these integral transforms + linear transforms "stacked" on top of each other.

Where do our parameters of the model $\theta$ go? They implicitly are in the $W_t$ and $b_t$ terms, the $\kappa^{(t)}$ terms, as well as in defining the $\mathcal{P}, \mathcal{Q}$ lifting and projection operators too.

Computation

Ok, now we have our model, but we still have a big issue -- above, we've defined some really complicated integrals, but it could very well be that these integrals can't be analytically solved. The naive solution is to just discretize every integral. That is, take a discretization of $D$ as $\{x_1, \cdots, x_j\}$ and so

becomes

As the kernel is a matrix operator $\kappa(x_j, x_l)$, so as it stands this becomes intractable especially as we scale up our refinement by increasing $J$. A good amount of the paper is devoted to finding ways to efficiently and accurately calculate this integral, including truncation (if we're evaluating a function at $x$, don't integrate over all of $D$, but rather a neighborhood of $x$ that lies in $D$), graph neural networks (explanation: cooked), low rank approximations, etc. The one that we will focus on is the Fourier Neural Operator

Fourier Neural Operator

There's a fundamental theorem in Fourier analysis that convolution (basically equivalent to the operation our integral transform does) of two functions can be transformed into a pointwise multiplication of these functions if we take their Fourier transform (if you want a better understanding of the Fourier transform, there are millions of sources online that'll explain it better than me). Plus, with FFT (Fast fourier transform), it's faster to compute convolutions by taking the Fourier transform, doing pointwise multiplication, and inverse Fourier transforming rather than performing convolutions directly. This is the essence of what the Fourier neural operator does -- every integral transform layer described above is replaced with a fourier transform, multiplication by the "fourier modes" (aka fourier transform of) the kernel operator, and inverse transform. Note that this procedure calculates only the integral term in each hidden layer. To get the full $t+1$th output, we also perform the $W$ linear transformation, add the bias function, and pass them through our nonlinearity. For sake of constraining computational complexity and not introducing high frequency modes, we also "truncate" the Fourier transform at some value (think of Fourier space as the set of frequencies that define our function, then each component is how strong that specific frequency is -- on average, we mostly care about the lower frequency terms)

That's basically it, refer to the blog post for more details on what experiments they ran, comparisons with previous models, and other stuff.

Fourier Neural Operator Implementation

When I first read through the neural operator paper, I was really interseted in seeing how the model is actually implemented in code, given that its structure is quite different from a standard feed-forward neural network, especially the FNO implementation.

I highly recommend you check out https://github.com/neuraloperator/neuraloperator/blob/main/neuralop/models/fno.py. In particular, the feed forward method:

def forward(self, x, output_shape=None, **kwargs):
        """FNO's forward pass
        
        1. Applies optional positional encoding

        2. Sends inputs through a lifting layer to a high-dimensional latent
            space

        3. Applies optional domain padding to high-dimensional intermediate function representation

        4. Applies `n_layers` Fourier/FNO layers in sequence (SpectralConvolution + skip connections, nonlinearity) 

        5. If domain padding was applied, domain padding is removed

        6. Projection of intermediate function representation to the output channels

        Parameters
        ----------
        x : tensor
            input tensor
        
        output_shape : {tuple, tuple list, None}, default is None
            Gives the option of specifying the exact output shape for odd shaped inputs.
            
            * If None, don't specify an output shape

            * If tuple, specifies the output-shape of the **last** FNO Block

            * If tuple list, specifies the exact output-shape of each FNO Block
        """

        if output_shape is None:
            output_shape = [None]*self.n_layers
        elif isinstance(output_shape, tuple):
            output_shape = [None]*(self.n_layers - 1) + [output_shape]

        # append spatial pos embedding if set
        if self.positional_embedding is not None:
            x = self.positional_embedding(x)
        
        x = self.lifting(x)

        if self.domain_padding is not None:
            x = self.domain_padding.pad(x)

        for layer_idx in range(self.n_layers):
            x = self.fno_blocks(x, layer_idx, output_shape=output_shape[layer_idx])

        if self.domain_padding is not None:
            x = self.domain_padding.unpad(x)

        x = self.projection(x)

        return x

Ok this roughly correlates with what we defined above symbolically. What is each FNO_block? It is a series of "spectral convolutions," which seems to be exactly this convolution with the kernel operator "weights" implementing by doing FFT, doing some sorta computation akin to element-wise multiplication (todo: need to read the code better to understand what it's doing), and an iFFT to return to spatial dimension.

Appendix

Will dive into the math here

What is a function space?

I've been throwing around the term "function space" a lot in these notes, but it might be a little unclear what that actually means.