P HYSICS -I NFORMED N EURAL O PERATOR FOR L EARNING P ARTIAL D IFFERENTIAL E QUATIONS
Anonymous authors
Paper under double-blind review
A BSTRACT
Machine learning methods have recently shown promise in solving partial differential equations (PDEs).
They can be classiﬁed 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 ﬂows and
Kolmogorov ﬂows, while PINN and other methods fail to converge.
I NTRODUCTION
Machine learning-based methods are starting to show promise in scientiﬁc computing and especially in solving partial differential equations (PDEs).
They have demonstrated advantages in both efﬁciency 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 speciﬁc 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 difﬁculty 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.
[Figure]
the solution operator from a family of equations
[Figure]
[Figure]
the learned operator as ansatz to solve for a speciﬁc instance.
Operator learning and neural operators.
A recent alternative approach is to learn the solution operator of a family of PDEs, deﬁned 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-efﬁcient 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).
Speciﬁcally, the neural operator generalizes the neural network to the operator setting where the input and output spaces are inﬁnite-dimensional.
The framework has shown success in learning resolution-invariant solution operators for highly non-linear problems such as turbulence ﬂow (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 efﬁcient 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 ﬂow 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 ﬂow, 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 ﬂows, 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 coefﬁcient function in the Darcy ﬂow 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 ﬁne-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 ﬁne-tuning at test-time to further improve the ﬁdelity of the operator ansatz.
(2) Efﬁcient Fourier-space methods for computing derivatives present in the PDE loss. (3) Efﬁcient learning through the design of data augmentation and loss functions. 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 ﬁne-tunes the solution function parameterized by the solution operator.
Thus, PINO is much faster and more accurate compared to PINN.
P RELIMINARIES AND P ROBLEM SETTINGS 2. P ROBLEM SETTINGS
We consider two natural class of PDEs.
In the ﬁrst, we consider the stationary system
[Equation]
where D is a bounded domain,
a ∈ A ⊆ V is a PDE coefﬁcient/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 ﬁxed boundary condition (potentially can be entered as a parameter).
This formulation gives rise to the solution operator G † :
A → U deﬁned by a (cid:55)→ 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
[Equation]
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 u 0 ∈ U .
This formulation gives rise to the solution operator G † : A → C (cid:0) (0 , T ]; U (cid:1) deﬁned by a (cid:55)→ u .
Prototypical examples include the Burgers’ equation and the Navier-Stokes equation.
2. S OLVING EQUATION USING THE PHYSICS - INFORMED LOSS (PINN)
Given an instance a and a solution operator G † deﬁned 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 ﬁnite element, ﬁnite 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 deﬁned by minimizing the l.h.s.
of equation (1) in
[Equation]
In the case of a dynamical system, it minimizes the residual of equation (2) in some natural norm up to a ﬁxed ﬁnal time T > 0 .
A typical choice is the L 2 (cid:0) (0 , T ]; L 2 ( D ) (cid:1) norm, yielding
[Equation]
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 difﬁculty 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.
2. L EARNING THE SOLUTION OPERATOR ( NEURAL OPERATOR )
An alternative setting is to learn the solution operator G .
Given a PDE as deﬁned 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 { a j , u j } Nj =1 is available, where
G † ( a j ) =
u j and a j ∼ µ 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
[Equation]
where we assume the setting of (1) for simplicity of the exposition.
The operator data loss is deﬁned as the average error across all possible inputs
[Equation]
Similarly, one can deﬁne the operator PDE loss as
[Equation]
In general, it is non-trivial to compute the derivatives
d G θ ( a ) /dx
and d G θ ( a ) /dt for model G θ .
In the following section, we will discuss how to compute these derivatives for Fourier neural operator.
2. N EURAL OPERATORS
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.
Deﬁnition 1 (Neural operator G θ ) Deﬁne the neural operator
[Equation]
where P : R d a → R d 1 , Q : R d L → R d u are the pointwise neural networks that encode the lower dimension function into higher dimensional space and vice versa.
The model stack L layers of σ ( W l + K l ) where W l ∈ R d
l +1 ×
d l are pointwise linear operators (matrices), K l : { D → R d l } → { D → R d l +1 } are integral kernel operators, and σ are ﬁxed activation functions.
The parameters θ consists of all the parameters in P , Q , W l , K l .
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 efﬁciently compute K .
This leads to a fast architecture that obtains state-of-the-art results for PDE problems.
Deﬁnition 2 (Fourier convolution operator K ) Deﬁne the Fourier convolution operator
[Equation]
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 coefﬁcient (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.
P HYSICS - INFORMED NEURAL OPERATOR (PINO)
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 J data and/or the PDE loss J pde which can be additional to the dataset.
• Test-time optimization : use G θ ( a ) as the ansatz to approximate u † with the pde loss L pde and/or an additional operator loss L op obtained from the pre-train phase.
3. P RE - TRAIN : OPERATOR LEARNING WITH PDE LOSS
Given a ﬁxed amount of data { a j , u j } , the data loss J data offers a stronger constraint compared to the PDE loss J pde .
However, the PDE loss does not require a pre-existing dataset.
One can sample virtual PDE instances by drawing additional initial conditions or coefﬁcient conditions a j ∼ µ for training, as shown in Algorithm 1.
In this sense, we have access to the unlimited dataset by sampling new a j in each iteration.
3. T EST - TIME OPTIMIZATION : SOLVING EQUATION WITH OPERATOR ANSATZ
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 L pde 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 L op (anchor loss) to bound the further optimized model from the pre-trained model
[Equation]
u
=
Neural
A
U .
do Compute L pde Update neural G ; end 7 end where G θ i ( a ) is the model at i th training epoch.
We update the operator G θ using the loss L pde + α L op .
It is possible to further apply optimization techniques to ﬁne-tune the last fewer layers of the neural operator and progressive training that gradually increase the grid resolution and use ﬁner 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 coefﬁcients 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 ﬁne-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 speciﬁc 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 ﬁner grid will reduce the truncation error.
However, it may make the optimization instable.
Using hard constraints such as the anchor loss L op relieve such problem.
3. D ERIVATIVE OF NEURAL OPERATORS
In order to use the equation loss L pde , one of the major technical challenge is to efﬁciently compute the derivatives D ( G θ a ) =
d ( G θ a )( x ) /dx for neural operators.
In this section, we discuss three efﬁcient methods to compute the derivatives of neural operator G θ as deﬁned in 8.
Numerical differentiation.
A simple but efﬁcient approach is to use conventional numerical gradients such as ﬁnite difference and Fourier gradient (Zhu et al., 2019; Gao et al., 2021).
These numerical gradient are fast and memory-efﬁcient: given a n -points grid, ﬁnite difference requires O ( n ) and Fourier method require O ( n log n ) .
However, they face the same challenges as the corresponding numerical solvers: ﬁnite difference methods require a ﬁne-resolution uniform grid; spectral methods require smoothness and uniform grids.
Especially.
These numerical errors on the gradient will be ampliﬁed 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 (cid:55)→ 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 ◦
( W L + K L ) ◦ · · · ◦ σ ( W 1 +
K 1 ) ◦ P and u =
G θ a = Q v L = Q ( W L + K L ) v L − 1 . . .
Since Q is pointwise,
[Equation]
a) the long-temporal transient with Re ∼
T = 50 .
PINO can output the full trajectory in one step which leads to 400x solver cannot converge to a reasonable error rate due to
PINO retains 10x speed up compared to the solver, while converges much faster compared to PINN.
To query any location x , we need to evaluate the Fourier convolution K L v L − 1 ( x ) and the bias part ( W L v L − 1 )( x ) .
For the ﬁrst part, we can directly query the Fourier coefﬁcient R · F v t at spatial location x without doing the inverse FFT F .
For the second part, we can deﬁne 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 efﬁcient.
Since the number of parameters | θ | is usually much greater than the grid size n , the numerical methods are indeed signiﬁcantly faster.
Empirically, the autograd method is usually slower and memory-consuming.
Exact gradients.
We develop an efﬁcient 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 ( K L v L − 1 ) can be directly compute on the Fourier space.
The linear part D ( W L v L − 1 ) can be interpolated using the Fourier method.
Therefore to exactly compute the derivative of v L , one just needs to run the numerical Fourier gradient.
Then we just need to apply chain rule for Jacobian Ju =
J ( Q ◦ v L ) =
J ( Q ( v L )) ◦ J ( v L ) and Hessian Hu
= H ( Q ◦ v L ) =
J ( v L ) ◦ H ( Q ( v L )) ◦ J ( v L ) + J ( Q ( v L )) ◦ H ( v L ) .
Since Q is pointwise, this can be much simpliﬁed.
In the experiments, we mostly use the exact gradient and the numerical Fourier gradient.
E XPERIMENTS
In this section, we conduct empirical experiments to examine the efﬁcacy 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 u 0 ∈ L 2 per ((0 , 1); R ) is the initial condition and ν = 0 .
01 is the viscosity coefﬁcient.
We aim to learn the operator mapping the initial condition to the solution, G † : u 0 (cid:55)→ u | [0 , 1] .
[Equation]
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 coefﬁcient and f = 1 is a ﬁxed forcing function.
We are interested in learning the operator mapping the diffusion coefﬁcient to the solution, G † : a (cid:55)→ u .
Note that although the PDE is linear, the operator G † is not.
[Equation]
Navier-Stokes Equation.
We consider the 2-d Navier-Stokes equation for a viscous, incompressible ﬂuid in vorticity form on the unit torus, where u ∈ C ([0 , T ]; H r per ((0 , l ) 2 ; R 2 )) for any r > 0 is the velocity ﬁeld, w = ∇ × u is the vorticity, w 0 ∈
L 2 per ((0 , l ) 2 ; R ) is the initial vorticity, ν ∈ R + is the viscosity coefﬁcient, and
f ∈ L 2 per ((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 † : w 0 (cid:55)→ w | [0 ,T ] .
[Equation]
Specially, we consider two problem settings:
• Long temporal transient ﬂow : we study the build-up of the ﬂow from the initial condition unear-zero velocity to u T 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 ﬂow : In this case u lies in the attractor where arbitrary starting time t 0 .
We choose t ∈ [ t 0 , t 0 + 0 . 5] or [ t 0 , t 0 + 1] , l = 1 , Re = 500 similar to Li et al.
The main challenge is to capture the small details that evolve chaotically.
4. O PERATOR LEARNING WITH PDE LOSSES
We ﬁrst show the effectiveness of applying the physics constraints J pde 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 ﬂow (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.
4. S OLVE EQUATION USING OPERATOR ANSATZ
Long temporal transient ﬂow.
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 L op .
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 ﬂow.
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 ﬁeld 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 ﬂow 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.
4. I NVERSE PROBLEM
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 coefﬁcient function a † from the given solution function u † .
We assume a dataset { a j , u j } 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 (cid:55)→ u with data.
Initialize ˆ a to approximate a † .
Optimize ˆ a using
[Equation]
• Backward model : Learn the backward operator F θ : u (cid:55)→ a with data.
Use F θ ( u † ) to approximate a † .
Optimize F θ using
[Equation]
Where R ( a ) is the regularization term.
We use the PDE loss L pde to deal with the small error in G θ and the ill-deﬁning issue of F θ .
We perform a showcase example on the Darcy equation, where we are given the full-ﬁeld observation of u (clean data) to recover the coefﬁcient function a .
The coefﬁcient function a is piecewise constant (representing two types of media), so the inverse problem can be viewed as a classiﬁcation problem.
We deﬁne 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% classiﬁcation 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 coefﬁcient function.
Similar to the forward problem, the operator ansatz has an easier optimization landscape while being expressive.
A I MPLEMENTATION DETAILS
In this section, we list the detailed experiment setups and parameter searching for each experiment in Section 4.
Without speciﬁcation, we use Fourier neural operator backbone with width = 64 , mode = 8 , L = 4 and GeLU activation.
The numerical experiments are performed on Nvidia V
GPUs.
A. B URGERS E QUATION
We use the 1000 initial conditions u 0 ∼ µ where µ = N (0 , 625( − ∆ + 25 I ) − 2 ) to train the solution operator on PINO with width = 64 , mode = 20 , and GeLU activation.
We use the numerical method to take the gradient.
We use Adam optimizer with the learning rate 0 .
001 that decays by half every 100 epochs.
500 epochs in total.
The total training time is about 1250 s on a single Nvidia
GPU.
PINO achieves 0.37% relative l2 error averaged over 200 testing instances.
PINN-DeepONet achieves 1.38% which is taken from Wang et al.
(2021b) which uses the same problem setting.
A. D ARCY F LOW
We use the 1000 coefﬁcient conditions a to train the solution operator where a ∼ µ where µ = ψ # N (0 , ( − ∆ + 9 I ) − 2 ) , ψ ( a ( x ))
=
if a ( x ) ≥ 0 ; ψ ( a ( x )) =
if a ( x ) < 0 .
The zero boundary condition is enforced by multiplying a molliﬁer m ( x ) = sin( πx ) sin( πy ) for all methods.
The parameter of PINO on Darcy Flow is the same as in the Burgers equation above.
Regarding the implementation detail of the baselines: as for FNO, we use the same hyperparameters as its paper did (Li et al., 2020c); DeepONet (Lu et al., 2019) did not study Darcy ﬂow so we grid search the hyperparameters of DeepONet: depth from 2 to 12, width from 50 to 100.
The best result of
DeepONet is achieved by depth 8, width 50.
The results are shown in Table
[Table]
A. L ONG TEMPORAL TRANSIENT FLOW .
We study the build-up of the ﬂow from the initial condition u 0 near-zero velocity to u T that reaches the ergodic state.
We choose T = 50 , l = 1 as in Li et al.
(2020a).
We choose the weight parameters of error α
= β = 5 .
The initial condition w 0 ( x ) is generated according to w 0 ∼ µ where µ = N (0 , 7 3 / 2 ( − ∆ + 49 I ) − 2 . 5 ) with periodic boundary conditions.
The forcing is kept ﬁxed f ( x ) = 0 .
1(sin(2 π ( x 1 + x 2 ))
+ cos(2 π ( x 1 + x 2 ))) .
We compare FNO, PINO (no test-time optimization), and PINO (with test-time optimization).
They get 3 .
87% , and 1 .
84% relative l2 error on the last time step u (50) over 5 testing instance.
[Table]
A. C HAOTIC K OLMOGOROV FLOW .
For this experiment, u lies in the attractor.
We choose T = 0 .
5 or 1 , and l = 1 , similar to
Li et al.
The training set consists of 4000 initial condition functions and corresponding solution functions with a spatial resolution of 64 × 64 and a temporal resolution of 65 .
Extra initial conditions are generated from Gaussian random ﬁeld N (cid:0) 0 , 7 3 / 2 ( − ∆ + 49 I ) − 5 / 2 (cid:1) .
We estimate the generalization error of the operator on a test set that contains 300 instances of Kolmogorov ﬂow and reports the averaged relative L 2 error.
Each neural operator is trained with 4k data points plus a number of extra sampled initial conditions.
The Reynolds number in this problem is 500.
The reported generalization error is averaged over 300 instances.
Comparison study.
The baseline method PINN is implemented using library DeepXDE (Lu et al., 2021c) with TensorFlow as backend.
We use the two-step optimization strategy (Adam (Kingma & Ba, 2014) and L-BFGS) following the same practice as NSFNet (Jin et al., 2021), which applies
PINNs to solving Navier Stokes equations.
We grid search the hyperparameters: network depth from 4 to 6, width from 50 to 100, learning rate from 0.01 to 0.0001, the weight of boundary loss from 10 to 100 for all experiments of PINNs.
Comparison between PINO and PINNs on test-time optimization.
The results are averaged over 20 instances of the Navier Stokes equation with Reynolds number 500.
The best result is obtained by PINO using pre-trained operator ansatz and virtual sampling.
The neural operator ansatz used here is trained over a set of 400 data points.
The authors acknowledge that there could exist more sophisticated variants of PINN that performs better in our test cases.
[Table]
A. T RANSFER LEARNING ACROSS R EYNOLDS NUMBERS
We study the test-time optimization with different Reynolds numbers on the 1s Kolmogorov ﬂow.
For the higher Reynolds number problem Re = 500 , 400 , the pre-training operator shows better
[Table]
The Reynolds number is 500.
The reported generalization error is averaged over 300 instances.
Training on additional initial conditions boosts the generalization ability of the operator.
convergence accuracy.
In all cases, the pre-training operator shows better convergence speed as demonstrated in Figure 5.
The results are shown in Table 5 where the error is averaged overinstances.
Each row is a testing case
[Table]
B A DDITIONAL EXPERIMENTS B. A DDITIONAL 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 ﬂow 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.
We observe that a ll the operator ansatzs pretrained over PDE instances with different Reynolds number can boost the test-time optimization accuracy and speed compared to training from scratch.
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 signiﬁcantly.
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. L ID C AVITY FLOW .
We demonstrate an addition example using PINO to solve for lid-cavity ﬂow 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 ×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 ﬁeld 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 ﬁeld in x , y , and the pressure ﬁeld.
We set width = 32 , mode = 20 with learning rate 0 .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 ﬁeld at t = 10 where the PINO accurately predicts the ground truth.
C F OURIER CONTINUATION
The Fourier neural operator can be applied to arbitrary geometry via Fourier continuations.
Given any compact manifold M , we can always embed it into a periodic cube (torus),
[Equation]
where we can do the regular FFT.
Conventionally, people would deﬁne the embedding i as a continuous extension by ﬁtting 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 ﬁrst 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 efﬁcient and accurate, except if there is shock (in this case, we will use the exact gradient method). D C ONCLUSION AND F UTURE WORKS
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 ﬁrst learns a solution operator (pre-training) and does test-time optimization in the control phase to make sure it satisﬁes desired constraints (Hwang et.
al.).
(2) Multiscale modeling: in multi-scale modeling, one learns a ﬁne-scale solution operator, which is then regulated by a coarse-scale solver.
Note that the ﬁne-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.