Neural operators

Neural operators are a class of deep learning architectures designed to learn maps between infinite-dimensional function spaces. Neural operators represent an extension of traditional artificial neural networks, marking a departure from the typical focus on learning mappings between finite-dimensional Euclidean spaces or finite sets. Neural operators directly learn operators between function spaces; they can receive input functions, and the output function can be evaluated at any discretization. ^[1]

The primary application of neural operators is in learning surrogate maps for the solution operators of partial differential equations (PDEs), ^[1] which are critical tools in modeling the natural environment. ^{[2] [3]} Standard PDE solvers can be time-consuming and computationally intensive, especially for complex systems. Neural operators have demonstrated improved performance in solving PDEs ^[4] compared to existing machine learning methodologies while being significantly faster than numerical solvers. ^{[5] [6] [7]} Neural operators have also been applied to various scientific and engineering disciplines such as turbulent flow modeling, computational mechanics, graph-structured data, ^[8] and the geosciences. ^[9] In particular, they have been applied to learning stress-strain fields in materials, classifying complex data like spatial transcriptomics, predicting multiphase flow in porous media, ^[10] and carbon dioxide migration simulations. Finally, the operator learning paradigm allows learning maps between function spaces, and is different from parallel ideas of learning maps from finite-dimensional spaces to function spaces, ^{[11] [12]} and subsumes these settings when limited to fixed input resolution.

Operator learning

Understanding and mapping relationships between function spaces has many applications in engineering and the sciences. In particular, one can cast the problem of solving partial differential equations as identifying a map between function spaces, such as from an initial condition to a time-evolved state. In other PDEs this map takes an input coefficient function and outputs a solution function. Operator learning is a machine learning paradigm to learn solution operators mapping the input function to the output function.

Using traditional machine learning methods, addressing this problem would involve discretizing the infinite-dimensional input and output function spaces into finite-dimensional grids and applying standard learning models, such as neural networks. This approach reduces the operator learning to finite-dimensional function learning and has some limitations, such as generalizing to discretizations beyond the grid used in training.

The primary properties of neural operators that differentiate them from traditional neural networks is discretization invariance and discretization convergence. ^[1] Unlike conventional neural networks, which are fixed on the discretization of training data, neural operators can adapt to various discretizations without re-training. This property improves the robustness and applicability of neural operators in different scenarios, providing consistent performance across different resolutions and grids.

Definition and formulation

Architecturally, neural operators are similar to feed-forward neural networks in the sense that they are composed of alternating linear maps and non-linearities. Since neural operators act on and output functions, neural operators have been instead formulated as a sequence of alternating linear integral operators on function spaces and point-wise non-linearities. ^[1] Using an analogous architecture to finite-dimensional neural networks, similar universal approximation theorems have been proven for neural operators. In particular, it has been shown that neural operators can approximate any continuous operator on a compact set. ^[1]

Neural operators seek to approximate some operator ${\displaystyle {\mathcal {G}}:{\mathcal {A}}\to {\mathcal {U}}}$ between function spaces ${\displaystyle {\mathcal {A}}}$ and ${\displaystyle {\mathcal {U}}}$ by building a parametric map ${\displaystyle {\mathcal {G}}_{\phi }:{\mathcal {A}}\to {\mathcal {U}}}$. Such parametric maps ${\displaystyle {\mathcal {G}}_{\phi }}$ can generally be defined in the form

${\displaystyle {\mathcal {G}}_{\phi }:={\mathcal {P}}\circ \sigma (W_{T}+{\mathcal {K}}_{T}+b_{T})\circ \cdots \circ \sigma (W_{1}+{\mathcal {K}}_{1}+b_{1})\circ {\mathcal {Q}},}$

where ${\displaystyle {\mathcal {P}},{\mathcal {Q}}}$ are the lifting (lifting the codomain of the input function to a higher dimensional space) and projection (projecting the codomain of the intermediate function to the output codimension) operators, respectively. These operators act pointwise on functions and are typically parametrized as multilayer perceptrons. ${\displaystyle \sigma }$ is a pointwise nonlinearity, such as a rectified linear unit (ReLU), or a Gaussian error linear unit (GeLU). Each layer ${\displaystyle t=1,\dots ,T}$ has a respective local operator ${\displaystyle W_{t}}$ (usually parameterized by a pointwise neural network), a kernel integral operator ${\displaystyle {\mathcal {K}}_{t}}$, and a bias function ${\displaystyle b_{t}}$. Given some intermediate functional representation ${\displaystyle v_{t}}$ with domain ${\displaystyle D}$ in the ${\displaystyle t}$-th hidden layer, a kernel integral operator ${\displaystyle {\mathcal {K}}_{\phi }}$ is defined as

${\displaystyle ({\mathcal {K}}_{\phi }v_{t})(x):=\int _{D}\kappa _{\phi }(x,y,v_{t}(x),v_{t}(y))v_{t}(y)dy,}$

where the kernel ${\displaystyle \kappa _{\phi }}$ is a learnable implicit neural network, parametrized by ${\displaystyle \phi }$.

In practice, one is often given the input function to the neural operator at a specific resolution. For instance, consider the setting where one is given the evaluation of ${\displaystyle v_{t}}$ at ${\displaystyle n}$ points ${\displaystyle \{y_{j}\}_{j}^{n}}$. Borrowing from Nyström integral approximation methods such as Riemann sum integration and Gaussian quadrature, the above integral operation can be computed as follows:

${\displaystyle \int _{D}\kappa _{\phi }(x,y,v_{t}(x),v_{t}(y))v_{t}(y)dy\approx \sum _{j}^{n}\kappa _{\phi }(x,y_{j},v_{t}(x),v_{t}(y_{j}))v_{t}(y_{j})\Delta _{y_{j}},}$

where ${\displaystyle \Delta _{y_{j}}}$ is the sub-area volume or quadrature weight associated to the point ${\displaystyle y_{j}}$. Thus, a simplified layer can be computed as

${\displaystyle v_{t+1}(x)\approx \sigma \left(\sum _{j}^{n}\kappa _{\phi }(x,y_{j},v_{t}(x),v_{t}(y_{j}))v_{t}(y_{j})\Delta _{y_{j}}+W_{t}(v_{t}(y_{j}))+b_{t}(x)\right).}$

The above approximation, along with parametrizing ${\displaystyle \kappa _{\phi }}$ as an implicit neural network, results in the graph neural operator (GNO). ^[13]

There have been various parameterizations of neural operators for different applications. ^{[5] [13]} These typically differ in their parameterization of ${\displaystyle \kappa }$. The most popular instantiation is the Fourier neural operator (FNO). FNO takes ${\displaystyle \kappa _{\phi }(x,y,v_{t}(x),v_{t}(y)):=\kappa _{\phi }(x-y)}$ and by applying the convolution theorem, arrives at the following parameterization of the kernel integral operator:

${\displaystyle ({\mathcal {K}}_{\phi }v_{t})(x)={\mathcal {F}}^{-1}(R_{\phi }\cdot ({\mathcal {F}}v_{t}))(x),}$

where ${\displaystyle {\mathcal {F}}}$ represents the Fourier transform and ${\displaystyle R_{\phi }}$ represents the Fourier transform of some periodic function ${\displaystyle \kappa _{\phi }}$. That is, FNO parameterizes the kernel integration directly in Fourier space, using a prescribed number of Fourier modes. When the grid at which the input function is presented is uniform, the Fourier transform can be approximated using the discrete Fourier transform (DFT) with frequencies below some specified threshold. The discrete Fourier transform can be computed using a fast Fourier transform (FFT) implementation.

Training

Training neural operators is similar to the training process for a traditional neural network. Neural operators are typically trained in some Lp norm or Sobolev norm. In particular, for a dataset ${\displaystyle \{(a_{i},u_{i})\}_{i=1}^{N}}$ of size ${\displaystyle N}$, neural operators minimize (a discretization of)

${\displaystyle {\mathcal {L}}_{\mathcal {U}}(\{(a_{i},u_{i})\}_{i=1}^{N}):=\sum _{i=1}^{N}\|u_{i}-{\mathcal {G}}_{\theta }(a_{i})\|_{\mathcal {U}}^{2}}$,

where ${\displaystyle \|\cdot \|_{\mathcal {U}}}$ is a norm on the output function space ${\displaystyle {\mathcal {U}}}$. Neural operators can be trained directly using backpropagation and gradient descent-based methods.

Another training paradigm is associated with physics-informed machine learning. In particular, physics-informed neural networks (PINNs) use complete physics laws to fit neural networks to solutions of PDEs. Extensions of this paradigm to operator learning are broadly called physics-informed neural operators (PINO), ^[14] where loss functions can include full physics equations or partial physical laws. As opposed to standard PINNs, the PINO paradigm incorporates a data loss (as defined above) in addition to the physics loss ${\displaystyle {\mathcal {L}}_{PDE}(a,{\mathcal {G}}_{\theta }(a))}$. The physics loss ${\displaystyle {\mathcal {L}}_{PDE}(a,{\mathcal {G}}_{\theta }(a))}$ quantifies how much the predicted solution of ${\displaystyle {\mathcal {G}}_{\theta }(a)}$ violates the PDEs equation for the input ${\displaystyle a}$.

References

1. Kovachki, Nikola; Li, Zongyi; Liu, Burigede; Azizzadenesheli, Kamyar; Bhattacharya, Kaushik; Stuart, Andrew; Anandkumar, Anima (2021). "Neural operator: Learning maps between function spaces" (PDF). Journal of Machine Learning Research. 24: 1–97. arXiv: 2108.08481 .
2. Evans, L. C. (1998). Partial Differential Equations. Providence: American Mathematical Society. ISBN   0-8218-0772-2.
3. "How AI models are transforming weather forecasting: A showcase of data-driven systems". phys.org (Press release). European Centre for Medium-Range Weather Forecasts. 6 September 2023.
4. Russ, Dan; Abinader, Sacha (23 August 2023). "Microsoft and Accenture partner to tackle methane emissions with AI technology". Microsoft Azure Blog.
5. Li, Zongyi; Kovachki, Nikola; Azizzadenesheli, Kamyar; Liu, Burigede; Bhattacharya, Kaushik; Stuart, Andrew; Anima, Anandkumar (2020). "Fourier neural operator for parametric partial differential equations". arXiv: 2010.08895 [cs.LG].
6. Hao, Karen (30 October 2020). "AI has cracked a key mathematical puzzle for understanding our world". MIT Technology Review.
7. Ananthaswamy, Anil (19 April 2021). "Latest Neural Nets Solve World's Hardest Equations Faster Than Ever Before". Quanta Magazine.
8. Sharma, Anuj; Singh, Sukhdeep; Ratna, S. (15 August 2023). "Graph Neural Network Operators: a Review". Multimedia Tools and Applications. 83 (8): 23413–23436. doi:10.1007/s11042-023-16440-4.
9. Wen, Gege; Li, Zongyi; Azizzadenesheli, Kamyar; Anandkumar, Anima; Benson, Sally M. (May 2022). "U-FNO—An enhanced Fourier neural operator-based deep-learning model for multiphase flow". Advances in Water Resources. 163: 104180. arXiv: 2109.03697 . Bibcode:2022AdWR..16304180W. doi:10.1016/j.advwatres.2022.104180.
10. Choubineh, Abouzar; Chen, Jie; Wood, David A.; Coenen, Frans; Ma, Fei (2023). "Fourier Neural Operator for Fluid Flow in Small-Shape 2D Simulated Porous Media Dataset". Algorithms. 16 (1): 24. doi: 10.3390/a16010024 .
11. Jiang, Chiyu Lmaxr; Esmaeilzadeh, Soheil; Azizzadenesheli, Kamyar; Kashinath, Karthik; Mustafa, Mustafa; Tchelepi, Hamdi A.; Marcus, Philip; Prabhat, Mr; Anandkumar, Anima (2020). "MESHFREEFLOWNET: A Physics-Constrained Deep Continuous Space-Time Super-Resolution Framework". SC20: International Conference for High Performance Computing, Networking, Storage and Analysis. pp. 1–15. doi:10.1109/SC41405.2020.00013. ISBN   978-1-7281-9998-6.
12. Lu, Lu; Jin, Pengzhan; Pang, Guofei; Zhang, Zhongqiang; Karniadakis, George Em (18 March 2021). "Learning nonlinear operators via DeepONet based on the universal approximation theorem of operators". Nature Machine Intelligence. 3 (3): 218–229. arXiv: 1910.03193 . doi:10.1038/s42256-021-00302-5.
13. Li, Zongyi; Kovachki, Nikola; Azizzadenesheli, Kamyar; Liu, Burigede; Bhattacharya, Kaushik; Stuart, Andrew; Anima, Anandkumar (2020). "Neural operator: Graph kernel network for partial differential equations". arXiv: 2003.03485 [cs.LG].
14. Li, Zongyi; Hongkai, Zheng; Kovachki, Nikola; Jin, David; Chen, Haoxuan; Liu, Burigede; Azizzadenesheli, Kamyar; Anima, Anandkumar (2021). "Physics-Informed Neural Operator for Learning Partial Differential Equations". arXiv: 2111.03794 [cs.LG].

External links

• neuralop – Python library of various neural operator architectures