A Neural PDE Solver with Temporal Stencil Modeling
Authors: Zhiqing Sun
Yiming Yang
Shinjae Yoo
Venue: Submitted to ICLR 2023 (2023)
Keywords: neural PDE solver
Navier-Stokes equation
turbulent flow
Computational Fluid Dynamics
CFD

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.

Under review as a conference paper at ICLR 2023
A NEURAL PDE SOLVER WITH TEMPORAL STENCIL MODELING
Anonymous authors Paper under double-blind review
ABSTRACT
Numerical simulation of non-linear partial differential equations plays crucial role in modeling physical science and engineering phenomena, such as weather, climate, and aerodynamics. Recent Machine Learning (ML) models trained on low-resolution apatio-tempors signals have shown new promises in capturing im- portant dynamics in high-resolution signals, under the condition that the models can effectively recover the missing details. However, this study shows that signifi- cant information is often lost in the low-resolution down-sampled features. To ad- dress such issues, we propose a new approach, namely Temporal Stencil Modeling (TSM), which combines the strengths advanced time-series sequence modeling (with the HiPPO features) and state-of-the-art neural PDE solvers (with learn- able stencil modeling). TSM aims to recover the lost information from the PDE trajectories and can be regarded ume 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% 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 diitttiiiittiitti turbulent flow settings.
temporal generalization of classic finite vol-
INTRODUCTION
Complex physical systems described by non-linear partial differential equations (PDEs) are ubiqui- tous 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 equa- tions of importance is usually computationally intractable with direct numerical simulations and finest features in high resolutions.
ceaming-accelerated Kochkov et al. [021] (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 pscudo-spectral methods that requires 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-qualit, simulation with significantly reduced computational cost.
Recent advances in machine learning- accelerated PDE solvers (Bar-Sinat et al. [209 Li et al. 2020c
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, 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 com- bines the best of two worlds: stencil learning (i.c., Learned Interpolation Kochkov et al. (02T) that used in state-of-the- neural PDE solver for conservation-fome PDEs, and HiPPO (Gu

---

Under review as a conference paper at ICLR 2023
Veactt by tryedtory
e
Latest vulocky
uTT)
HPРО
Schwmen WENG)
Stencil Modeling (KoehNov
Temporal (ours)
Convaactive Tux approximatin
Convective fux approsimatio
Carvective fux approximaton
Divergence
Expliatt timeslop
External feedng F(?)
Presoure projection
New volocity o + Ar)
Calculeiing atancil Interpolation confiiierts
Figure I: 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 aata-adaptive (see Sec. 3.1 for more details).
al. 2020) as a state-of-the -ft-----tth--tttt art time series sequence model. Specifically, in this paper we focus on trajectory- -enhance high-quality approximation of the convective flux within finite volume method framework. As illustrated in Fig. [ TSM can be regarded as a temporal generalization classic finite volume methods such as WENO (Liu et al. 1994 Jiang & Shu 1996) and recently proposed learned interpolation solvers (Kochkov et al. al.] 202 l), both which adaptively weight or interpolate the stencils based the latest states only. On the other hand, in TSM
the HiPPO-based temporal features to calculate the interpolation coefficients for approximating the integrated veloc- ity 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 al. T02T 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 Newtonian fluid. Our empirical evaluation shows that TSM achieve both state- of-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 BACKGROUND & RELATED WORK
2. NAVIER-STOKES EQUATION
A time-dependent PDE in the conservation form can be written as
0,u + V J(u)= = 0
(1)
where u [0,T] x X + R" is the density of the conserved quantity (i.c., the solution), t E [o. T] the temporal dimension, X IR" is the spatial dimension, and J R R is the flux, which
2

---

Under review as a conference paper at ICLR 2023
represents the quantity that pass or travel (whether it actually moves not) through surface or substance. V is the divergence of J. Specifically, the incompressible, constant density 2-D Navier Stokes equation for fluids has conservation form of:
uu + V. (u@u) = vV Vp+f
(2)
(3)
where @ denotes the tensor product, V is the kinematic viscosity, p is the fluid density, p is the pressure filed, and f is the external forcing. In Eq. 2 the left-hand side describes acceleration and convection, and the right-hand side is in effect summation of diffusion, internal forcing source, and external forcing source. Eq. enforres the incompressibility of the fluid.
A common technique to solve time-dependent PDEs is the method of lines (MOL) (Schiesser 2012) where the basic idea is to replace the spatial (boundary value) derivatives in the PDE with algebraic approximations. Specifically, we discretize the spatial domain into grid X G", where G is a set of grids on F. Each grid cell in G denote a small non non-overlapping volume, whose center Xg. and the average solution value is calculated as uy = I, ull, x)dx. We then solve a.u', for E C" and & E [0. T}. Since E G is a set of pre-defined grid points, the only derivative operator now in time, making an ordinary differential equations (ODEs)-based system that approximates the original PDE.
2.2 CLASSICAL SOLVERS FOR COMPUTATIONAL FLUID DYNAMICS
In Computational Fluid Dynamics (CFD) (Anderson & Wendt 1995 Pope & Pope 2000), the Reynolds number Re ULVv dictates the balance between convection and diffusion, where U and are the typical velocity and characteristic linear dimension. When the Reynolds number Re >> 1, the fluids exhibit time-dependent chaotic behavior, known turbulence, where the small- scale changes in the initial conditions can lead to a large difference in the outcome. The Direct Numerical Simulation (DNS) method solve Eq. 2 directly, and is general-purpose solver with high stability. However, as Re determines the smallest spatio-temporal feature scale that need to be cap- tured by DNS, DNS faces a computational complexity as high as (Re)) (Choi & Moin 2012) and cannot scale to Reynolds number flows or large-size computation domains.
2.3 NEURAL PDE SOLVERS
A wide range of neural network-based solvers have recently been proposed to at the intersection of PDE solving and machine learning. We roughly classify them into four categories:
Physics-Informed Neural Networks (PINNs) PINN directly parameterizes the solution u as X
neural network F:
R" (Weinan Yu 2018: Raissi al. 2019 Bar & Sochen 2019) 00 T x
Smith et al. 2020 Wang et al. 2022). They are closely related to the classic Galerkin methods (Matthies & Keese 2005), where the boundary condition date-fitting losses and physics-informed
losses are introduced to train the neural network. These methods suffers from the parametric depen- dence issue, that is, for any new inital and boundary conditions, the optimization problem needs to be solved from scratch, thus limit their applications especially for time-dependent PDEs.
Neural Operator Learning Neural Operator methods learn the mapping from any functional parametric dependence to the solution as F (0, T" x X + R") - ([,, + AT] X - RT) Lu
et al. 2019: BBhattacharya et al.) 2020 Patel et al. 2021). These methods are usually not bounded by fixed resolutions, and Tearn to directly predict any solution at time step / E [T. T AT). Fourier transform Li al. 2020c Tran et al. (021). wavelet transform (Gupta et al. 2021), random fea- tures (Nelsen Stuart 202 l), attention mechanism (Cao 202T), or graph neural networks al.) 2020alby are often used in the neural network building blocks. Compared to neural methods that mimic the method of lines, the operator learning methods are not designed to generalize to dynam- ics for out-of-distribution t [T' +oo], and only exhibit limited accuracy for long trajectories.
Neural Method-of-Lines Solver Neural ofthod-of-Linnn Solvers are autoregressive models that solve the PDE iteratively, where the difference from the latest state time T to the state time
3

---

Under review as a conference paper at ICLR 2023
T + At is predicted by neural network F ([0,T] x X + X (X - +R*))=++r
+-T+Ar The
typical choices for F include modeling the absolute difference: Vg E X G" u + u +
F Fo (u(or))))) [Wang at. 2020 Sanchez- at. 2020; Stachenfed al.
2021 and modeling
the relative difference: " + u + St F ((o..T) (Brandstetter et al. 2021 where the latter
is believed to have the better consistency property, i.e., limar=0 1
0.
Hybrid Physics-MI. Physics-MI. hybrid models is recent line of work that uses neural network to correct the errors in the classic (typically low-resolution) numerical simulators. Most of these approaches seek to learn the corrections of the numerical simulators' outputs (Mishra 2019 Um
et al. 20200 List et al. 2022) Dresdner et al. 2022,] Frezat el al., 2022, Bruno et al. 2002). while Bar- Sinai et al. 2019); Kochkov et al. (2021) learn infer the stencils of advection -iifuuionn problems in a Finite Volume Method (FVM) framework. The proposed Temporal Stencil Modeling (TSM)
method belongs to the latter category.
TEMPORAL STENCIL MODELING FOR PDES
3.1 NEURAL STENCIL MODELING IN FINITE VOLUME SCHEME
Finite Volume Method (FVM) is a special MOL technique for conservation form PDEs, and can be derived from Eq. via Gauss' theorem, where the integral of (i.e., the averaged vector field of volume) over unit cell increases only by the net flux into the cell. Recall that the incompressible, constant density Navier Stokes equation for fluids has a conservation form of:
0u + V (u H u) = u 1 Vp+f
(4)
We can see that in FVM, the cell-average divergence can be calculated by summing the surface flux, so the problem boils down to estimating the convective flux u & u each face. This only requires estimating u by interpolating the neighhoring discretized velocities, called stencils. The beauty of FVM is that the integral of u is exactly conserved, and it can preserve accurate simulation as long as the flux u & u is estimated accurately. Fig. Iillusrrates an implementation (Kochkov et al. 2021] 1) of classic FVM for the Navier-Stokes equation, where the convection and diffusion operators are based on fiite-difference approximations and modeled by explicit time integration, and the pressure is implicitly modeled by the projection method (Chorin) 1967). The divergence operator enforces local conservation of momentum according to a finite volume method, and the pressure projection enforces incompressibility. The explicit time-step operator allows for the incorporation of additional time-varying forces. We refer the readers to (Kochkov et al.) 202T) for more details.
Classic FVM solvers use manually designed n* ooder accurate method to calculate the interpo- lation coefficients of stencils to approximate the convective flux. For example, linear interpola- tion, upwind interpolation (Lax 1959). and WENO5 method (Shu 2003) Gottlieb et al., 2006) can achieve first, second, and fifth-order accuracy, respectively. Besides, adjusting the interpolation weights adaptively based on the input data is not new in numerical simulating turbulent flows. For example, given the interpolation axis, the upwind interpolation adaptively uses the value from the previous/next cell along that axis plus a correction for positive/negative velocity. WENO (weighted essentially non-oseillatory) scheme also computes the derivative estimates by taking an adaptively- weighted average over multiple estimates from different neighborhoods. However, the classic FVM solvers are designed for general cases and can only adaptively adjust the interpolation coefficients with simple patterns, and thus are sub-optimal when abundant PDE observation data is available.
In this paper, we follow Kochkov et al. (2021) and aim to learn more accurate flux approximation in the FVM framework by predicting the Iearnable interpolation coefficients for the stencils with neural networks. In principle, with x convolutional kernels, a 1, 2, 3-layer neural network is able to perfectly mimic the linear interpolation, Lax- Wendroff method, and WENO method, respectively (Brandstetter et al. 2021). Such observation well connects the learnable neural stencil modeling methods to the classical schemes. However, previous work (Bar-Sinai et al. 2019 Kochkov et al. 2021 investigates the learned interpolation scheme that only adapts to the latest state u' which still uses the same information as classical solvers. In this paper, further generalize both classical and previous learnalle stencil interpolation schemes by predicting the interpolation coefficients with the abundant information from all the previous trajectories {u'l (t [0, T)]
4

---

Under review as a conference paper at ICLR 2023
Table 1: 64 x 64 - 2048 x 2048
super-resolution MSE with differ- ent approaches.
Method
MSE
Fross we HPPO frotures
Enora Roalurs
Bicubic Interpolation
2.246
CNN w/ 1-step raw features
0.029
CNN w/ 32-step raw features
0.015
Figure 2: 64 64 - 2048 x 2048 super-resolution emrors with 32ep-step HiPPO features and 32At-step raw features.
w/ 32-step HiPPO features
0.007
3.2 TEMPORAL RESPORRNSOOOUTIONN WITH HIPPO FEATURES
fundamental question in neural stencil modeling is why ML models can predict more accurate flux approximations, and previous work (Kochkov et al., 2021) attribute their success the super- resolution power of neural nteworks, that is, the machine learning models can recover the missing details from the coarse features. In this paper, empirically verify this hypothesis by explicitly training a super-resolution model.
Specifically, we treat the 2-D fluid velocity map as W image, and train a CNN-based U-Net decoder (Ronneberger et al. (2155) to generate the super-resolution vorticity results. Our results are reported in Tab. (right), and we find that the super-resolution results generated by neural networks is nearly 100x better than bicubic interpolation, which verify the super-resolution power of MI. models to recover the details.
Next, since the temporal information is always available for PDE simulation we investigate whether the temporal information can further reduce the super-resolution errors, i.e., recovering more details from the coarse features. After preliminary experiments, we decide to keep the convo- lutional neural networks as the spatial encoder module for their efficient implementation on GPUs and translation invariant property, and only change the temporal input features. Inspired by the cent progress in time series modeling, we consider the following two types of features as the CNN model's 's inputs:
Raw features Following the previous work video super-resolution (Liao al. the W xT'x x2 -shaped velocity trajectory as an W image with feature channels C
2015
we treat T 2.
HiPPO features HiPPO (High-order Polynomial Projection Operators) (Gu et al., 2020 2021)
recently proposed framework for the online compression of continuous time series by projection onto polynomial bases. It computes the optimal polynomial coefficients for the scaled Legendre state-of-the-art autoregressive sequence model for raw images (Ty et al. 2020) and audio (Goel et al. 2022). In this paper, we propose to adopt HiPPO to encode the raw fluid velocity trajectories. Due the space limit, we leave the general description of the HiPPO technique to Appendix
polynomial basis. It has been shown
Results We report the temporal super-resolution results in Tab. From the table, we can see that the 32At-step raw features can reduce the super-resolution error half, and using the HiPPO to encode the time series can further reduce the error scale by half. We also visualize the errors of 32A/-step HiPPO features and 32A/-step raw features in Fig. 2 Our temporal super-resolution results show that plenty of information in the PDE low-resolution temporal trajectories is missed with vanilla stencil modeling, will be underutilized with raw features. and HiPPO can better exploit the temporal information in the velocity trajectories.
3.3 TEMPORAL STENCIL MODELING WITH HIPPO FEATURES
As better super-resolution performance indicate that more details can be recovered from the low- resolution features, we propose compute the interpolation coefficients in convective flux with
Even if we only have access to step of the ground-truth trajectory, can still use the high-resolution DNS to generate long enough trajectory initialize the spatial-temporal model.
5

---

Under review as a conference paper at ICLR 2023
the HiPPO-based temporal information, which should lead to accurate flux approximation. Incorporating HiPPO features in the neural stencil modeling framework is forwart with ground-truth initial velocity trajectories V[o.7], first recurrently encode the trajectory (step by step with Eq. 7 in Appendix and use the resulted features (v)) to compute the interpolation coefficients for the stencils. Given model-geeneaate new velocity VT+At, due to the recurrence property of HiPPO, feature HYPPO(y1o.T+a)) which is very efficient. Fig. 9 in the appendix illustrates such process.
only apply single update on HiPPO and get the new encoded
4 EXPERIMENTS
4.1 EXPERIMENTAL SETUP
Simulated data Following previous work (Kochkov et al., 2021), we train our method with 2- D Kolmogorov flow, a variant of incompressible Navier-Stokes flow with constant forcing sin(4y) 0.1u. All training and evaluation data are generated with a JAX-based finite volume- based direct numerical simulator in stagerredssuaae mesh (McDonough) 2007) as briefly de-
scribed in Sec. [3.] We refer the readers to the appendix of (Kochkov et al. 2021 for more data generation details.
0.001 on a 2n x 2 domain, which results a time-step of A/ = 7.0125 10-3 according to the the 64 x 64 simulation grid. For training, we generate 128 trajectories fluid dynamics, each starts with different random initial conditions and simulates with 2048 2048 resolution for 40.0 time units. We use 16 trajectories for evaluation.
We train the neural models on Re 1000 flow data with density P and viscosity
(CFD) condition
Unrolled training All the learnable solvers are trained with the Mean Squared Error (MSE) loss on the velocities. Following previous work (Li et al. [020c] Brandstetter et al. 2021 l Kochkov et al.) 202T Dresdner al. 2022), we adopt the unrolled training technique, which requires the learnable solvers mimic the ground-truth solution for more than one unrolled decoding step:
N
C(u gt (o.T) Σ 1=1
MSE (14), prred (t:))
(5)
N
where f, [TT At,..,, Nat) is the unrolled time steps. Uyt and Upred are the ground-truth solution and learnable solver's prediction, respectively. Unrolled training can improve the inference performance but makes the training less stable (due to bad initial predictions), so we use N 32
unrolling steps for training.
Temporal bundling In our preliminary experiments, we found that due the abundant informa- tion in the trajectories, the TSM solvers are more prone to over-fit. Therefore, we adopt the tem- poral bundling technique (Brandstetter et al. 2021) the learned interpolation coefficients in TSM solvers. Assume that in step-by- prediction scheme, we predict u(o.7] - c(T At), where c is the stencil interpolation coefficients in the convective flux approximation. In temporal bundling, we predict K steps of the interpolation coefficients U|o.T] + {c(T+Att) c(T+2At), c(T K At)] advance, and then time-step forward the FVM physics model for K steps with pre-computed sten- cil interpolation coefficients.
Neural network architectures & hyper-parameters In TSM-64 x 64 with a T-length trajectory,
the input and output shapes of TSM are 64 x 64 x T x 2 and 64 x 64 x C x 2, while the input and output shapes of CNN are 64 64 (C 2) and 64 64 x (8 x (42 - 1)). which represents that for each resolution, we predict 8 interpolations that need 15 inputs each. For HiPPQS) we set the hyper- syper-parameters
-0.5, 1.0, dt 1.0. For CNN, we use 6-layer network with 3 x kernels and 256 channels with periodic padding

---

Under review as a conference paper at ICLR 2023
DNE
324
84
TEWH p405
Figure 3: Qualitative results of predicted vorticity fields for reference (DNS 2048 2048), previous learnalle sota model (LI 64 x 64) (Kochkov et al. (021), and method (TSM 64 x 64), starting from the same initial condition. The yellow box denotes a vortex that is not captured by LI
D15 (reriiee
1or sutii 54-8 &
04:04 Mochoryct 2011
X0
Simiiatioon 2018
Figure 4: (left) Comparison of the vorticity correlation between prediction and the ground-trut trrnh so- lution (i.e., DNS 2048 x 2048). (middle) Energy spectrum scaled by k averaged between simulation time 6.0 to 20.0. (right) Comparison high vorticity correlation duration v.s. inference latency.
4.2
MAIN RESULTS
The classic and neural methods evaluated be roughly classified into four categories: 1) pure Physics models, i.e., the FVM-based Direct Numerical Simulation (DNS), 2) pure Machine Learning (ML.)based neural method-of-lines models, 3) Learned Correction (L.C) models, which correct the final outputs (i.e., velocities) of physics models with neural networks (Um et al., 2020). and 4) Neural Stencil Modeling models, including single-time-step Learned Interpolation (LI) Kochkov et al. 2021), and our Temporal Stencil Modeling (TSM).
For neural network baselines, except the periodic Convolutional Neural Networks (CNN) (LeCun et al., 1999; Kochkov et al., 202T) with raw and HiPPO features we already described, we also compare with Fourier Neural Operators (FNO) Li et al. 2020c) and Multiwavelet-basee model (MWT), two recent -fate-of-the-arr pure-ML PDE solver based on spectral and wavelet features.
We evaluate all the solvers based on their Pearson correlation P with the ground-truth (i.e., DNS with the highest 2048 x 2048 resolution) flows in terms of the scalar vorticity field w = Ō2 lly Uy Furthermore, ease comparing all the different solvers quantitatively, we focus on their high-correlation duration time, i.e., the duration of time until correlation p drops below 0.8.
The comparison between HiPPO-based TSM and I-time-step raw-feature (Kochkov et al. 2021) is shown in Fig- and Fig. f (left). We can see that our HiPPO feature-based TSM significantly outperforms the previous ofate-of-the-art ML-physics model, especially when trained with a 4-step temporal bundling. From Fig. (middle))
that all the learnable solvers can better capture the frequecyy features with similar energy spectrum E(k) = [u(k)) pattern as the high- resolution ground-truth trajectories. From Fig. (right),
that with the help of temporal
https::/gii http::://gtthub..com////////////

---

Under review as a conference paper at ICLR 2023
Table 2: Quantitative comparisons with the metric of high-coreelation (p > 0.8) duration (w.rt the reference DNS-2043 2048 trajectories). All learnable solvers use the 64 64 grids.
raw features
HIPO features
Method
Type
High-com: duration
righ cum 200)
DNS-64 64
Physics
2.805
85
DNS-128 x 128
Physics
3,983
DNS-256 x 256
Physics
5,386
DNS-512 x 512
Physics
6,788
DNS-1024 1024
Physics
8.752
75-
1-step-raw -NNN
ML
4.824
step--ra
ML
7.517
32-step-FNO
6,283
Inttal veloc traeecory longth
Ta
32-step-WMB WMT
ML
5.890
1-step-raw-CN CNN
LC
6,900
Figure 5: Temporal stencil modeling perfor-
32-step-FNO
LC
7.630
mance nigh-correlation duration) with dif-
1-step-raw-CNN
LI
7,910
ferent feature types and different initial
32-step-FNO
TSM
7.798
jectory steps.
4-step-raw-CN CNNN
TSM
8.359
+ 4-step lemporal-bundle
TSM
8,303
32-step-HiPPO-CCE
TSM
9.256
+ stett temporal-hundle
TSM
9.481
bundling, HiPPO-TSM can achieve high simulation accuracy while reduces the inference latency into 80% when compared original LI.
4.3 ABLATION STUDY
We present a quantitative comparison for more methods in Tab. We can that the learned correction models are always better than pure-ML models, and neural stencil modeling models are always better than learned correction models. In terms of neural network architectures, we find that under the same ML--physies framework, raw-feature CNN is always better than FNO, and HiPPO- feature CNNs are always better than raw-feature CNNs. Finally, when adopting temporal bundling, we can see that only HiPPO-TSM can benefit from alleviating the over-fitting problem, while the performance of rewtrreture TSM can be hurted by temporal bundling.
We also study the impact of the trajectory length raw-feature and HiPPO-featee TSM models in Fig. Notice that in raw features, the temporal window size fixed during unrolling. while in HiPPO Teatures, the temporal window size is expanded during unrolling decoding. From the figure, we can see that the raw-feature CNN achieves the best performance with a window size of 4, while HiPPO-feature keep increasing the performance, and reach the peak with 32 initial trajectory length.
4.4 GENERALIZATION TESTS
We evaluate the generalization ability of our HiPPO-TSM, (4-step bundle) model and LI trained on trained on Kolmogorov flows (Re- = 1000). Specifically, we consider the following test cases: (A) decaying flows (starting Re 1000), (B) more turbulent Kolmogorov flows (Re 4000), and (C) 2x larger domain Kolmogorov flows (Re = 1000). Our results are shown in Fig. 66 from which can see that HiPPO-TSM achieves consistent improvement over LI. HiPPO-TSM also achieves competitive performance to DNS- 1024 x 1024 or DNS-2048 x 2048, depending on the ground-truth (i.e., highest resolution) being DNS-2048 2048 or DNS-4096 4096).
CONCLUSION & FUTURE WORK
In this paper, we propose a novel Temporal Stencil Modeling (TSM) method for solving time- dependent PDEs in conservation form. TSM can be regarded the temporal generalization of classic finite volume solvers such as WENO and vanilla neural stencil modeling methods (Kochkov) et al. 2021), in that TSM leverages the temporal information from trajectories, instead of only using
8

---

Under review as a conference paper at ICLR 2023
MUITERy cermaati
1251220
D16 256:250
2244
51d4 (oacloy di 2001
Saruz
DR5 302323633
75%4 797 04 4bonde (5.25)
Inrion 17m6
Verbubyy cormaaattoo
DRG 250G2E
512-521
095 234502200
(R1 4235-4333
200
Learndd Interiolation (fochkov et al. 2021)
Temporal Stenci Modelng
Mighestreacktioon DNS
(eurs)
Figure 6: Generalization results of neural methods trained Kolmogorov flows (Re 1000) and evaluated on (A) decaying flows (starting from Re 1000), (B) more turbulent Kolmogorov flows (Re 4000), and (C) 2x larger domain Kolmogorov flows (Re 1000).
the latest states, to approximate the (convective) flux more accurately. Our empirical evaluation 2-D incompressible Navier-Stokes turbulent flow data show that both the temporal information and its temporal feature encoding scheme crucial achieve oftte--f--he-rtt simulation accuracy. We also show that TSM have strong generalization ability on various out-of-distribution turbulent flows.
For future work, we plan evaluate TSM method on the more challenging and realistic 3-D turbulent flows, as well as 2-D turbulent flows with non-periodic boundary conditions. We are also interested in leveraging the Neural Architecture Search (NAS) technique automatically find better features and neural architectures for solving Navier-Stokes equation in the TSM framework.
9

---

Under review as a conference paper at ICLR 2023
ETHIC STATEMENT
We do not see obvious negative social impact of our work.
REPRODUCIBILITY STATEMENT
We release the source code at hnttps://anonymouuuls after the acceptance of the manuscript. Since the training and evaluation data in this paper are all simulated, they can also be faithfully reproduced from our release code.
We plan to open-source all the code
REFERENCES
Giancarlo Alfonsi. eeynolds-averagee stvier-stokes equations for turbulence modeling. Applied Mechanics Reviews, 62(4), 2009.
John David Anderson and John Wendt. Computational fluid dynamics, volume 206. Springer, 1995.
Leah Bar and Nir Sochen. Unsupervised deep learning algorithm for pde-based forward and inverse problems. arXiv preprint 1rXiv:1904.054:: 2019.
Yohai Bar-Sinai, Stephan Hoyer, Jason Hickey, and Michael P Brenner. Learning data-driven dis- cretizations for partial differential equations. Proceedings of the National Academy of Sciences, 116631)11334415549) 2019.
Kaushik Bhattacharya, Bamdad Hosseini, Nikola B Kovachki, and Andrew M Stuart. Model reduc- tion and neural networks for parametric pdes. arXiv preprint arXiv:2005.031 20000000000,,,,. 2020.
J Bouussineq. Theorie de I'ecoulement tourbillant. Mem. Acad. Sci., 23:46, 1877.
Johannes Brandstetter, Daniel E Worrall, and Max Welling, Message passing neural pde solvers. In International Conference on Learning Representations, 2021.
Oscar P Bruno, Jan S Hesthaven, and Daniel V Leibovici. Fe-based shock-dynamics solver with neural-network localized artificial-viscosity assignment. Journal of Computational Physics: X, pp. 1001 10, 2022.
Shuhao Cao. Choose transformer: Fourier or galerkin. Advances Neural Information Processing Systems, 34:24924 24940, 2021.
Haecheon Choi and Parviz Moin. Grid-point requirements for large eddy simulation: Chapman's estimates revisited. Physics of fluids, 24(1):011702 2012.
Alexandre Joel Chorin. The numerical solution of the navier-stokes equations for incompressible fluid. Bulletin of the American Mathematical Society, 73(6):928 -011,,911111 1967.
Richard Courant, Kurt Friedrichs, and Hans Lewy. On the partial difference equations of mathemat- ical physics. IBM journal of Research and Development, 11(2):15-22 1967.
Gideon Dresdner, Dmitrii Kochkov, Peter Norgaard, Leonardo Zepeda-Núñez, Jamie A Smith, Michael Brenner, and Stephan Hoyer. Learning to correct spectral methods for simulating turbulent flows. arXiv preprint arXiv:2207.00556 2022.
Hugo Frezat, Julien Le Sommer, Ronan Fablet, Guillaume Balarac, and Redouane Lguen- sat. A posteriori learning for guasi-gcostrophic turbulence parametrization. arXiv preprint arXiv:2204.0391/ 2022.
Karan Goel, Albert Gu, Chris Donahue, and Christopher Ré. It's raw! audio generation with state- space models. arXiv preprint arXiv:2202.09729, 2022.
Sigal Gottlieb, Julia Mullen, and Steven J Ruuth. fifth order flux implicit weno method. Journal of Scientific Computing, 27(1):271-287, 2006.
10

---

Under review as a conference paper at ICLR 2023
Albert Gu, Tri Dao, Stefano Ermon, Atri Rudra, and Christopher Ré. Hippo: Recurrent memory with optimal polynomial projections. Advances Neural Information Processing Systems, 33: 1474-1487, 2020.
Albert Gu, Karan Goel, and Christopher Re. Efficiently modeling long sequences with structured state spaces. In International Conference on Learning Representations, 2021.
Gaurav Gupta, Xiongye Xiao, and Paul Bogdan. Multiwavelet-based operator learning for differen- tial equations. Advances in Neural Information Processing Systems, 34:24048-24062, 2021.
Guang-Shan Jiang and Chi-Wang Shu. Efficient implementation of weighted eno schemes. Journal of computational physics, 126(1): (62222200- 1996.
Dmitrii Kochkov, Jamie A Smith, Ayya Alieva, Qing Wang, Michael P Brenner, and Stephan computational fluid dynamics. Proceedings of the National
Hoyer. Machine Academy of Sciences, (18(21):e2101784) 2021.
Peter Lax. Systems of conservation laws. Technical report, LOS ALAMOS NATIONAL LAB NM, 1959.
Yann LeCun, Patrick Haffner, Léon Bottou, and Yoshua Bengio. Object recognition with gradient- based learning. In Shape, contour and grouping in computer vision, pp. 319-345. Springer, 1999.
Tony Lelievre and Gabriel Stoltz. Partial differential equations and stochastic methods in molecular dynamics. Acta Numerica, 25:681-880 2016.
Marcel Lesieur and Olivier Metais. New trends in large-eddy simulations of turbulence. Annual review of fluid mechanics, 28(1): -22 1996.
Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, An- drew Stuart, and Anima Anandkumar. Neural operator: Graph kernel network for partial differ- ential equations. arXiv preprint 000.:20003.034888. 2020a.
Zongyi Li, Nikola Kovachki, Kamyar Arizzadenesheli, Burigede Liu, Andrew Stuart, Kaushik Bhat- tacharya, and Anima Anandkumar. Multipole graph neural operator for parametric partial differ- ential equations. Advances in Neural Information Processing Systems, 33:6755-6766, 2020b.
Zongyi Li, Nikola Borislavov Kovachki, Kamyar Azizzadenesheli, Kaushik Bhattacharya, Andrew Stuart, Anima Anandkumar, et al. Fourier neural operator for parametric partial differential equa- tions. In International Conference on Learning Representations, 2020c.
Renjie Liao, Xin Tao, Ruiyu Li, Ziyang Ma, and Jiaya Jia. Video super-resolution via deep draft- ensemble learning. Proceedings of the IEEE international conference on computer vision, pp- 531-539, 2015.
Bjöm List, Li-Wei Chen, and Nils Thuerey. Learned turbulence modelling with differentiable fluid solvers. arXiv preprint arXiv:2202.06980 2022.
Xu-Dong Liu, Stanley Osher, and Tony Chan. Weighted essentially non-oscillatory schemes. Jour- nal of computational physics, 115(1)2000-214 1994.
Lu Lu, Pengzhan Jin, and George Em Karniadakis. Deeponet: Learning nonlinear operators for iden- tifying differential equations based on the universal approximation theorem of operators, arXiv 2019.
preprint
Hermann G Matthies and Andreas Keese. Galerkin methods for linear and nonlinear elliptic stochas- tic partial differential equations. Computer methods applied mechanics and engineering, 194 (12-16):1295-1333 2005.
James M McDonough.
Lectures in computational fluid dynamics of incompressible flow: Mathe- matics, algorithms and implementations. 2007.
Siddhartha Mishra. A machine learning framework for data driven acceleration of computations of differential equations. Mathematics in Engineering, 1(1):118-144) 2019.

---

Under review as a conference paper at ICLR 2023
Nicholas H Nelsen and Andrew M Stuart. The random feature model for input- maps between banach spaces. SIAM Journal on Scientific Computing. 43(6):A3212-A3243 2021.
Ravi G Patel, Nathaniel A Trask, Mitchell A Wood, and Eric C Cyr. A physics-informed operator re- gression framework for extracting data-driven continuum models. Computer Methods in Applied Mechanics and Engineering, 373:1 13500, 2021.
Stephen B Pope and Stephen B Pope. Turbulent flows. Cambridge university press, 2000.
Maziar Raissi, Paris Perdikaris, and George E. Karniadakis. Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations. Journal of Computational physics, 378:686-707, 2019.
Chae M Rhie and Wei-Liang Chow. Numerical study of the turbulent flow past airfoil with trailing edge separation. AIAA journal, 21(11):1525-1552 1983.
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedi- cal image segmentation. In International Conference Medical image computing and computer- assisted intervention, pp. 234-241. Springer. 2015.
Ahmed M Sallam and Ned HC Hwang. Human red blood cell hemolysis in a turbulent shear flow: contribution of reynolds shear stresses. Biorheology. 21(6):783-79) 797, 1984.
Alvaro Sanchez-Gonzales Jonathan Godwin, Tobias Pfaff, Rex Ying, Jure Leskovec, and Peter Battaglia. Learning to simulate complex physics with graph networks. In International Confer- ence on Machine Learning, pp. 8459-8468. PMLR, 2020.
William E. Schiesser. The numerical method of lines: integration of partial differential equations. Elsevier, 2012.
Chi- Chi-Wang Shu. High-order finite difference and finite volume weno schemes and discontinuous galerkin methods for cfd. International Journal of Computational Fluid Dynamics, 17(2):107- 118, 2003.
Jeffrey P Slotnick, Abdollah Khodadoust, Juan Alonso, David Darmofal, William Gropp, Elizabeth Lurie, and Dimitri Mavriplis. Cfd vision 2030 study: path to revolutionary computational aerosciences. Technical report, 2014.
Joseph Smagorinsky. General circulation experiments with the primitive equations: the basic experiment. Monthly weather review, 91(3):99-164, 1963.
Jonathan D Smith, Kamyar Azizzadenesheli, and Zachary E Ross. Eikonet: Solving the eikonal equation with deep neural networks. IEEE Transactions on Geoscience and Remote Sensing, 59 (12):10655-10666, 2020.
Kimberly Stachenfeld, Drummond B Fielding, Dmitrii Kochkov Miles Cranmer, Tobias Pfaff, Jonathan Godwin, Can Cui, Shirley Ho, Peter Battaglia, and Alvaro Sanchez-Gonzale Learned coarse models for efficient turbulence simulation. arXiv preprint arXiv:2112.15272 2021.
Makoto Takamoto, Timothy Praditia, Raphael Leiteritz, Dan MacKinlay, Francesco Alesiani, Dirk Pflüger, and Mathias Niepert. Pdehench: An extensive benchmark for scientific machine learning. In Thirty-sixth Conference on Neural Information Processing Systems Datasets and Benchmarks Track.
Yi Tay, Mostafa Dehghani. Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang, Sebastian Ruder, and Donald Metzler. Long range arena: A benchmark for efficient transformers. In International Conference on Learning Representations, 2020.
Alasdair Tran, Alexander Mathews, Lexing Xie, and Cheng Soon Ong. Factorized fourier neural
operators. arXiv preprint arXiv:2111.13802 2021.
Kiwon Um, Robert Brand, Yun Raymond Fei, Philipp Holl, and Nils Thuerey. tnoo: Learning from differentiable physics to interact with iterative poe-solvers. Advances in Neural Information Processing Systems, 33:6111-6122. 2020.
12

---

Under review as a conference paper at ICLR 2023
Chuwei Wang, Shanda Li, Di He, and Liwei Wang. Is 12 physics-informed loss always suitable for training physics-infhi informed neural network? arXiv preprint arXiv:2206.0200 02016, 2022.
Rui Wang, Karthik Kashinath, Mustafa Mustafa, Adrian Albert, and Rose Yu. Towards physics- informed deep learning for turbulent flow prediction. In Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, pp. 1457-1466, 2020.
E Weinan and Bing Yu. The deep ritz method: A deep learning-based numerical algorithm for solving variational problems. Communications Mathematics and Statistics, 1(6):1-12, 2018.
13

---

Under review as a conference paper at ICLR 2023
64
64 64
Temporal supor-resolute
Vanila superresoltion
Figure 7: Illustration of super-resolution process from trajectory of frames (i.e., temporal super- resolution) or single-frame (vanilla super-resolution).
A ADDITIONAL RELATED WORK
A.1 INDUSTRIAL SOLVERS FOR COMPUTATIONAL FLUID DYNAMICS
Industrial CFD typically relies either Reynolds-averaged Navier-Stoker (RANS) models (Boussi- nesq. 1877 Alfonsi 2009), where the fluctuations are expressed as a function of the eddy viscosity, or coarsely-resolved Large-Eddy Simulation (LES) (Smagorinsky 1963] Lesieur Mctais [996). where only the large scales are numerically simulated, and the small ones are modelled (with an a-priori physies assumption). However, there are severe limits both methods associated with their usage for general purpose. For example, RANS is too simple to model complex flows (Slot- nick et al.] 2014) and often exhibit significant errors when dealing with complex pressure-gradicn distributions and complicated geometries. LES can exhibit limited accuracy in their predictions of high-Re Re turbulent flows, due to the first-order dependence of LES on the subgrid-scale (SGS) model.
B INFRASTRUCTURES
We train and evaluate all the classic and neural Navier-Stokes solvers Nvidia Tesla V100-32G GPUs. The inference latency is measured by unrolling 2 trajectories for 25.0 simulation time on single V100 GPU.
C HIPPO FEATURES FOR TEMPORAL STENCIL MODELING
For scalar time series dgg == u(x)) the HiPPO (Gu et al. 2020) projection aims to find coefficient mapping for orthogonal polynomials such that
c(t) E R- = arg min |luge g (1(te)
(6)
where (0) = > en (1)g,()) and gn (1) is the th Legendre polynomial scaled to the [0, (] do-
main. put)) [o the uniform weight metric. By solving the corresponding ODE and its
discretization, Gu et al. (2020) showed that the optimal polynomial coefficients er can be calculated by the following recurrence:
cTT (1)++)== (1
A T/At
(T))
B u(T)
T/Lt
(2n + 1)2 (0n+1)1/2222) + 1)1/2
if
where
Ank
n 1+
if k, <
B, = (2n 1)
We refer the readers to (Gu et al. 2020) for the detailed derivations. When dealing with the multi- variate time series such as temporal PDE, we follow the original recipe and treat each scalar compo- nent independently, that is, we treat the H W T x shaped velocity trajectory H W x separate time series of length! T. Finally, we concatenate the resulted H W x 2 x C' features on the velocity dimension (i.e., H x W 2C) the input the CNN.
14

---

Under review as a conference paper at ICLR 2023
Eoee A
1127
8048x2046 (Drounderuth
2.03
1120
HPPPOST BM-6488
0PC 4M
PMT 3444
1mw 11.22
11.22
MEFROTTSM superrresolutoon
IPPPO-T
HPPO-15M-1BBIT Bapee--aach.ioon
Figure 8: Qualitative results with TSM evaluating on 32 x 32 and 16 x 16, and their 64 x 64 super-resolution results. From the super-resolution results of the T O, resolution models actually work quite well.
that the super-
D SOLVING RAVIER-STOKER WITH 32 32 AND 16 X 16 GRIDS
Since TSM shows significant improvements over vanilla neural stencil modeling the 64 x 64 grid, we are wondering whether TSM can further push the Pareto frontier by accurately approximate the convective flux in lower resolution grid.
Specifically, we still use ANS-2048>> x 2048 as the ground- truth training data, but train the TSM solver in the 32 32 and 16 16 down- down-sampled grids. In order to make fair and direction comparison to our original 64 x 64 model, we additionally train two 32 x 32 > 64 x 64 and 16 x 16 - 64 x 64 super-resolution model, and evaluate (on 64 x 64 vorticity) their super-resolved results.
We report the results in Fig. and Tab. We can see that the super-resolution models actually work quite well for 32 32 64 64 and 16 16 64 64. Besides, though the lower-resolution neural solvers' performance is significantly worse than 64 x 64, we can see that the improvement from temporal information in the trajectories is more significant in lower resolutions.
15

---

Under review as a conference paper at ICLR 2023
Table 3: Quantitative comparisons with the metric high-correlation (p 0.8) duration (W.r.t the reference DNS-2048 x 2048 trajectories). The results of TSM-32 x 32 and TSM-16 x 16 are evaluated on their corresponding 64 x 64 super-resolution results with additionally trained super- resolution models.
Method
High-corr. duration
DNS-64 x 64
2.805
DNS-128 128
3,983
DNS-256 256
5.386
DNS-512 512
6.788
DNS-1024 x 1024
8.752
LI-64 x 64
7.910
TSM-64 x 64
9.481 %)
LI-32 x 32
5,400
TSM-32 x 32
6,802 (+25.96 %)
LI-16 x
2.805
TSM-16 x
3.576 (+27.50%)
HxWxTx2
Hxw (Cxx)
HxWx
HiPPO encoding
Physics
u(ti.t [0,7]
HiPPO(u
u/ A0
HIPPO feature recurrent upcate
(Cx2(
HxWxC'
ofT A)
Physics Modal
CNN
HiiPO((
u(Z 24()
Figure 9: Illustration of using HiPPO features as the CNN inputs for neural stencil modeling. After encoding the HiPPO features for the initial velocity trajectory V Vjo.2], given model-generated new velocity, we only need an additional recurrent update step to get the HiPPO features for V(o,T+Ad (with Eq. 7 Notice that TSM, we use the fixed optimal HiPPO recurrence for temporal encoding and only the CNN part is learnable.
Table 4: The p-values in one-sample T-test for the differences between TSM and other baseline models. The differences on all 16 test trajectories are used for each significance test. We evaluate the significance for four high-correlation thresholds: 0.95, 0.9, 0.8, and 0.7.
p-value in one-sample T-test
Baseline
>>> 0.95
p>> 0.9
> 0.8
0 0.7
DNS 64x64
1.06 x 10-11
1.14 10 10
1.20
10 10-10
1.26
10 10
DNS 128x128
8.80 x 10
2.23 x 10-10
6.75
10
6.63
10-10
DNS 256x256
1.22 10 10
2.39
10 09
4.78
10 10-09 09
1.56
10
DNS 512x512
2.38 x 10
1.49 x 10
2.27
10
1.65
10
DNS 1024x
6.63 10-03
6.91
10
1.53
10 10-02 02
1.65
10 10-02
64x64 (Kochkov et al. 2021)
9.06 10-04 04
1.65
10 10-03
3.63
10-03
3.08
10-03
16

---

Under review as a conference paper at ICLR 2023
CNS-1024
DNS-512
CNS-256
20
10
40
60
80
20
40
time
time
time
DNS 128
DNS 64
DNS 32
time
time
time
LJ:52
TSM-32
50
20 40
60
20
-1.0
-0.5
velocity
Figure 10: Qualitative comparison of TSM and other baselines on ID KS equation. The solutions are down-sampled to grid in the space dimension for comparison. The dashed vertical yellow line denotes the time-step where the Pearson correlation between the model prediction and ground-truth (i.e., DNS-1024) is lower than threshold (p < 0.8).
E GENERALIZATION ON 1D CURAMOTO-SIVASHII
(KS) EQUATION
To verify the generalization ability of TSM. further evaluate it ID equations. Following previous work (Bar-Simai et al.) 2019 Stachenfeld et al., 202I), we choose Sanramoto-Sivashiee (KS) equation
representative 1D PDE that can generate unstable and chaotic dynamics. While KS-1D is not technically turbulent, it is a well-studied chaotic equation that can be used to assess the generalization ability of our models in ID cases.
Specifically, the KS equation can be written in the conservation form of
dv
ÔJ
O,
u(x = 0) = ux()))
(8)
Øt
dx
where
2
00
85
J
(9)
8x
d.
Following previous work (Bar- Sinai et al.
2019). we consider the ID KS equation with periodic boundaries. The domain size is set to L 200, and the initial condition is set to
vo(x)
Σ A sin (2Tl di)
(10)
where 10, and A, d. are sampled from the uniform distributions of [-0.5, 0.5), [- n). and (1, 2). respectively.
17

---

Under review as a conference paper at ICLR 2023
DNS-1024
DNS-5:2
50
30
20
10
20
40
60
80
20
time
time
DWS-225
UNS-128 129
DNS 04
tim
time
11.4
TSM-64
20 40
60
80
20
40
-1.0
-0.5
0.0
0.5
1.0
velacity
Figure 11: Qualitative comparison of TSM and other baselines on ID KS equation. The solutions are down-sampled to 64 grid in the space dimension for comparison. The dashed vertical yellow line denotes the time-step where the Pearson correlation between the model prediction and ground-truth (i.e., DNS-1024) is lower than threshold (p < 0.8).
Similar to the NS solution, we solve the nonlinear convection term by advecting all velocity compo- nents simultancously, using high order scheme based on Van-Leer flux limiter with a learnable interpolator, while the second and forth order diffusion are approximated using a second order cen- tral difference approximations. The Fourier spectral method is not used because of the precision issues in JAX FFT and 1FFT
use DNS-1024 as the ground-truth training data, and train TSM solver and solver in the 32 and 64 down-sampled grids. A time-step of A/ = 1.9635 10-2 and Al 9.81748 10-3 is used for and 64 grids, respectively. Following Bar-Sinai et al. (2019), the interpolation coefficients of 6-point steneil predicted by TSM and LI. For training, dynamics, each starts with different random initial conditions and simulates with 1024 resolution for 200.0 time units (after 80.0 time unit warmup). We use 16 trajectories for evaluation.
generate 1024 trajectories of fluid
When comparing with other DNS baselines, all solutions are down-sampled to 32 or 64 for com- parison. Fig. 10 and Fig. [1 show the results of TSM and LI compared with other baselines in the 32 and 64 grids, respectively. also presented in Fig. 12 We quantitative comparison of various solvers under 64-resolution is that TSM can outperforms LI with the same resolution, and outperform DNS with 4x ~ 8x resolutions.
httpss//ggtttbb.oom///o/////////
18

---

Under review as a conference paper at ICLR 2023
veloctt correlatiin
06
DNS-128
00
10--1024
L164 (Kochiev et 2021)
(ours)
100
5me
Figure 12: Quantitative comparison of TSM and other baselines on the velocity correlation of 1D KS equation. The solutions are down-sampled to 64 grid in the space dimension for comparison.
F GENERALIZATION ON 3D NAVIER-STOKER (NS) EQUATION
To verify the generalization ability of TSM, further evaluate it on 3D Navier Stokes equations. Recall that the incompressible, constant density Navier Stokes (NS) equation for fluids has a con- servation form of:
au + V (u = v/2 1an---vpp
1 vp
(11)
We train our method with 3D incompressible Navier-Stokes flow with a linear forcing = 0.05u to avoid the decaying of flows The viscosity of the fluid is set to = 6.65 x 10-4 and the density = For training, we generate 128 trajectories of fluid dynamics, each starts with different random initial conditions and simulates with 128 128 x 128 resolution for 10.0 time units (after 80.0 time unit warmup). We use trajectories for evaluation.
Following previous work (Stachenfeld et al. 2021 Takamoto et al.), the ground-truth solution is Friedrichs-Lewy (CFD) condition, have time-step A 1. 1.402497 x 10 on the 32 x 32 32 simulation grid. For TSM-32 x 32 x 32 and LI-32 x 32 x 32, most of the settings in 3D NS follows our setup for the 2D case, except that the interpolation coefficients is only calculated for 2 x 2-point stencil to avoid out-of-memory issues.
obtained by simulation on 128 x 128 x 128 grid. According to the
When comparing with other DNS baselines, all solutions are down-sampled to 32 32 32 for comparison. A quantitative comparison of various solvers under 32 x 32 x 32-resolution is presented in Fig. The qualitative results on the planes of 0, y 0, 2 are presented in Fig. [4] Fig. and Fig. [6] respectively. We can see that TSM can outperforms LI and DNS with the same resolution, but cannot beat DNS with 2x higher resolution. This is consistent with the results reported in previous work (Stachenfeld et al. 2021.
Sonfige/ineer_Forcinne
19

---

Under review as a conference paper at ICLR 2023
vortecy coresatonn
08
DNS 32X32332
DNS 6446444444
DNS: 1281122822228
(Kochkov 2021)
0.0
2.5
5.0
75
10.0
12.5
15.0
17.5
20.0
bme
Figure 13: Quantitative comparison of TSM and other baselines on the vorticity correlation (aver- aged over three directions) of 3D incompressible Navier-Stokes Stokes equation.
1.0
1.0
Figure 14: Qualitative 3D incompressible Navier-Stoke Stokes equation results of predicted vorticity fields on the plane of x = 0.
20

---

Under review as a conference paper at ICLR 2023
100
vortcity
-0.
1.0
Figure 15: Qualitative 3D incompressible Navier-Stokes equation results of predicted vorticity fields on the plane of 0.
21

---

Under review as a conference paper at ICLR 2023
10
vortccty
-00
1.0
Figure 16: Qualitative 3D incompressible Navier-Stoker equation results of predicted vorticity fields on the plane of = O.
22


Reviews:
[{"review_id": "r-8fe3TFiSc", "reviewer": "Reviewer_PBxt", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper suggests a Machine Learning (ML) driven method to support solving the Navier-Stokes equation of incompressible flows, a classical problem in Computational Fluid Dynamics (CFD). The submitted approach learns to infer the stencils of advection-diffusion fluids in a Finite Volume Method framework. \n\nThe authors introduce the mathematical setting of incompressible constant density Navier Stokes equation for fluids. They then report their ansatz by using embedded HIPPO features and learning a small Convolutional Neural Network (6 layer Unet) to estimate the spatio-temporal stencils. These stencils are used to perform a super-resolution-like task by which a reduced resolution input field (64x64) is regressed to a full-resolution solution of the aforementioned PDE solution (e.g. 2048x2048). \n\nThey then describe and show experimental results on low-resolution input data to compare their predictions against the ground truth at high resolution. The comparison is contrasted with one other ML-driven approach taken from [Kochkov et al 2021]. From these results, they infer strong indications that their approach is superior to the existing model by [Kochkov et al 2021].\n", "summary_of_the_review": "I congratulate the authors for their submission to ICLR'23. The presented paper demonstrates a novel approach to solving the Navier Stokes PDE in a highly relevant parameter domain (of incompressible flows). This work is hence very relevant for practical concerns and has the potential to push the CFD domain to new PDE or parameter space domains.\n\nThe ML-driven TSM approach under study is found to provide strong performance gains beyond one singular competing ML based approach. The evidence in favor of this can be considered weak. I encourage the authors to continue their work and add a more thorough statistical analysis of the predicted fluid fields with their setup. This may include a cross-validation study of the CNN, an estimation of the predicted estimates and a statistical analysis of the obtained confidence intervals. \n", "strength_and_weaknesses": "The paper is strong in introducing the domain of CFD to a reader from the ML community, i.e. not versed in CFD. The text is easy to follow (despite minor language issues) for an audience versed in partial differential equations (PDEs). The metrics under study are designed in an approachable manner, so that these metrics' interpretation is straight forward. The authors add with indicative visualizations (fig. 2) of the mathematical fields under study. The ML based approach is described in a condensed textual but also visual fashion (fig 1), which is helpful to understand the rational. Towards the end, the authors strive to underpin the trustworthiness of their approach by reporting an ablation study as well as tests for general applicability.\n\nA major weakness of this article is the lack of standard techniques in machine learning. A clear description of the input and output data as well as the network architecture is scattered throughout the paper. It woud be beneficial to have this under 'Experimental Setup'. Moreover, the network architecture used appears to be rather small (6 conv layers with 256 feature channels each). A clear description and comparison to the approach of LI would help.\n\nOne more major weakness along this line of thought: a cross-validation (standard technique in ML) could have been undertaken to sample the predictive performance of the network. If done with the competing approach too, this could provide a basis for a statistical treatment of the results and hence increase trustworthiness of any eventual interpretation. Related, the authors claim a \"significant\" improvement without conducting a statistical analysis of their predictions, i.e. a hypothesis test (or Bayes Factor computation) or estimation of a confidence/credible interval. This severely weakens the interpretations brought forward and raises severe doubts that the findings are significant in the narrow (statistical) sense of the term.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "mb8Sth4Frw", "reviewer": "Reviewer_EevF", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The authors proposed TSM that combines time series modeling method HiPPO and neural PDE solver. This proposed method has higher accuracy than existing models in 2-D Navier Stokes equation, and has lower inference latency. The authors generally identified these characteristics in various forcing and Reynolds number environments.", "summary_of_the_review": "Due to the following concerns, I think the paper is not at the acceptance level yet:\n\n- Did the authors compare other basic deep learning models besides CNN?\n\n- It seems necessary to compare it with other models that perform the Super-resolution task in the PDE Problem. Neural Operator models also perform super-resolution tasks. It is necessary to compare.\n\n- TSM does not seem to have significantly changed from the existing HiPPO and Neuarl PDE solver models. Can you explain this in more detail?\n\n- I wonder if using 64 x 64 grid for all learnable solvers in Table 2 is a fair comparison. How does the performance of the models change for different resolutions?\n\nHowever, I think this paper adds some value to the physics-informed machine-learning community. During the rebuttal period, I recommend they address as many issues as possible and I will reconsider my evaluation in a positive way.", "strength_and_weaknesses": "Strengths : \n\n- TSM-HiPPO encodes raw fluid velocity using the HiPPO, which operates by projecting the previous trajectories as polynomial bases. This approach can well extract features of time-varying fluid dynamic.\n\n- In addition to the ML based model, the authors also experimented with the physics-type model.\n\n- It was confirmed in the paper that TSM-HiPPO has better performance than baselines at various environment not in a single fixed environment.\n\n\nWeaknesses :\n\n- It is not possible to verify that the proposed model is available for general PDEs. In this paper, the authors verified the performance of the model only for Navier Stokes equation and did not compare it with other PDEs. Additional experiments are needed on various other PDEs that existing models are solving.\n\n- There is a lack of baseline in the neural operator field. In this paper, only Fourier Neural Operators (FNO) are compared as baseline. It is necessary to add models that appeared after FNO as baselines. Additional comparisons with wavelet transform based neural operator(MWNO), transformer based neural operator, and other models that the authors mentioned in Related work are needed.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "XHWr7ar-CYp", "reviewer": "Reviewer_41X5", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper introduces a method that combines the strengths of advanced time-series sequence modeling (with the HiPPO features) with neural PDE solvers. In a 2-D Kolmogorov flow, the method shows improved high-correlation duration, compared to baselines with learned correction or raw features.", "summary_of_the_review": "In summary, the paper introduces a method that improves upon prior SOTA neural solver with the temporal stencil modeling. My main concern is limited novelty and limited experiment evaluation of the paper, as stated in the weakness part of the paper.", "strength_and_weaknesses": "Strength:\nThe paper addresses an important problem. It is written clearly. The experiment is sound.\n\n\nWeaknesses:\n1. The use a the temporal stencil, and especially the use of HiPPO features is not well motivated. Since this is the most important part of the paper, it can be motivated better, and maybe shorten the background part. Also, it would be nice to give an intuition about the HiPPO features, even if the introduction is in the appendix.\n\n2. The main contribution is a simple combination of the time-series sequence modeling (with the HiPPO features) and SOTA neural PDE solvers, the novelty is limited.\n\n3. It would strengthen the paper, if the method is demonstrated in another (potentially more difficult, larger) dataset, to show the generality and/or scalability of the method.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "VBKPY80sKwW", "reviewer": "Reviewer_ay4D", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The work proposed a neural PDE solver with temporal stencil modeling. The neural solver incorporate the state of the art temporal HiPPO feature of low-resolution results as inputs for super resolution. The method was evaluated on 2-D Kolmogorov flows and consistently outperform a series of baselines.", "summary_of_the_review": "The idea of using HiPPO temporal features for neural PDE super-resolution solution is clear and well motivated. Comprehensive experiment results show strong and convincing performance of the proposed methods.", "strength_and_weaknesses": "Strength:\n\nThe core idea of the work is clear and simple. The motivation for using temporal feature for super-resolution PDE solution is intuitive and strong. The evaluation of the proposed method is comprehensive.I especially appreciate the generalization test experiment results and the method consistently outperforms existing approaches.\n\n\nWeakness:\n\nI do not have major concerns about the work except for two minor ones: 1) It’s clear HiPPO feature is used because it’s the SOTA temporal feature but have the authors considered other temporal features like learnable features from RNN and transformer models? 2) There’s a typo in the first paragraph of Section 4.1. Simulated Data-> Simulated Data", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject