A NEURAL PDE SOLVER WITH TEMPORAL STENCIL MODELING

Abstract

Numerical simulation of non-linear partial differential equations plays a crucial role in modeling physical science and engineering phenomena, such as weather, climate, and aerodynamics. Recent Machine Learning (ML) models trained on low-resolution spatio-temporal signals have shown new promises in capturing important dynamics in high-resolution signals, under the condition that the models can effectively recover the missing details. However, this study shows that significant information is often lost in the low-resolution down-sampled features. To address such issues, we propose a new approach, namely Temporal Stencil Modeling (TSM), which combines the strengths of advanced time-series sequence modeling (with the HiPPO features) and state-of-the-art neural PDE solvers (with learnable stencil modeling). TSM aims to recover the lost information from the PDE trajectories and can be regarded as a temporal generalization of classic finite volume methods such as WENO. Our experimental results show that TSM achieves the new state-of-the-art simulation accuracy for 2-D incompressible Navier-Stokes turbulent flows: it significantly outperforms the previously reported best results by 19.9% in terms of the highly-correlated duration time, and reduces the inference latency into 80%. We also show a strong generalization ability of the proposed method to various out-of-distribution turbulent flow settings.

1. INTRODUCTION

Complex physical systems described by non-linear partial differential equations (PDEs) are ubiquitous throughout the real world, with applications ranging from design problems in aeronautics (Rhie & Chow, 1983) , medicine (Sallam & Hwang, 1984) , to scientific problems of molecular modeling (Lelievre & Stoltz, 2016) and astronomical simulations (Courant et al., 1967) . Solving most equations of importance is usually computationally intractable with direct numerical simulations and finest features in high resolutions. Recent advances in machine learning-accelerated PDE solvers (Bar-Sinai et al. 2019; Li et al. 2020c; Kochkov et al. 2021; Brandstetter et al. 2021, inter alia) have shown that end-to-end neural solvers can efficiently solve important (mostly temporal) partial differential equations. Unlike classical finite differences, finite volumes, finite elements, or pseudo-spectral methods that requires a smooth variation on the high-resolution meshes for guaranteed convergence, neural solvers do not rely on such conditions and are able to model the underlying physics with under-resolved low resolutions and produce high-quality simulation with significantly reduced computational cost. The power of learnable PDE solvers is usually believed to come from the super-resolution ability of neural networks, which means that the machine learning model is capable of recovering the missing details based on the coarse features (Bar-Sinai et al., 2019; Kochkov et al., 2021) . In this paper, we first empirically verify such capability by explicitly training a super-resolution model, and then find that since low-resolution down-sampling of the field can lead to some information loss, a single coarse feature map used by previous work (Kochkov et al., 2021) is not sufficient enough. We empirically show that the temporal information in the trajectories and the temporal feature encoding scheme are crucial for recovering the super-resolution details faithfully. Motivated by the above observations, we propose Temporal Stencil Modeling (TSM), which combines the best of two worlds: stencil learning (i.e., Learned Interpolation in Kochkov et al. 2021) as that used in a state-of-the-art neural PDE solver for conservation-form PDEs, and HiPPO (Gu 

Raw

Figure 1 : Illustration of classic finite volume solvers (in red color), learnable solvers with vanilla stencil modeling (in blue color) and our temporal stencil modeling (in green color). While the convective flux approximation methods are different in each method, the divergence operator, the explicit time-step operator, and the pressure projection (in yellow color) are shared between classic solvers and learnable methods. Notice that the stencil interpolation coefficients in classic solvers such as WENO can also be data-adaptive (see Sec. 3.1 for more details). et al., 2020) as a state-of-the-art time series sequence model. Specifically, in this paper we focus on trajectory-enhanced high-quality approximation of the convective flux within a finite volume method framework. As illustrated in Fig. 1 , TSM can be regarded as a temporal generalization of classic finite volume methods such as WENO (Liu et al., 1994; Jiang & Shu, 1996) and recently proposed learned interpolation solvers (Kochkov et al., 2021) , both of which adaptively weight or interpolate the stencils based on the latest states only. On the other hand, in TSM we use the HiPPO-based temporal features to calculate the interpolation coefficients for approximating the integrated velocity on each cell surface. The HiPPO temporal features provide a good representation for calculating the interpolation coefficients, while the stencil learning framework ensures that the neural system's prediction exactly conserves the Conservation Law and the incompressibility of the fluid. With the abundant temporal information, we further utilize the temporal bundling technique (Brandstetter et al., 2021) to avoid over-fitting and improve the prediction latency for TSM. Following the precedent work in the field (Li et al., 2020c; Kochkov et al., 2021; Brandstetter et al., 2021) , we evaluate the proposed TSM neural PDE solver on 2-D incompressible Navier-Stokes equation, which is the governing equation for turbulent flows with the conservation of mass and momentum in a Newtonian fluid. Our empirical evaluation shows that TSM achieve both stateof-the-art simulation accuracy (+19.9%) and inference speed (+25%). We also show that TSM trained with steady-state flows can achieve strong generalization performance on out-of-distribution turbulent flows, including different forcings and different Reynolds numbers.

2.1. NAVIER-STOKES EQUATION

A time-dependent PDE in the conservation form can be written as ∂ t u + ∇ • J(u) = 0 (1) where u : [0, T ] × X → R n is the density of the conserved quantity (i.e., the solution), t ∈ [0, T ] is the temporal dimension, X ⊂ R n is the spatial dimension, and J : R n → R n is the flux, which

