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

