SEMI-SUPERVISED LEARNING OF PARTIAL DIFFEREN-TIAL OPERATORS AND DYNAMICAL FLOWS

Abstract

The evolution of many dynamical systems is generically governed by nonlinear partial differential equations (PDEs), whose solution, in a simulation framework, requires vast amounts of computational resources. In this work, we present a novel method that combines a hyper-network solver with a Fourier Neural Operator architecture. Our method treats time and space separately and as a result, it successfully propagates initial conditions in continuous time steps by employing the general composition properties of the partial differential operators. Following previous works, supervision is provided at a specific time point. We test our method on various time evolution PDEs, including nonlinear fluid flows in one, two, or three spatial dimensions. The results show that the new method improves the learning accuracy at the time of the supervision point, and can interpolate the solutions to any intermediate time.

1. INTRODUCTION

The evolution of classical and quantum physical dynamical systems in space and time is generically modeled by non-linear partial differential equations. Such are, for instance, Einstein equations of General Relativity, Maxwell equations of Electromagnetism, Schrödinger equation of Quantum Mechanics and Navier-Stokes (NS) equations of fluid flows. These equations, together with appropriate initial and boundary conditions, provide a complete quantitative description of the physical world within their regime of validity. Since these dynamic evolution settings are governed by partial differential operators that are often highly non-linear, it is rare to have analytical solutions for dynamic systems. This is especially true when the system contains a large number of interacting degrees of freedom in the non-linear regime. Consider, as an example, the NS equations, which describe the motion of viscous fluids. In the regime of high Reynolds numbers of the order of one thousand, one observes turbulences, in which all symmetries are broken. Despite much effort, two basic questions have remained unanswered, that is the existence and uniqueness of the solutions to the 3D NS equations and the anomalous scaling of the fluid observables in statistical turbulence. The solution to the (deterministic) NS equations seems almost random and is very sensitive to the initial conditions. Many numerical techniques have been developed for constructing and analysing the solutions to fluid dynamics systems. However, the complexity of these solvers grows quickly as the spacing in the grid that is used for approximating the solution is reduced and the degrees of freedom of the interacting fluid increases. Given the theoretical and practical importance of constructing solutions to these equations, it is natural to ask whether neural networks can learn such evolution equations and construct new solutions. The two fundamental questions are: (i) The ability to generalize to initial conditions that are different from those presented in the training set, and (ii) The ability to generalize to unseen time points, not provided during training.The reason to hope that such tasks can be performed by machine learning is that despite the seemingly random behaviour of, e.g. fluid flows in the turbulent regime, there is an underlying low-entropy structure that can be learnt. Indeed, in diverse cases, neural network-based solvers have been shown to provide comparable results to other numerical methods, while utilizing fewer resources. Our Contributions We present a hyper-network based solver combined with a Fourier Neural Operator architecture which is able to learn non-linear partial differential operators that govern the dynamics of chaotic and out of the equilibrium flows. 1. Our hyper-network architecture treats time and space separately. Utilizing a data set of initial conditions and the corresponding solutions at a labeled fixed time, the network learns a large class of time evolution PDEs. 2. Our approach enables interpolation to arbitrary (unlabelled) continuous times without additional data points. 3. Our solutions improve the learning accuracy at the supervision time-points. 4. We thoroughly test our method on various time evolution PDEs, including non-linear fluid flows in one, two and three spatial dimensions.

2. RELATED WORK

Hyper-networks While conventional networks employ a fixed set of pre-determined parameters, which is independent of the input, the hyper-network scheme, invented multiple times, and coined by Ha et al. ( 2017), allows the parameters of a neural network to explicitly rely on the input by combining two neural networks. The first neural network, called the hyper-network, processes the input or part of it and outputs the weights of a second neural network. The second network, called the primary network, has a fixed architecture and weights that vary based on the input. It returns, given its input, the final output. This framework was used successfully in a variety of tasks, ranging from computer vision (Littwin & Wolf, 2019), continual learning (von Oswald et al., 2020) , and language modeling (Suarez, 2017). While it is natural to learn functions with hyper-networks, since the primary network can be seen as a dynamic, input-dependent function, we are not aware of any previous work that applies this scheme for recovering physical operators.

Neural network-based PDE solvers

Due to the well-known limitations of traditional PDE solvers on one hand, and in light of new advances made in the field of neural networks on the other, lately we have witnessed very significant progress in the field of neural network-based PDE solvers (Karniadakis et al., 2021) . These solvers can be roughly divided into two groups according to the resource they utilize for learning: data-driven and model-based. Model-based solvers, known as Physics Informed Neural Networks (PINNs) (Raissi et al., 2019) , harness the differential operator itself for supervision. This is done by defining a loss, the residual of the PDE. These solvers also require a training dataset and can provide solutions for arbitrary times. Data-driven solvers are trained over a large dataset containing initial conditions and observed final states after the same time interval, T . Among these solvers there are encoder-decoder based ones, either a fully convolutional (Zhu & Zabaras, 2018) one, or a U-Net shaped one (Thuerey et al., 2020) , the PCANN (Bhattacharya et al., 2021) which applies Principal Component Analysis (PCA) as an auto-encoder and interpolates between the latent spaces using a neural network, The Multipole Graph Neural Operator (MGNO), the DeepONet which encodes input function and locations separately using two neural networks (Lu et al., 2019b ), Graph Neural Simulators Sanchez-Gonzalez et al. (2020) which utilize graph neural networks, and the MP-PDE (Brandstetter et al., 2021) , which learns a neural operator on the graph . These solvers learn solutions on a specific discretized grid, which poses a limitation for any practical applications. Recently, a mesh-invariant data-driven direction has been proposed (Lu et al., 2019a; Nelsen & Stuart, 2021; Anandkumar et al., 2020; Patel et al., 2021) . The mesh invariance is obtained by learning operators rather than mappings between initial and final states. This is achieved using architectures that enforce mesh invariant network parameters. Such networks have the ability train and infer on different meshes. As a result, these networks can be trained on small grids and run, during inference, on very large grids. Li et al. (2020b) have advanced the mesh-invariant line of works by introducing Fourier Neural Operators (FNO). FNOs utilize both a convolution layer in real space and a Fourier layer in the Fourier domain. Each Fourier layer first transforms its input, x ∈ R d to Fourier space, point-wise multiples the transformed input by a set of learned coefficients, w ∈ R d 2 (note that since x is a real vector, only d 2 coefficients are required), and applies an inverse Fourier transform to obtain the output, y = F -1 {w F {x}} ∈ R d . Since the learned parameters reside in Fourier space, it allows training on a specific scale, and inferring on other scales. It has been shown that the FNO solver outperforms previous solvers in a number of important PDEs. The current data-driven methods

