Year: 2021

Abstract:
Machine learning methods have recently shown promise in solving partial differential equations (PDEs). They can be classified into two broad categories: solution function approximation and operator learning. The Physics-Informed Neural Network (PINN) is an example of the former while the Fourier neural operator (FNO) is an example of the latter. Both these approaches have shortcomings. The optimization in PINN is challenging and prone to failure, especially on multi-scale dynamic systems. FNO does not suffer from this optimization issue since it carries out supervised learning on a given dataset, but obtaining such data may be too expensive or infeasible. In this work, we propose the physics-informed neural operator (PINO), where we combine the operating-learning and function-optimization frameworks, and this improves convergence rates and accuracy over both PINN and FNO models. In the operator-learning phase, PINO learns the solution operator over multiple instances of the parametric PDE family. In the test-time optimization phase, PINO optimizes the pre-trained operator ansatz for the querying instance of the PDE. Experiments show PINO outperforms previous ML methods on many popular PDE families while retaining the extraordinary speed-up of FNO compared to solvers. In particular, PINO accurately solves long temporal transient flows and Kolmogorov flows, while PINN and other methods fail to converge.

heading: 1 INTRODUCTION
text: Machine learning-based methods are starting to show promise in scientific computing and especially in solving partial differential equations (PDEs). They have demonstrated advantages in both efficiency and accuracy compared to conventional solvers. They are even able to tackle previously intractable problems such as higher-dimensional, multi-scale, high-contrast, and chaotic PDE systems (Um et al., 2020; Brunton et al., 2020; Fan et al., 2018; Long et al., 2018; Han et al., 2018; Bruno et al., 2021). Broadly, ML-based approaches for PDEs can be divided into two categories: optimizing to solve for a specific solution function of PDE vs. learning the solution operator over a family of PDEs.
Optimization of solution function and PINN. Most ML-based methods, as well as the conventional solvers, fall into this category. Conventional solvers such as FDM and FEM usually discretize the domain into a grid and optimize/approximate the solution function on the grid, which imposes a truncation error. The Physics-Informed Neural Network (PINN)-type methods are proposed to overcome the discretization issue (Raissi et al., 2019). They use a neural network as the ansatz of the solution function and take advantage of auto-differentiation to compute the exact, mesh-free derivatives. Recently, researchers have developed numerous variations of PINN with promising results on inverse problems and partially observed tasks (Lu et al., 2021a; Zhu et al., 2019; Smith et al., 2021). However, compared to conventional solvers, PINNs face several optimization issues: (1) the challenging optimization landscape from soft physics or PDE constraints (Wang et al., 2021a), (2) the difficulty to propagate information from the initial or boundary conditions to unseen parts of the interior or to future times (Dwivedi & Srinivasan, 2020), and (3) the sensitivity to hyper-parameters selection (Sun et al., 2020). As a result, PINNs are still unable to compete with conventional solvers in most cases, and they often fail to converge on high-frequency or multi-scale PDEs (Wang et al., 2020b; Fuks & Tchelepi, 2020; Raissi et al., 2020). In this work, we propose to overcome these optimization challenges by integrating operator learning with PINN.
(a) learn the solution operator from a family of equations. (b) use the learned operator as ansatz to
solve for a specific instance.
Operator learning and neural operators. A recent alternative approach is to learn the solution operator of a family of PDEs, defined by the map from the input–initial conditions and boundary conditions, to the output–solution functions. In this case, usually, a dataset of input-output pairs from an existing solver is given. There are two main aspects to consider (a) model: to design models for learning highly complicated PDE solution operators, and (b) data: to be data-efficient and to improve generalization. Recent advances in operator learning replace traditional convolutional neural networks and U-Nets from computer vision with operator-based model tailored to PDEs with greatly improved model expressiveness (Li et al., 2020c; Lu et al., 2019; Patel et al., 2021; Wang et al., 2020a; Duvall et al., 2021). Specifically, the neural operator generalizes the neural network to the operator setting where the input and output spaces are infinite-dimensional. The framework has shown success in learning resolution-invariant solution operators for highly non-linear problems such as turbulence flow (Li et al., 2020b;a). However, the data challenges remain: (1) the need for training data, which assumes an existing solver or experimental setup, (2) the non-negligible generalization error, and (3) extrapolation to unseen conditions. These issues can be addressed by adding physics or PDE constraints to operator learning (Zhu et al., 2019; Wang et al., 2021b; Zhang et al., 2021).
Our contributions. To overcome the shortcomings of both physics-informed optimization and data-driven operator learning, we propose the physics-informed neural operator (PINO) that combines operator learning with equation solving (test-time optimization). It requires fewer or no data points to learn the operator and generalizes better. In PINO, we use the pre-trained operator as the ansatz to optimize for the solution function at test time, which reduces the generalization error. Compared to PINN, PINO has a much better optimization landscape and representation space, and hence, PINO converges faster and more accurately. Our contributions can be summarized as follows:
• We propose the physics-informed neural operator (PINO), combining the operator-learning and physics-informed settings. We introduce the pre-training and test-time optimization schemes that utilize both the data and equation constraints (whichever are available). We develop an efficient method to compute the exact gradient for neural operators to incorporate the equation constraints.
• By utilizing pre-trained operator ansatz, PINO overcomes the challenge of propagating information from the initial condition to future time steps with (soft) physics constraints. It can solve the 2d transient flow over an extremely long time period, where PINN and DeepONet (Lu et al., 2019) fail. Even without any pre-training and using only PDE constraints for the given instance, PINO still outperforms PINN by 20x smaller error and 25x speedup on the chaotic Kolmogorov flow, demonstrating superior expressivity of the neural operator over standard neural networks.
• By utilizing the equation constraints, PINO requires fewer or no training data and generalizes better compared to FNO (Li et al., 2020c). On average it has 7% smaller error on the transient and Kolmogorov flows, while matching the speedup of FNO (400x) compared to the GPU-based pseudo-spectral solver (He & Sun, 2007), matching FNO. Further, the pre-trained PINO model on the Navier Stokes equation can be easily transferred to different Reynolds numbers ranging from 100 to 500 using test-time optimization.
• We propose the forward and backward PINO models for inverse problems. Our approach accurately recovers the coefficient function in the Darcy flow which is 3000x faster than the conventional solvers using accelerated MCMC (Cotter et al., 2013).
Our major novelty and contributions are to use the pre-trained operator ansatz with instance-wise fine-tuning to overcome the optimization challenges in PINN and the generalization challenges
in operator learning. Previous works such as PINN-DeepONet (Wang et al., 2021b) and Physicsconstrained modeling (Zhu et al., 2019) use the PDE constraints in operator learning, like we do during the pre-training phase in PINO. However, we propose several methodological advances as well as extensive experiments to understand the optimization and generalization challenges. Our methodological advances include: (1) Instance-wise fine-tuning at test-time to further improve the fidelity of the operator ansatz. (2) Efficient Fourier-space methods for computing derivatives present in the PDE loss. (3) Efficient learning through the design of data augmentation and loss functions. (4) Novel formulation for inverse problems that results in accurate recovery as well as good speedups.
PINN vs. PINO: pointwise vs. function-wise optimization. The neural operator ansatz in PINO has an easier optimization landscape and a more expressive representation space compared to the neural networks ansatz in PINN. The neural operator parameterizes the solution function as an aggregation of basis functions, and hence, the optimization is in the function space. This is easier than just optimizing a single function as in PINN. Further, we can learn these basis functions in the pre-training phase which makes the test-time optimization on the querying instance even easier. In this case, PINO does not need to solve from scratch. It just fine-tunes the solution function parameterized by the solution operator. Thus, PINO is much faster and more accurate compared to PINN.

heading: 2 PRELIMINARIES AND PROBLEM SETTINGS

heading: 2.1 PROBLEM SETTINGS
text: We consider two natural class of PDEs. In the first, we consider the stationary system
P(u, a) = 0, in D ⊂ Rd
u = g, in ∂D (1)
where D is a bounded domain, a ∈ A ⊆ V is a PDE coefficient/parameter, u ∈ U is the unknown, and P : U ×A → F is a possibly non-linear partial differential operator with (U ,V,F) a triplet of Banach spaces. Usually the function g is a fixed boundary condition (potentially can be entered as a parameter). This formulation gives rise to the solution operator G† : A → U defined by a 7→ u. A prototypical example is the second-order elliptic equation P(u, a) = −∇ · (a∇u) + f . In the second setting, we consider the dynamical system
du dt = R(u), in D × (0,∞)
u = g, in ∂D × (0,∞) u = a in D̄ × {0}
(2)
where a = u(0) ∈ A ⊆ V is the initial condition, u(t) ∈ U for t > 0 is the unknown, and R is a possibly non-linear partial differential operator with U , and V Banach spaces. As before, we take g to be a known boundary condition. We assume that u exists and is bounded for all time and for every u0 ∈ U . This formulation gives rise to the solution operator G† : A → C ( (0, T ];U ) defined by a 7→ u. Prototypical examples include the Burgers’ equation and the Navier-Stokes equation.

heading: 2.2 SOLVING EQUATION USING THE PHYSICS-INFORMED LOSS (PINN)
text: Given an instance a and a solution operator G† defined by equations (1) or (2) , we denote by u† = G†(a) the unique ground truth. The equation solving task is to approximate u†. This setting consists of the ML-enhanced conventional solvers such as learned finite element, finite difference, and multigrid solvers (Kochkov et al., 2021; Pathak et al., 2021; Greenfeld et al., 2019), as well as purely neural network-based solvers such as the Physics-Informed Neural Networks (PINNs), Deep Galerkin Method, and Deep Ritz Method (Raissi et al., 2019; Sirignano & Spiliopoulos, 2018; Weinan & Yu, 2018). Especially, these PINN-type methods use a neural network uθ with parameters θ as the the ansatz to approximate the solution function u† . The parameters θ are found by minimizing the physics-informed loss with exact derivatives computed using automatic-differentiation (autograd). In the stationary case, the physics-informed loss is defined by minimizing the l.h.s. of equation (1) in
the squared norm of F . A typical choice is F = L2(D), giving the loss function
Lpde(a, uθ) = ‖P(a, uθ)‖2L2(D) + α‖uθ|∂D − g‖ 2 L2(∂D)
= ∫ D |P(uθ(x), a(x))|2dx+ α ∫ ∂D |uθ(x)− g(x)|2dx (3)
In the case of a dynamical system, it minimizes the residual of equation (2) in some natural norm up to a fixed final time T > 0. A typical choice is the L2 ( (0, T ];L2(D) ) norm, yielding
Lpde(a, uθ) = ∫ T 0 ∫ D |duθ dt (t, x)−R(uθ)(t, x)|2dxdt+ α ∫ T 0 ∫ ∂D |uθ(t, x)− g(t, x)|2dxdt
+ β ∫ D |uθ(0, x)− a(x)|2dx
(4)
The PDE loss consists of the physics loss in the interior and the data loss on the boundary and initial conditions, with hyper-parameters α, β > 0. It can be generalized to variational form as in (Weinan & Yu, 2018).
Challenges of PINN PINNs take advantage of the universal approximability of neural networks, but, in return, suffer from the low-frequency induced bias. Empirically, PINNs often fail to solve challenging PDEs when the solution exhibits high-frequency or multi-scale structure (Wang et al., 2021a; 2020b; Fuks & Tchelepi, 2020; Raissi et al., 2020). Further, as an iterative solver, PINNs have difficulty propagating information from the initial condition or boundary condition to unseen parts of the interior or to future times (Dwivedi & Srinivasan, 2020). For example, in challenging problems such as turbulence, PINNs are only able to solve the PDE on a relatively small domain (Jin et al., 2021), or otherwise, require extra observational data which is not always available in practice (Raissi et al., 2020; Cai et al., 2021). In this work, we propose to overcome the challenges posed by the optimization by integrating operator learning with PINNs.

heading: 2.3 LEARNING THE SOLUTION OPERATOR (NEURAL OPERATOR)
text: An alternative setting is to learn the solution operator G. Given a PDE as defined in (1) or (2) and the corresponding solution operator G†, one can use a neural operator Gθ with parameters θ as a surrogate model to approximate G†. Usually we assume a dataset {aj , uj}Nj=1 is available, where G†(aj) = uj and aj ∼ µ are i.i.d. samples from some distribution µ supported on A. In this case, one can optimize the solution operator by minimizing the empirical data loss on a given data pair
Ldata(u,Gθ(a)) = ‖u− Gθ(a)‖2U = ∫ D |u(x)− Gθ(a)(x)|2dx (5)
where we assume the setting of (1) for simplicity of the exposition. The operator data loss is defined as the average error across all possible inputs
Jdata(Gθ) = ‖G† − Gθ‖2L2µ(A;U) = Ea∼µ[Ldata(a, θ)] ≈ 1
N N∑ j=1 ∫ D |uj(x)− Gθ(aj)(x)|2dx. (6)
Similarly, one can define the operator PDE loss as
Jpde(Gθ) = Ea∼µ[Lpde(a,Gθ(a))]. (7)
In general, it is non-trivial to compute the derivatives dGθ(a)/dx and dGθ(a)/dt for model Gθ. In the following section, we will discuss how to compute these derivatives for Fourier neural operator.

heading: 2.4 NEURAL OPERATORS
text: In this work, we will focus on the neural operator model designed for the operator learning problem. The neural operator, proposed in (Li et al., 2020c), is formulated as an generalization of standard deep neural networks to operator setting. Neural operator compose linear integral operator K with pointwise non-linear activation function σ to approximate highly non-linear operators.
Definition 1 (Neural operator Gθ) Define the neural operator
Gθ := Q ◦ (WL +KL) ◦ · · · ◦ σ(W1 +K1) ◦ P (8)
where P : Rda → Rd1 , Q : RdL → Rdu are the pointwise neural networks that encode the lower dimension function into higher dimensional space and vice versa. The model stack L layers of σ(Wl +Kl) where Wl ∈ Rdl+1×dl are pointwise linear operators (matrices), Kl : {D → Rdl} → {D → Rdl+1} are integral kernel operators, and σ are fixed activation functions. The parameters θ consists of all the parameters in P,Q,Wl,Kl.
Recently, Li et al. (2020a) proposes the Fourier neural operator (FNO) that deploys convolution operator for K. In this case, it can apply the Fast Fourier Transform (FFT) to efficiently compute K. This leads to a fast architecture that obtains state-of-the-art results for PDE problems.
Definition 2 (Fourier convolution operator K) Define the Fourier convolution operator( Kvt ) (x) = F−1 ( R · (Fvt) ) (x) ∀x ∈ D (9)
where R is part of the parameter θ to be learn.
Challenges of Operator learning. Operator learning is similar to the supervised learning in computer vision and language where data play a very important role. One needs to assume the training points and testing points follow the same problem setting and the same distribution. Especially, the previous FNO model trained on one coefficient (e.g. Reynolds number) or one geometry cannot be easily generalized to another. Moreover, for more challenging PDEs where the solver is very slow or the solver is even not existent, it is hard to gather a representative dataset. On the other hand, since FNO doesn’t use any knowledge of the equation, it is cannot get arbitrary close to the ground truth by using higher resolution as in conventional solvers, leaving a gap of generalization error. These challenges limit FNO’s applications beyond accelerating the solver. In the following section, we will introduce the PINO framework to overcome these problems by using the equation constraints.

heading: 3 PHYSICS-INFORMED NEURAL OPERATOR (PINO)
text: We propose the PINO framework that uses one neural operator model Gθ for solving both operator learning problems and equation solving problems. It consists of two phases
• Pre-train the solution operator: learn a neural operator Gθ to approximate G† using either/both the data loss Jdata and/or the PDE loss Jpde which can be additional to the dataset.
• Test-time optimization: use Gθ(a) as the ansatz to approximate u† with the pde loss Lpde and/or an additional operator loss Lop obtained from the pre-train phase.

heading: 3.1 PRE-TRAIN: OPERATOR LEARNING WITH PDE LOSS
text: Given a fixed amount of data {aj , uj}, the data loss Jdata offers a stronger constraint compared to the PDE loss Jpde. However, the PDE loss does not require a pre-existing dataset. One can sample virtual PDE instances by drawing additional initial conditions or coefficient conditions aj ∼ µ for training, as shown in Algorithm 1. In this sense, we have access to the unlimited dataset by sampling new aj in each iteration.

heading: 3.2 TEST-TIME OPTIMIZATION: SOLVING EQUATION WITH OPERATOR ANSATZ
text: Given a learned operator Gθ, we use Gθ(a) as the ansatz to solve for u†. The optimization procedure is similar to PINNs where it computes the PDE loss Lpde on a, except that we propose to use a neural operator instead of a neural network. Since the PDE loss is a soft constraints and challenging to optimize, we also add an optional operator loss Lop (anchor loss) to bound the further optimized model from the pre-trained model
Lop ( Gθi(a),Gθ0(a) ) := ‖Gθi(a)− Gθ0(a)‖2U
Algorithm 1: Pre-training scheme for Physics-informed neural operator learning Data: Data: input output function pairs {aj , uj}Nj=1 Result: Neural operator G : A → U .
1 for i = 0, 1, 2 . . . do 2 Compute Ldata and Lpde on (ai, ui). Update neural operator G; 3 for j = 1 to K do 4 Sample a′ from distribution µ; 5 Compute Lpde on a′. Update neural operator G; 6 end 7 end
where Gθi(a) is the model at ith training epoch. We update the operator Gθ using the lossLpde+αLop. It is possible to further apply optimization techniques to fine-tune the last fewer layers of the neural operator and progressive training that gradually increase the grid resolution and use finer resolution in test time.
Optimization landscape. Using the operator as the ansatz has two major advantages: (1) PINN does point-wise optimization, while PINO does optimization in the space of functions. In the linear integral operation K, the operator parameterizes the solution function as a sum of the basis function. Optimization of the set of coefficients and basis is easier than just optimizing a single function as in PINNs. (2) we can learn these basis functions in the operator learning phase which makes the later test-time optimization even easier. In PINO, we do not need to propagate the information from the initial condition and boundary condition to the interior. It just requires fine-tuning the solution function parameterized by the solution operator.
Trade-off (1) complexity and accuracy: the test-time optimization is an option to spend more computation in exchange of better accuracy. When the speed is most desired, one should use the pre-trained model to directly output the prediction. However, if one wants to solve for a specific instance accurately, then the person can use test-time optimization to correct the generalization error posed in operator learning. (2) resolution effects on optimization landscape and truncation error: using higher resolution and finer grid will reduce the truncation error. However, it may make the optimization instable. Using hard constraints such as the anchor loss Lop relieve such problem.

heading: 3.3 DERIVATIVE OF NEURAL OPERATORS
text: In order to use the equation loss Lpde, one of the major technical challenge is to efficiently compute the derivatives D(Gθa) = d(Gθa)(x)/dx for neural operators. In this section, we discuss three efficient methods to compute the derivatives of neural operator Gθ as defined in 8.
Numerical differentiation. A simple but efficient approach is to use conventional numerical gradients such as finite difference and Fourier gradient (Zhu et al., 2019; Gao et al., 2021). These numerical gradient are fast and memory-efficient: given a n-points grid, finite difference requiresO(n) and Fourier method require O(n log n). However, they face the same challenges as the corresponding numerical solvers: finite difference methods require a fine-resolution uniform grid; spectral methods require smoothness and uniform grids. Especially. These numerical errors on the gradient will be amplified on the output solution.
Autograd. Similar to PINN (Raissi et al., 2019), the most general way to compute the exact gradient is to use the auto-differentiation library of neural networks (autograd). To apply autograd, one needs to use a neural network to parameterize the solution function û : x 7→ u(x). However, it is not straightforward to write out the solution function in the neural operator which directly outputs the numerical solution u = Gθ(a) on a grid, especially for FNO which uses FFT. To apply autograd, we design a query function û that input x and output u(x). Recall Gθ := Q◦ (WL +KL) ◦ · · · ◦ σ(W1 + K1) ◦ P and u = Gθa = QvL = Q(WL +KL)vL−1 . . .. Since Q is pointwise,
u(x) = Q(vL(x)) = Q((WLvL−1)(x) +KLvL−1(x)) (10)
(a) the long-temporal transient flow with Re ∼ 20, T = 50. PINO can output the full trajectory in one step, which leads to 400x speedup compared to the solver. PINN cannot converge to a reasonable error rate due to the long time window. (b) the chaotic Kolmogorov flow with Re =∼ 500, T = 0.5. PINO retains 10x speed up compared to the solver, while converges much faster compared to PINN.
Figure 2: The accuracy-complexity trade-off on PINO, PINN, and the pseudo-spectral solver.
To query any location x, we need to evaluate the Fourier convolution KLvL−1(x) and the bias part (WLvL−1)(x). For the first part, we can directly query the Fourier coefficient R · Fvt at spatial location x without doing the inverse FFT F . For the second part, we can define the query function as an interpolation or a low-rank integral operator as in (Kovachki et al., 2021). Then we can use autograd to take the exact derivative for (10). The autograd method is general and exact, however, it is less efficient. Since the number of parameters |θ| is usually much greater than the grid size n, the numerical methods are indeed significantly faster. Empirically, the autograd method is usually slower and memory-consuming.
Exact gradients. We develop an efficient and exact gradient method based on the architecture of the neural operator. The idea is similar to the autograd, but we explicitly write out the gradient on the Fourier space and apply the chain rule. Given the explicit form (10), D(KLvL−1) can be directly compute on the Fourier space. The linear part D(WLvL−1) can be interpolated using the Fourier method. Therefore to exactly compute the derivative of vL, one just needs to run the numerical Fourier gradient. Then we just need to apply chain rule for Jacobian Ju = J(Q ◦ vL) = J(Q(vL)) ◦ J(vL) and Hessian Hu = H(Q ◦ vL) = J(vL) ◦H(Q(vL)) ◦ J(vL) + J(Q(vL)) ◦H(vL). Since Q is pointwise, this can be much simplified. In the experiments, we mostly use the exact gradient and the numerical Fourier gradient.

heading: 4 EXPERIMENTS
text: In this section, we conduct empirical experiments to examine the efficacy of the proposed PINO. In 4.1, we show the physics constraints helps operator-learning with fewer data and better generalization. Then in 4.2, we investigate how PINO uses operator ansatz to solve harder equations with improved speed and accuracy. We study on three concrete cases of PDEs on Burgers’ Equation, Darcy Equation, and Navier-Stokes equation.
Burgers’ Equation. The 1-d Burgers’ equation is a non-linear PDE with periodic boundary conditions where u0 ∈ L2per((0, 1);R) is the initial condition and ν = 0.01 is the viscosity coefficient. We aim to learn the operator mapping the initial condition to the solution, G† : u0 7→ u|[0,1].
∂tu(x, t) + ∂x(u 2(x, t)/2) = ν∂xxu(x, t), x ∈ (0, 1), t ∈ (0, 1] u(x, 0) = u0(x), x ∈ (0, 1)
(11)
Darcy Flow. The 2-d steady-state Darcy Flow equation on the unit box which is the second order linear elliptic PDE with a Dirichlet boundary where a ∈ L∞((0, 1)2;R+) is a piecewise constant diffusion coefficient and f = 1 is a fixed forcing function. We are interested in learning the operator
mapping the diffusion coefficient to the solution, G† : a 7→ u. Note that although the PDE is linear, the operator G† is not.
−∇ · (a(x)∇u(x)) = f(x) x ∈ (0, 1)2
u(x) = 0 x ∈ ∂(0, 1)2 (12)
Navier-Stokes Equation. We consider the 2-d Navier-Stokes equation for a viscous, incompressible fluid in vorticity form on the unit torus, where u ∈ C([0, T ];Hrper((0, l)2;R2)) for any r > 0 is the velocity field, w = ∇× u is the vorticity, w0 ∈ L2per((0, l)2;R) is the initial vorticity, ν ∈ R+ is the viscosity coefficient, and f ∈ L2per((0, l)2;R) is the forcing function. We want to learn the operator mapping the vorticity from the initial condition to the full solution G† : w0 7→ w|[0,T ].
∂tw(x, t) + u(x, t) · ∇w(x, t) = ν∆w(x, t) + f(x), x ∈ (0, l)2, t ∈ (0, T ] ∇ · u(x, t) = 0, x ∈ (0, l)2, t ∈ [0, T ] w(x, 0) = w0(x), x ∈ (0, l)2
(13)
Specially, we consider two problem settings:
• Long temporal transient flow: we study the build-up of the flow from the initial condition u0 near-zero velocity to uT that reaches the ergodic state. We choose t ∈ [0, 50], l = 1, Re = 20 as in Li et al. (2020a). The main challenge is to predict the long time interval.
• Chaotic Kolmogorov flow: In this case u lies in the attractor where arbitrary starting time t0. We choose t ∈ [t0, t0 + 0.5] or [t0, t0 + 1], l = 1, Re = 500 similar to Li et al. (2021). The main challenge is to capture the small details that evolve chaotically.

heading: 4.1 OPERATOR LEARNING WITH PDE LOSSES
text: We first show the effectiveness of applying the physics constraints Jpde to learn the solution operator.
Burgers equation and Darcy equation. PINO can learn the solution operator without any data on simpler problems such as Burgers and Darcy. Compared to other PDE-constrained operators, PINO is more expressive and thereby achieves better accuracy. On Burgers (11), PINN-DeepONet achieves 1.38% (Wang et al., 2021b); PINO achieves 0.37%. Similarly, on Darcy flow (12), PINO outperforms FNO by utilizing physics constraints, as shown in Table 1. For these simpler equation, test-time optimization may not be needed. The implementation detail and and the search space of parameters are included in the Appendix A.1 and A.2.

heading: 4.2 SOLVE EQUATION USING OPERATOR ANSATZ
text: Long temporal transient flow. It is extremely challenging to propagate the information from the initial condition to future time steps over such a long interval T = [0, 50] just using the soft physics constraint. None of the PINN, DeepONet, and PINO (from scratch without pre-training) can handle this case (error > 50%), no matter solving the full interval at once or solving per smaller steps. However, when the pre-training data is available for PINO, we can use the learned neural operator ansatz and the anchor loss Lop. The anchor loss is a hard constraint that makes the optimization much easier. Providing N = 4800 training data, the PINO without test-time optimization achieves 2.87% error, lower than FNO 3.04% and it retains a 400x speedup compared to the GPU-based pseudo-spectral solver (He & Sun, 2007), matching FNO. Further doing test time optimization with the anchor loss and PDE loss, PINO reduces the error to 1.84%.
Chaotic Kolmogorov flow. We show an empirical study on how PINO can improve the generalization of neural operators by enforcing more physics. For this experiment, the training set consists of 4000 data points of initial condition and corresponding solution. Using Algorithm 1, we can sample unlimited additional initial conditions from Gaussian random field at any resolution. Table 2 and 4 compare the generalization error of neural operators trained by different schemes and different amounts of simulated data. The result shows that training neural operator with additional PDE instances consistently improves the generalization error on all three resolutions we are evaluating.
Based on the solution operators learned in the above operator-learning section, we continue to do test-time optimization. The results are shown in Figure and Table 3. Overall, PINO outperforms PINN by 20x smaller error and 25x speedup. Using pre-trained model make PINO converge faster. The implementation detail and the search space of parameters are included in the Appendix B.2.
Transfer Reynolds numbers. The extrapolation of different parameters and conditions is one of the biggest challenges for ML-based methods. By doing test-time optimization, the pre-trained PINO model on the Kolmogorov flow can be easily transferred to different Reynolds numbers ranging from 100 to 500 in test-time optimization as shown in Table 5. Such property envisions broad applications.

heading: 4.3 INVERSE PROBLEM
text: One of the major advantages of the physics-informed method is to solve the inverse problem. In this case, we investigate PINO on the inverse problem of the Darcy equation to recover the coefficient function a† from the given solution function u†. We assume a dataset {aj , uj} is available to pre-train the operator. We propose two formulations of PINO for the inverse problem:
• Forward model: Learn the forward operator Gθ : a 7→ u with data. Initialize â to approximate a†. Optimize â using Jforward := Lpde(â, u†) + Ldata(Gθ(â)) +R(â). • Backward model: Learn the backward operatorFθ : u 7→ awith data. UseFθ(u†) to approximate a†. Optimize Fθ using
Jbackward := Lpde(Fθ(u†), u†) + Lop(Fθ(u†),Fθ0(u†)) +R(Fθ(u†))
Where R(a) is the regularization term. We use the PDE loss Lpde to deal with the small error in Gθ and the ill-defining issue of Fθ. We perform a showcase example on the Darcy equation, where we are given the full-field observation of u (clean data) to recover the coefficient function a. The coefficient function a is piecewise constant (representing two types of media), so the inverse problem can be viewed as a classification problem. We define R(a) as the total variance. As shown in Figure 3, the PINO backward model has the best performance: it has 2.29% relative l2 error on the output u and 97.10% classification accuracy on the input a; the forward model has 6.43% error on the output and 95.38% accuracy on the input. As a reference, we compare the PINO inverse frameworks with PINN and the conventional solvers using the accelerated MCMC method with 500,000 steps (Cotter et al., 2013). The posterior mean of the MCMC has 4.52% error and 90.30% respectively ( Notice the Bayesian method outputs the posterior distribution, which is beyond obtaining a maximum a posteriori estimation) Meanwhile, PINO methods are 3000x faster compared to MCMC. PINN does not converge in this case. The major advantage of the PINO backward model compared to the PINO forward model is that it uses a neural operator Fθ(u†) as the ansatz for the coefficient function. Similar to the forward problem, the operator ansatz has an easier optimization landscape while being expressive.

heading: B ADDITIONAL EXPERIMENTS
text: B.1 ADDITIONAL BASELINES
We add a comparison experiment against the Locally adaptive activation functions for PINN (LAAFPINN) (Jagtap et al., 2020) and Self-Adaptive PINN (SA-PINN) (McClenny & Braga-Neto, 2020). For the Kolmogorov flow problem, we set Re=500, T=[0, 0.5]. We search among the following hyperparameters combinations: LAAF-PINN: n: 10, 100, learning rate: 0.1, 0.01, 0.001, depth 4, 6. SA-PINNs: learning rate 0.001, 0.005, 0.01, 0.05, network width 50, 100, 200, depth 4, 6, 8.
As shown in Figure 6, both LAAF-PINN and SA-PINN converge much faster compared to the original PINN method, but there is still a big gap with PINO. LAAF-PINN adds learnable parameters before the activation function; SA-PINN adds the weight parameter for each collocation point. These techniques help to alleviate the PINNs‘ optimization problem significantly. However, they didn’t alter the optimization landscape effectively in the authors’ opinion. On the other hand, by using operator ansatz, PINO optimizes in a function-wise manner where the optimization is fundamentally different.
Note that the contribution of PINO is orthogonal to the above methods. One can apply the adaptive activation functions or self-adaptive loss in the PINO framework too. All these techniques of PINNs can be straightforwardly transferred to PINO. We believe it would be interesting future directions to study how all these methods work with each other in different problems.
B.2 LID CAVITY FLOW.
We demonstrate an addition example using PINO to solve for lid-cavity flow on T = [5, 10] with Re = 500. In this case, we do not have the pre-train phase and directly solve the equation (test-time optimization). We use PINO with the velocity-pressure formulation and resolution 65 × 65 × 50 plus the Fourier numerical gradient. It takes 2 minutes to achieve a relative error of 14.52%. Figure 4 shows the ground truth and prediction of the velocity field at t = 10 where the PINO accurately predicts the ground truth.
We assume a no-slip boundary where u(x, t) = (0, 0) at left, bottom, and right walls and u(x, t) = (1, 0) on top, similar to Bruneau & Saad (2006). We choose t ∈ [5, 10], l = 1, Re = 500. We use the velocity-pressure formulation as in Jin et al. (2021) where the neural operator output the velocity field in x, y, and the pressure field. We set width = 32, mode = 20 with learning rate 0.0005 which decreases by half every 5000 iterations, 5000 iterations in total. We use the Fourier method with Fourier continuation to compute the numerical gradient and minimize the residual error on the velocity, the divergence-free condition, as well as the initial condition and boundary condition. The weight parameters (α, β) between different error terms are all chosen as 1. Figure 4 shows the ground truth and prediction of the velocity field at t = 10 where the PINO accurately predicts the ground truth.

heading: C FOURIER CONTINUATION
text: The Fourier neural operator can be applied to arbitrary geometry via Fourier continuations. Given any compact manifoldM, we can always embed it into a periodic cube (torus),
i :M→ T n
where we can do the regular FFT. Conventionally, people would define the embedding i as a continuous extension by fitting polynomials (Bruno et al., 2007). However, in Fourier neural operator, it can be simply done by padding zeros in the input. The loss is computed at the original space during training. The Fourier neural operator will automatically generate a smooth extension to do padded domain in the output, as shown in Figure 7.
This technique is first used in the original Fourier neural operator paper (Li et al., 2020a) to deal with the time dimension in the Navier-Stokes equation. Similarly, Lu et al. (2021b) apply FNO with extension and interpolation on diverse geometries on the Darcy equation. In the work, we use Fourier continuation widely for non-periodic boundary conditions (Darcy, time dimension). We also added an example of lid-cavity to demonstrate that PINO can work with non-periodic boundary conditions.
Furthermore, this Fourier continuation technique helps to take the derivatives of the Fourier neural operator. Since the output of FNO is always on a periodic domain, the numerical Fourier gradient is usually efficient and accurate, except if there is shock (in this case, we will use the exact gradient method).

heading: D CONCLUSION AND FUTURE WORKS
text: In this work, we develop the physics-informed neural operator (PINO) that bridges the gap between physics-informed optimization and data-driven neural operator learning. We introduce the pre-training and test-time optimization schemes for PINO to utilize both the data and physics constraints. In the pre-training phase, PINO learns an operator ansatz over multiple instances of a parametric PDE family. The test-time optimization scheme allows us to take advantage of the learned neural operator ansatz and solve for the solution function on the querying instance faster and more accurately.
There are many exciting future directions. PINO can be used to enhance operator-learning with some constraints, for examples: (1) Control problem: one first learns a solution operator (pre-training) and does test-time optimization in the control phase to make sure it satisfies desired constraints (Hwang et. al.). (2) Multiscale modeling: in multi-scale modeling, one learns a fine-scale solution operator, which is then regulated by a coarse-scale solver. Note that the fine-scale problems do not follow Gaussian distribution as in the training dataset (Liu et. al.). (3) Weather prediction: beyond the observation dataset, one can add physics constraints such as the conservation of mass (Jiang et. al.). (4) Airfoil design: the design problem can be formulated as an inverse problem as shown in Section 4.3. PINO backward formulation is particularly effective for the inverse problem (Thuerey et. al.).
PINO is also designed to work on most of the working scenarios of PINNs: if the optimization is benign, then pre-train is not needed. The major advantage compared to PINN-based methods is the optimization landscape. The operator ansatz, even randomly initialized, still has a much better
optimization landscape compared to the PINN-based model. Most of the techniques and analysis of PINN can be transferred to PINO.
It is interesting to ask how to overcome the hard trade-off of accuracy and complexity, and how the PINO model transfers across different geometries. Furthermore, we can develop a software library of pre-trained models. PINO’s excellent extrapolation property allows it to be applied on a broad set of conditions, as shown in the Transfer Reynold’s number experiments.

Review summaries:
The paper proposes the physics-informed neural operator (PINO). It combines the operating-learning and function-optimization frameworks, which improves convergence rates and accuracy over traditional  methods. Experiments test the advantage of PINO.

The authors introduce physics-informed neural operator for learning partial differential equations (PDEs) from data with deep learning. This approach combines two recent and popular methods in the field of machine learning and PDEs, the physics-informed neural network (PINNs) for solving PDEs and Fourier neural operator (FNO) for learning solution operators. The paper takes advantages of both approaches to overcome convergence issues arising in PINNs by incorporating PDE constraints into FNO. Then, the authors apply their method to a number of challenging test cases and obtain better results than PINNs on PDE solving problems and FNO for learning solution operators.

This paper proposes Physics-Informed Neural Operator (PINO) by combining two previous methods, Fourier neural operators (FNO) and physics-informed neural networks (PINNs). Similar to one use case of FNO, PINO learns the solution operator over multiple instances of a parametric PDE family in the first phase. In the second phase, given an instance from the parametric PDE family, PINO fine-tunes the learned solution operator as a pre-trained model to work well for the particular instance. The proposed method is evaluated by using numerical experiments.