Physics-Informed Neural Operator for  Learning Partial Differential Equations

1 INTRODUCTION . 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 . 2 PRELIMINARIES AND PROBLEM SETTINGS . 2.1 PROBLEM SETTINGS . 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 . 2.2 SOLVING EQUATION USING THE PHYSICS-INFORMED LOSS ( PINN ) . 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 .

Summary: 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.