Neural Operators¶

Many problems in science and engineering require solving complex PDEs that govern the behaviour of these systems. These PDEs are often difficult and inefficient to solve, requiring days and months to solve using traditional PDE solvers using numerical discretization schemes discussed previously, such as finite element method (FEM), finite difference method (FDM), finite volume method (FVM). Moreover, identifying the underlying equations governing the behaviour of these systems requires expert knowledge in the field. However, even with this knowledge, some systems may be too complicated to model requiring modelling assumptions and simplifications to best approximate the system. Data-driven and physics-guided methods, such as machine learning (ML), have been shown to solve many of these problems faster and with similar accuracy compared to these traditional solvers. These ML models have also been used to fill the missing gaps in knowledge from data.

In this notebook, we introduce a new type of ML models, neural operators (NOs), which are neural networks (NNs) designed to learn mappings between infinite-dimensional function spaces. They differ from conventional NNs which only operate on finite-dimensional vector spaces. The advantages of NOs is the ability to work directly with the functions describing our problem rather than fixed and finite vector representations. Let's first introduce some key concepts before highlight some differences between of NOs and NNs.

Key Concepts: Functions, Function Spaces and Operators¶

A function, $f$, can be interpreted as a relationship between a variable, the independent variable, from one domain to another variable, the dependent variable, to another domain. That is for any variable $x \in X$, there exists a unique mapping, $f: X \rightarrow Y$, to another variable $y \in Y$. The heat equation introduced previously is one such example, describing the relationship between the independent variables, $(x, y, t)$ to the dependent variable, $u$. We observed that for every location in the 2D domain, $(x, y)$, and time step, $t$, their existed a mapping to a unique temperature $u$, i.e. $u = f(x, y, t)$.

A function space consists of a set of functions sharing common properties and subjected to certain mathematical rules. For example, we can define multiple 'heat functions' on the same 2D domain, $\Omega$, and $t>0$, but with different thermal diffusivity, $\alpha$. Each of these heat functions, $f_{\alpha}(x, y, t) \rightarrow u$, make up a function space for our heat equation.

Operators, $G$, can be thought of as a mapping between input functions in one function space, $a \in A$, to output functions in the same or different function space, $u \in U$, i.e. $G: A \rightarrow U$. These input and output functions can represent different physical quantities, such as thermal diffusivity and temperature distributions.

Background: Neural Networks and the Universal Approximation Theorem¶

Many real-world problems in science and engineering can be solved by first identifying the governing PDE describing its behavior, and then solving these sets of complex, nonlinear equations. However, the ability to identify and compose this set of PDEs to adequately model the problem requires expert knowledge in the specific field. Additionally, some systems may be too complicated to design an appropriate predictive model or the governing set of equations may be too complex to efficiently solve by traditional means requiring simplifications. This has led to the emergence of scientific machine learning (SciML), which combines the physical laws governing a system (such as PDEs) with data-driven ML techniques.

But why can ML be used to approximate these nonlinear functions? A key property of NNs is the Universal Approximation Theorem, which states that a feedforward NN with a single hidden layer containing a finite number of neurons can approximate any continuous function, given an appropriate activation function. In general, each neuron in the network performs some operation on its input, then applies an activation function and produces some output. The complexity of the functions that the network can model depends on the architecture, the number of hidden layers, and the number of neurons in each layer. In mathematical terms, we can represent a NN for learning our heat function as, $\hat{f}(x, y, t; \theta) \rightarrow u$, taking inputs, $(x, y, t)$, produces an output, $u$, with trainable weights, $\theta$. Then for some arbitrary level of accuracy, $\epsilon$, their exists some NN, $\hat{f}$ with parameters, $\theta$, such that: \begin{equation} \tag{1} | \hat{f}(x, y, t; \theta) - f(x, y, t) |_{\infty} < \epsilon \end{equation}

This property makes NNs incredibly powerful tools for modeling nonlinear systems in science and engineering. However, it should be stated that it doesn’t guarantee that the NN will be able to learn the specific function, nor does it indicate the best choice of architecture for learning functions. As a result, training NNs to approximate complex functions can still be challenging, often requiring careful consideration of the architecture, loss function, and hyperparameters.

Motivation: Neural Operators vs Traditional Neural Networks¶

A major caveat for using traditional NNs for learning continuous functions is the fixed resolution of the training data, such as images with consistent pixel dimensions. For example, to learn our heat equation, $u = f(x, y, t)$, we would first need to define a mesh for our spatial coordinates $(x, y)$ which must be consistent across all time steps, $t$. Each set of inputs, $(x, y, t)$, must then be mapped to a temperature, $u$, also defined on some consistent mesh. Changing the dimensions of the grid, and removing, adding, or moving grid points makes these models unusable requiring costly re-training. These standard networks, such as Multilayer perceptron (MLP), convolution neural network (CNN), residual network (Resnet), are only able to train on input and output vectors of fixed dimensions at consistent grid locations.

NOs are designed to work directly with continuous domains by learning the operators mapping two functions. Training still requires discretization of the input and output functions (as shown in Figure 1), however, they learn the overall shape and behavior of the data. This allows NOs to handle functions at different resolutions without the need for costly re-training. This property, commonly referred to as discretization invariance, makes NOs more robust and applicable to different scenarios not seen during training, providing consistent performance at different resolutions and grids.

Another popular architecture in SciML is the Physics-Informed Neural Networks (PINNs), which is discretization invariant, as it only requires point-wise evaluations, i.e. $(x, y, t) \rightarrow u$. But how does it compare to NOs? We aim to learn mappings between function spaces, i.e. $f_{\alpha}(x, y, t) \rightarrow u$, where our PDE, $f$ is parametrized by some parameter, $\alpha$. PINNs are only able to learn the relationship for one version of this PDE, i.e. a fixed $\alpha$. We could add this as an additional input to our PINN, $(x, y, t, \alpha) \rightarrow u$, however this adds complexity to our learning process, as an additional set of weights must now be optimized during training.

Operator Learning¶

NOs typically learn mappings between infinite dimensional spaces from a set of observables $\{a^{(i)}, u^{(i)}\}^{N}_{i}$, consisting of pairs of input and output functions. As shown in Figure 1, each input-output pair is discretized at $m$ finite locations for performing point-wise evaluations of the solution operator, $G$. NOs approximate $G$ by optimizing the model parameters, $\theta$, to minimize the difference between the predicted and true output functions, $G_{\theta} \approx G$, (refer to equation 2). \begin{equation} \tag{2} | G(a)(y) - G_{\theta}(a)(y) | < \epsilon \end{equation} where $\epsilon$ is an acceptable user-defined tolerance. Note that $G(a)(y) = u(y)$. By iteratively training on these pairs and measuring the closeness of the predictions to the true solutions, these networks eventually learn the patterns and relationships between them. Notice the similarity between equation 1 and 2. The Universal Approximation Theorem can also be applied for approximating nonlinear operators!

Figure 1: NO training data adapted from LuLu2020

Deep Operator Networks¶

In this course, we will focus on the NO architecture called Deep Operator Networks (DeepONets), first introduced by Lu Lu et al in 2021 [LuLu2020]. The DeepONet architecture consists of two NNs, a branch network that learns an embedding of input function space, $A$, and a trunk network that learns an embedding of output function space, $U$, as shown in Figure 2. To train these networks, the branch network takes each $\{a^{(i)}\}^{m}_{i}$ at $m$ finite locations and learns a representation, $\{b_k\}^{p}_{k}$, while the trunk network takes an evaluation location for each solution, $y$ and learns a representation, $\{t_k\}^{p}_{k}$. An approximation for $G$ is then obtained using equation 3: \begin{equation} \tag{3} G_{\theta}(a)(y) \approx \sum^{p}_{k} b_k (u(x_1), u(x_2), ..., u(x_m)) \cdot t_k(y) \end{equation}

These networks are flexible in the choice of branch and trunk network architectures, and can be chosen depending on the problem being solved.

Figure 2: DeepONet architecture adapted from LuLu2020

References¶

1. Neural Operator: Learning Maps Between Function Spaces With Applications to PDEs
2. Neural Operators for Accelerating Scientific Simulations and Design
3. Learning nonlinear operators via DeepONet based on the universal approximation theorem of operators
4. Universal Approximation Theorem