Towards Multi-spatiotemporal-scale Generalized PDE Modeling
Authors: Jayesh K Gupta
Johannes Brandstetter
Venue: Accepted by TMLR ()

Abstract: Partial differential equations (PDEs) are central to describing complex physical system simulations. Their expensive solution techniques have led to an increased interest in deep neural network based surrogates. However, the practical utility of training such surrogates is contingent on their ability to model complex multi-scale spatio-temporal phenomena. {In recent years}, various neural network architectures have been proposed to target such phenomena, most notably Fourier Neural Operators (FNOs), which give a natural handle over local \& global spatial information via parameterization of different Fourier modes, and U-Nets which treat local and global information via downsampling and upsampling paths. However, large-scale comparisons between these convolution-based approaches are notoriously sparse. In this work, we make such comprehensive comparisons regarding performance, runtime complexity, memory requirements, and generalization capabilities. Concretely, we stress-test various FNO, (Dilated) ResNet, and U-Net like approaches to fluid mechanics problems in both vorticity-stream and velocity function form. For U-Nets, we transfer recent architectural improvements from computer vision, most notably from object segmentation and generative modeling. Next, we use our insights on design considerations, and introduce U-FNets, i.e., modern U-Nets that are augmented with FNO downsampling layers. Those architectures further improve performance without major degradation of computational cost. Finally, we ablate and discuss various choices for parameter conditioning}, and show promising results on generalization to different PDE parameters and time-scales with a single surrogate model. Source code for our PyTorch benchmark framework is available at https://anonymous.4open.science/r/tmlr-pdemulti-6677/.

Published in Transactions on Machine Learning Research (07/2023)
Towards Multi-Spatiotemprra--Scae Generalized PDE Modeling
Jayesh K. Gupta*
jky stnsss.standfor.cco edu
Microsoft Autonomous Systems and Robotics Research
Johannes Brandstetter'
johanmeshd@mmccrrooff..om
Microsoft Research AlfScience
Reviewed on OpenReview: https: // openreview. net/ forum? id- dPSTDbGtBY
Abstract
Partial differential equations (PDEs) are central to describing complex physical system sim- ulations. Their expensive solution techniques have led an increased interest in deep neural network based surrogates. However, the practical utility of training such surrogates is con- tingent on their ability to model complex multi-scale spatio-temporal phenomena. In recent years, various neural network architectures have been proposed to target such phenomena, most notably Fourier Neural Operators (FNOs), which give a natural handle over local & global spatial information via parameterization of different Fourier modes, and U-Nets which treat local and global information via downsampling and upsampling paths. How- ever, large-scale comparisons between these convolution-base approaches are notoriously sparse. In this work, we make such comprehensive comparisons regarding performance, runtime complexity, memory requirements, and generalization capabilities. Concretely, stress-test various FNO, (Dilated) ResNet, and U-Net like approaches to fluid mechanics problems in both vorticity-stream and velocily function form. For U-Nets, we transfer recent architectural improvements from computer vision, most notably from object segmentation and generative modeling, Next, we use our insights on design considerations, and intro- duce U-FNets, i.e., modern U-Nets that are augmented with FNO downsampling layers. Those architectures further improve performance without major degradation of computa- tional cost. Finally, we ablate and discuss various choices for parameter conditioning, and show promising results on generalization to different PDE parameters and time-scales with a single surrogate model. Source code for our PyTorrh benchmark framework is available at https::/ /tttps://githan //ggttt......../// httpaaa...oo/accoooooo/////reeee
1 Introduction
Many mathematical models of physical phenomena are expressed in differential equation forms (Olver, 1986), generally as temporal partial differential equations (PDEs). Their expensive solution techniques have led to an increased interest in deep neural network based surrogates (Bar- Sinai et al., 2019; Raissi et al., 2019; Lu et al., 2021; et al., 2020a; Brandstetter et al., 2022c; Um et al., 2020); especially in the studies that relate to fluid dynamics (Guo al., 2016; Kochkov et al., 2021; Rasp & Thuerey, 2021; Keisler, 2022; Weyn et al., 2020; Sonderby et al., 2020; Wang et al., 2020a; Pathak et al., 2022). However, generalizing across different PDE parameters, and different time-scales is a notoriously hard problem. For example, in fluid mechanics, slightly different values of a single parameter like Reynolds numbers can make all the difference for flows being laminar or turbulent. Another source of challenge stems from the fact that physical phenomena appear at different spatial and temporal scales. For example, blizzards are rather local weather phenomena, whereas heat waves are rather global ones, both resulting from the same underlying principles. For exactly these
* "Equal contributions.

---

Published in Transactions on Machine Learning Research (07/2023)
reasons, fluid mechanics in general (Munson et al., 2013), and weather forecasting in particular (Jolliffe & Stephenson, 2012) have always posed a great scientific challenge.
Prominent examples of neural PDE sur- rogates are Fourier Neural Operators (FNOs) (Ti et al., 2020a). At core, FNO building blocks consist of Fast Fourier transforms (FFTs) (Cooley & Tukey, 1965; Van Loan, 1992) and weight multiplication in the Fourier space, where low Fourier modes provide global information, and high Fourier modes provide local information. An FNO layer processes global and local infor- mation simultancously via weight multipli- cation of the different modes. On the other hand, U-Nets (Ronneberger et al., 2015) are standard architectures in the context of image modeling, image segmentation, and image generation. U-Nets are constructed as a spatial downsampling pass followed by a spatial upsampling pass, with additional skip connections present between the downsampling pass activations and corresponding upsampling layers. Local and global information is therefore treated in pling corresponds to scquentially processing information more globally, whereas upsampling corresponds to fine-graining the global information and adding local information via skip connections. Figure 2 contrasts local and global information flows in FNO and U-Net like architectures. Given the recent success of modern U-Net architectures in complex generative image modeling tasks (Ho et al., 2020; Nichol & Dhariwal, 2021; Ramesh et al., 2021) it's pertinent that these are evaluated on PDE Operator learning tasks and compared FNO like approaches. Furthermore, given the different nature of FNO and U-Net like approaches, is worth surveying their respective advantages and performance on different tasks, as well as investigating under which cireumstances combining them might be beneficial. The third line of models 2016a) like archilectures, which priori have natural handle processing local and global information in contrast to recently introduced Dilated ResNets (Stachenfeld et al., 2021) which adapt filter sizes at different layers via dilated convolutions.
Ground truth
Abs. eror Predietion
Figure 1: Example rollout trajectories of the best-performing U-Net model, which is trained to generalize across different timesteps (At) and different force terms.
distributed fashion than in FNO like architectures. Downsam-
ResNet (He et al.,
To summarize our contributions: (1) To our knowledge, we are the first to present a side-by-side analysis of FNO, (Dilated) ResNet, and U-Net like architectures on their ability model complex multi-scale spatio- temporal phenomena. In doing so, we present new architecture designs based on modern updates to U-Nets. (2) We generalize to different PDE parameters and time-scales showing promising results for single surrogate models as exemplified in Figure 1. (3) We propose unified PyTorc based framework for enabling easy side-by-side comparisons of various PDE operator lcarning methods which is available at https: htt///ggthh.../// com/mi crosoft t/ pdearena.
2 Preliminaries
Common parameterization of Fourier transform layers. The discrete Fourier transform (DFT) to- gether with point-wise multiplication in the Fourier space is the heart of Fourier Neural Operator (FNO) layers. DF'ss convert an n-dimensional complex signal f(x) = f(11,..., Fx)) ) TR - C at M1 x... x M, grid points into its complex Fourier modes f(E.....,E.) via:
Mi
M 2 mn +...t my En.
f(E En) = FFUEE,,......................... 6
Σ
f(x)+
m2=0
m,
2

---

Published in Transactions on Machine Learning Research (07 (022))
FNO Layer
FNO Layer
FNO Layer
Local spatial information
Local
Global
1.n A. 1.1 1400004 M
Spatial information
Global spatial information
Fourier modes
Figure 2: Information flow in Fourier based (left) and U-Net based architectures (right). FNO layers (Li et al., 2020a) consist. of Fast, Fourier transforms and weight multiplication ion in the Fourier space. Low Fourier modes provide global and high Fourier modes provide local information. U-Nets (Ronneberger et al., 2015) are constructed as a spatial downsampling pass, followed by spatial upsampling pass, where information from the downsampling pass is added via skip-connections.
where (c,, E EM, ....X... x... ZM. In FNO layers, discrete Fourier transforms on real-valued input fields and respective back transforms implemented as Fast Fourier Transforms' on real-valued inputs (RFFTs)? are interleaved with a weight multiplication by a complex weight matrix of shape Gn Cout for each mode, which results in complex-valuee valued weight tensor of the form W C% C x(E] xfgmax "...x"".") where Fourier modes above cut-off frequencies (Emmax ..... <max) are set to zero. These cut-off frequencies turn out to be important hyperparameters. Additionally, residual connection is usually implemented convolution layer (sce Figure 2).
with kernel size
Common parameterizations of convolution layers. Regular convolutional neural network (CNN) (Fukushima & Miyake, 1982; LeCun et al., 1998) layers are the basic building blocks of U-Net like architec- tures. CNNs take as input feature maps f 2 + Rea and convolve" them with a set of Cout filters {w (ea with w nn R
If w')=((x)
Σ (f(e), 2)) 9EZ"
(2)
which be interpreted inner product of input feature maps with the corresponding filters at every point y z". The filter size of convolutional layer is crucial choice in neural network design since it defines the regions from which information is obtained. Common practice is to use rather small filters (Simonyan & Zisserman, 2014; Szegedy et al., 2015; 5; He et al., 2016 Contimuous formulations of filters were introduced to handle irregularly sampled data (Schütt et al., 2018; Simonovsky & Komodakis, 2017; Wang et al., 2018; Wu. et al., 2019) and to match the resolution of the underlying data (Peng al., 2017; Cordonnier et al., 2019; Romero et al., 2021b). A promising direction to adapt filter sizes different layers using dilation (Dai al, 2017). For example, dilated convolutions in the context of PDE modeling were proposed in Stachenfeld et al. (2021). As downside, dilations might limit the bandwidth of the filters, and thus the amount of collected detail. A rather new direction is therefore to adapt filter sizes either via learnable dilation (Pintea et al., 2021) or via flexible sized continuous convolutions (Romero et al., 2021a; 2022).
Fast Fourier transforms (FFTs) immensely accelerate DFT computation by factorizing Equation into product of spanse (mostly zero) factors.
2The FFT of real-valued signal is Hermitian-smmmetric, the output contains only the positive frequencies below the Nyquist frequency for the spatial dimension.
In deep learning, convolution operation in the forward pass is implemented as cross-correlation.
3

---

Published in Transactions on Machine Learning Research (07 (07/2023)
sst block
2nd block
3rd block
4th block
-10
-10-2
-10
-10
y-modes
20
40
20
10
60
30
15
50
100
20
40
60
10
20
30
0 5
10 15
x-modes
Figure 3: Analyzing filter properties of trained U-Net architectures. Absolute values of Fourier modes of the filters in each first. layer of the respective down-sampling blocks are shown, where for each mode the average is taken over all filters.
Connecting Fourier transform and convolution. Starting with the 1-dimensional case and omitting channel dimensions, assume signal consisting of imput points, and we further assume circular padding, We can now rewrite Equation 2 into discrete circular convolution (Bamich, 2018; Bronstein et al., 2021) of two n-dimensional vectors f, w E R":
W1
W2
Wp-1
Wp
n-1 W(ish) mod nfz- (Cw))sss C =
W0
W1
W2
Wn-1
[f * w), = Σ j=0
3
U2
Wp-1
Wg
W1
The indexing (i j) mod n returns circular shifts, which can be combined into a circulant matrix Cw- is general practice use rather small filters which only consist of & non-zero clements where usually << n. The remaining elements of C are filled up with zeros. The action of C.', on f. or equivalently the convolution of f with can be expressed via the convolution theorem:
Cw - (
(
(4)
VT
where the matrix W consists of the eigenvectors of C. and W is its complex conjugate. All circulant matrices have the same eigenvectors, which if multiplied with signal yields the discrete Fourier transform (DFT) of the signal. That is, multiplication (from left) with W is the discrete Fourier transform (of f); and multiplication by W is the inverse Fourier transform. The matrix D has the Fourier modes of the vector w its diagonal. Thus, we can analyze convolution by expressing its filters vectors w E R", which comprise the actual k filter values (A: corresponds to the kernel size) and additional n - k zeros. When we extend the circular convolution approach to two dimensions, w itself becomes a matrix w E R"**, Figure 3, we plot the Fourier modes of the two dimensional filters a trained U-Net. We take the absolute values of the modes and average for each mode over all filters the first convolution layer of different down- sampling blocks. Although precise statements are difficult to make, it evident that Fourier mode averages of different blocks are downsampled versions of each other, which complies with the interpretation that the downsampling blocks of U-Nets process information at different scales. This is therefore in contrast. to FNO like architectures which process different scales within each FNO layer.
Fourier transform for downsampling. Bandlimited pre-subsampling (Mallat, 1999), i..., suppressing high-frequencies before down-sampling, a well know technique in signal processing; for an illustrative example see c.g., Figure in Worrall & Welling (2019). We hypothesize that replacing convolutions with FNO layers which set Fourier modes above cut-off frequencies to zero is advantageous, especially in the lower downsampling blocks of U-Net architectures where the spatial resolution is large. In the lower U-Net blocks, convolutions might be all what is needed to learn efficient coarso-grained representations.
Tested partial differential equations. A partial differential equation (PDE) relates solutions u X + R" and respective derivatives for all points æ in the domain X' E R" where u²(z) are initial conditions at time 0 and [tu] æ) 0 are boundary conditions with boundary operator B when x is on the boundary ax of the domain. In this work, we investigate PDEs of fluid mechanics problems. To be more precise, we focus on the incompressible Navier-Stokes equations (Temam, 2001), and the shallow water
4

---

Published in Transactions on Machine Learning Research (07/2023)
equations(Vrengge (rrugdennhi,,,, 1994) as important variant thereof. For the latter, we test for data obtained in velocity function and vorticity stream formulation. Such comparisons, although never done for deep learning PDE surrogates, are important since the characteristies of these two formulations are different. While vorticity stream formulation is often casier to solve, it has limited applicability in complex and realistic flow scenarios, especially when dealing with complex geometries and resultant boundary conditions. Comparatively, velocity function formulation more intuitive, directly dealing with physical measurable quantities, i.e., velocity fields, with easier implementation of boundary conditions (Gatski, 1991; Kundu et al., 2015).
In 2 dimensions, the Navier-Stokes equations in vector velocity form conserve the velocity flow fields v : X' + I2 where X' E IB² via:
dv v VU V V v==00
(5)
Øt
where v Vo is the convection, i.e. the rate of change of along v, / V the viscosity, i.e. the diffusion or net movement of Vp the internal pressure and f an external force. An additional incompressiiility constraint V v = 0 yields mass conservation of the Navier-Stokes equations.
E' By introducing the vorticity w : E R as the curl of the flow velocity, ie. w V x v, we can rewrite the incompressible 2-dimensional Navier-Stokes equations scalar vorticity stream function form (Kundu et al., 2015; Guyon et al., 2001; Acheson, 1991) as:
Ow Oip ow
bb ow ØL dy dx dx Oy +
( (ou pa
pv p.b 2x2 +
+
Re Or dx2
Ox2
= -
(6)
where the streamfunction is defined via the relations Pe which is indirectly proportional to the viscosity and proportional to the absolute velocity. As a result, the 2D incompressible Navier-Stokes equations are turned into one parabolic equation, i.e. the vorticity transport equation (Equation left), and one elliptic equation, i.e. the Poisson equation (Equation right). Since the streamfunction is directly obtained from the vorticity via the Poisson equation one usually solves for the scalar vorticity.
By v2 and By = uy and Re is the Reynolds number
dz
The shallow water equations (Vreugdenhil, 1994) can be derived from integrating the incompressible Navier Stokes equations, in cases where the horizontal length scale is much larger than the vertical length scale. As such, shallow water equations describe a thin layer of fluid of constant density in hydrostatic balance, bounded from below by the boltom lopography and from above by a free surface. For simplified weather modeling, the shallow water equations express the velocity in x- direction termed zonal velocity, the velocity in the y- direction termed meridional velocity, and the vertical displacement of free surface, which subsequently is used to derive pressure fields. Since the shallow water equations derived from the Navier-Stokes equations, vorticity-stream function formulation exists well. Note however that when it comes to describing flows in e.g. more complex geometries, the velocity formulation is in general easier to deal with (Kundu et al., 2015).
3 PDE Surrogates
We implement ResNet architectures using 8 residual blocks, where each block consists of two convolution layers with 3 kernels, shortout connections, group normalization (Wu & He, 2018), and GeLU activation functions (Hendrycks & Gimpel, 2016). In contrast standard ResNets for image classification, we don't use any down-projection techniques, e.g. convolution layers with strides larger than 1 or via pooling layers. In doing so, ResNets have no natural built-in handle over local and global informations, and therefore serve as important bascline to ablate effects of local and global information flow which is fundamental in e.g. FNO and U-Net. like archilectures. Recently, Stachenfeld et al. (2021) introduced Dilated ResNets, which adapt filter sizes at different layers using dilated convolutions, and thus are an alternative way of subsequently aggregating global information. The models consist of 4 residual blocks where each block individually consists of 7 dilated CNN layers with dilation rates of [1, 2, 4, 8, 4,2, 1]. We implement Dilated ResNets with and without group normalization layers.
(Dilated) ResNets.
5

---

Published in Transactions on Machine Learning Research (07/2023)
Fourier Neural Operators. We implement FNO architectures where the mumber of FNO layers, the number of channels, and the number of non-zero Fourier modes are hyperparameters. All architectures consist of two embedding and two output layers. Each FNO layer comprises a convolution path with a I x 1 kernel and a Fourier path where pointwise weight multiplication is done for the lower modes in the Fourier domain. We use GeLU activation functions, and no normalization scheme.
U-Nets. U-Nets have already been used as PDE surrogates in Ma et al. (2021); Chen & Thuerey (2021). U-Nets are constructed as spatial downsampling followed by a spatial upsampling pass, where each down- and upsampling block consists of two convolutional layers. A particularity of U-Nets is the presence of skip connections between the downsampling pass activations and corresponding upsampling layers. Orignially, downsampling is achieved via max-pooling operations. We term the 2015 U-Net implementation as U- Net2o1s, which is based on the PDEbench repository of Takamoto et al. (2022). Furthermore, we include a slightly different version which we term U-Netbase which has bias weights and group normalization instead of batch normalization be comparable with modern U-Net versions. To match the number of weights of U-Nel2015; the bottleneck layer in h- is omitted. Modern versions of the architecture (Ho et al., 2020; Nichol & Dhariwal, 2021; Ramesh et al., 2021) often use Wide ResNet (Zagoruyko & Komodakis, 2016) style 2D convolutional blocks, cach of which can be followed by a spatial attention block (Vaswani et al., 2017). Other notable changes are the substitution of max-pooling operations by downsampling layers. We term the respective implementations U-Netmod and U-Netaut in our experiments.
Fourier U-Nets. Based on the insights of Section 2, we replace lower blocks both in the downsampling and in the upsampling path of U-Net architectures by Fourier blocks, where each block consists of 2 FNO layers and residual connections. We test substituting only the lowest block (U-FINet), and the lowest two blocks (U-F2Net) of the U-Netmod architecture. Substituting all blocks would yield architecture which resembles the UNO architecture (Rahman et al., 2022b), with the difference that in UNO architectures downsampling is done individually via linear layers along the x- and y- dimension, and that "mode scheduling" reduces the number of modes for higher blocks in the respective downsampling and upsampling paths. For complete comparison, we therefore also implement the UNO architecture".
3.1 Operator learning
Major practical benefits of neural PDE surrogates come from amortizing the cost of their compute-expensive training process which depends on the surrogates' ability to effectively generalize across different parameter settings as well as across different time discretizations. Operator learning is a popular term for training these neural surrogates. Theoretical grounding arises from Chen & Chen (1995) who extend the universal approximation theorem in neural networks (Hornik et al., 1989; Cybenko, 1989) to operator approximation, forming the basis for DeepONets (Lu al., 2019) with theoretical extensions in Lu et al. (2021), graph kernel networks (Li et. al., 2020b), and FNOs. An impressive comparison of DeepONets and FNOs can be found in Lu et al. (2022).
Operator learning (Lu et al., 2019; Li et al., 2020b;a; Lu et al., 2021; 2022) relates solutions u X + R". u' : X" - RR" defined on different domains X' E Rmm X" E RRm" via operators G:
G (u (u + (u' EU")
where U and U' are the spaces of solutions u and u', respectively.
Parameter conditioning. We evaluate ENO and U-Net like architectures on their generalization ca- pabilitics across PDE parameters and different time-scales. Data points consist of solution pairs u, u' e u where the pair itself is from the same solution space U, but different pairs {u, u'}1 and (u, are from different solution spaces Ui and U2 characterized by different PDE. parameters, e.g., force terms. Simply put, an input-output pair is always drawn from one trajectory and therefore naturally shares the same PDE parameters. Further, the mapping u + u' should generalize across different time windows At, which is taken into account. by conditioning on the lead time, i.e., the time the neural surrogate is asked to predict into the future. Taking all these roquirements into account, we train neural surrogates to generalize across different
We based our implementation on https://gttaab..ccm/snisss/a///
com

---

Published in Transactions on Machine Learning Research (07/2023)
initial conditions, different PDE parameters (force terms) and different time windows. Both time windows At and force terms are continuous scalar parameters, and thus can be encoded into vector representation by using simusoidal embeddings as is common in Transformers (Vaswani et al., 2017) and various neural implicit representation learning techniques (Mildenhall et al., 2021).
4 Experiments
We establish the following set of desiderata for our benchmarks: (i) simplicity: the tasks should be easy to setup, while being backed by actual PDE solvers written by domain experts, (ii) challenging: the tasks should be diffieult. enough, (iii) diverse: the tasks should be diverse, both in their formulation well as in their requirements, and (iv) gemeralizability: the tasks should probe generalization across different time horizons well as different parameter settings. Following these desiderata, we assessed the described architectures in four experimental settings to probe (i) Fourier vs. U-Net based approaches, (ii) differences due to the voriticty-stream vs. velocity function formulation of the datasets, and (iii) parameter conditioning performance. Results of the main paper are complemented by comprehensive studies and various ablations in Appendix B.
Concretely, we first probe selected architectures the shallow water equations, both in vorticity-stream and velocity function form. We do second comprehensive evaluation on the Navier-Stokes equations in velocity function formulation since this formulation turns out to be harder lo model for neural surrogates. Finally, we probe parameter conditioning choices on the Navier-Stokes equations. We discuss results and differences between architectures.
All datasets contained multiple input and output fields. More precisely, one scalar and one velocity vector field in case of the velocity formulation, and two scalars in case of the vorticity formulation. Inputs to the neural PDE surrogates were respective fields at previous timesteps, where t varies for different PDEs. The one-slep loss is the mean-squared error al the next limestep summnd over fields. The rollout loss (reported in Appendix B) is the mean-squared error after applying the neural PDE surrogate 5 times, summing over fields and time dimension. We alternatively test the relative MSE loss as used in Li et al. (2020a). We optimized models using the AdamW optimizer (Kingma & Ba, 2014; Loshchilov & Hutter, 2019) for 50 epochs and minimized the summed mean squared error. We used cosine annealing as learning rate scheduler (Loshchilov & Hutter, 2016) with a lincar warmup. Table 1 compares parameter count, runtime and memory requirement of the tested architectures, showing that runtime and memory requirements are in the same ballpark for both architecture families if the number of parameters is kept similar.
Shallow water equations. We modified the implementation in SpeedyWeather. -1 (Klöwrr et al., 2022), obtaining data on a grid with spatial resolution of 192 x 96 (Ar = 1.875 Ay = 3.75%), and temporal resolution of At = h8 We first evaluated the different architectures on the shallow water equations in velocity function formulation, prodicting scalar pressure field and vector wind velocity field. Figure 4 (Ieft) shows results obtained by various models. In general, all methods which have a dedicated local and global information flow, Le. Dilated ResNet, FNO, and U-Net architectures, perform rather well. heless, across all tested models, performance differences of an order of magnitude arise, where U-Nets in general perform best. Adding FNO blocks to U-Net architectures (U-FINet, U-F2Net) seems to be beneficial. We further evaluate on the shallow water equations in vorticity stream formulation, and predict the scalar pressure field and the scalar wind vorticity field. Figure 4 (middle) shows results obtained by various models. Performance-wiiet a similar pattern arises, where again the lowest losses are observed for U-Net architectures.
Velocity function formulation of Navier-Stokes equations. We further tested on Navier-Stoker equations in velocity function form, which is more common in the real world than the vorticity stream function form. In addition to the velocity field v of Equation 5, we introduced scalar field representing a scalar quantity, i.e particle concentration, that is being transported via the velocity field. The scalar field is advected by the vector field, i.e. as the vector field changes, the scalar field is transported along with it. Complementary, the scalar field influences the vector field only via an external buoyancy force term in y-direction, i.e. f = (0, f)". We obtained data on a grid with spatial resolution of 128 x 128 (Ar = 0.25,
'nttps://githubb wwtp:/hthu..cm/mmaaiiiieeeeeeeee www.m///waanlll/mmemmmmmm htmm///annaaaatttteetaaaaaaeeer

---

Published in Transactions on Machine Learning Research (07 (07/2023)
Table 1: Comparison of parameter count, runtime, and memory requirement of various architectures. Sub- script numbers indicate the used number of Fourier modes. For U-FNet experiments subscript numbers indicate the number of Fourier modes in the lowest and second-lower block.
Runtine [s]
Mem. [MB]
MEMoOD
Channels
Res. Res.Layers/ Blocks Blocks
Params.
Fwd.
bwd.
Peak usage
ResNei128
128
2.4M
0.084
0.180
4273
ResNet256
256
9.6
0.231
0.497
38
8500
DelResNet.125
128
4.2
0.1
0.342
16
4849
DiRReNNet128-oom
128
4.2 M
0.183
0.423
16
6922
FNO 128-8
FNO128-8 PNO128-8mcdestt
128
134 M
0.059
0.171
537
2953
FNQ FNO128-44mca 28-4medeet6
128
67.2
0.031
0.089
268
1852
11000--mmdded22
64
07.1
0.016
0.050
268
1204
FNO96-4mades32 1modec32
FNO96-144
96
151 M
0.026
0.080
601
2179
FNO128- 400128-4mode332
128
268
UNO64
64
110 M
0.070
0.134
440
1925
UNO128
440
0.341
5513
U4Netgo1:64
64
31
0.013
0.037
124
1305
U-Nel2018
128
496
U-Nefpae 64
64
31.1 M
0.021
0.046
124
1277
U-Nefbwee
124
0.056
496
3000
U-Netmod 64
61
144 M
0.079
0.184
577
3900
U-Net 64
64
148
0.081
0.190
393
3975
U-FI D-FeNeemeeden
64
154 M
0.083
0. 0.205
617
3936
U-FIN eomddeet6
185
U-F2Net F.N2Netpddss..
64
163 M
0.085
0.213
652
3961
61
193
0.085
0,216
772
1016
U-P2N m
64
224
0.086
0.219
897
4149
1- m-12o11
64
344
0.090
0.232
1400
4496
Shallow water vorticity
Shallow water velocity
Navier-Stoke velocity
10
10
Re5Nes256
DilResNet 12-nnrm
10
PNO122 (601128-0 1801128--edettB
PNO96
MSE
UNO66
10
U-Net 64
10-2
10
U-Netaoo
10 0
U-Netar
U- PiNeeeeddeet
10-3
U- U-F2Nemmode115.
10
10-4
U-F2Net
5600
5600
2080
5200
Num. Train Trajectories
Figure 4: One-step errors for modeling different PDEs, shown for different number of training trajectories. Results are averaged over three different random seeds, and are obtained for the velocity function and vorticity stream formulation of the shallow water equations on 2-day prediction (left, middle), and for the Navier- Stokes equation (right). For better visibility only selected architectures are displayed, for full comparisons see Appendix B. Note the logarithmic scale of the y-axes.
Ay 0.25), and temporal resolution of At = 1.5s using @ (Holl et al., 2020). Figure 4 (right) shows results obtained by different architectures. In contrast to the shallow water experiments, the compute- expensive Dilated ResNet architectures perform on par with the best U-Net and U-FNet architectures. The reason for the performance difference of Dilated ResNets is hard pin down. Our primary guess is on the importance of aliasing vs. patterns in the data. Concretely, for data with different proportions of higher frequency modes, the importance of aliasing might shift.
ohttpe://github. http://ttthtuuuccccuuoeeu wwt.t/omtt/eeeeeeeeeeee http:/tttw..cccmmmmmmpp///PPPPPa
8

---

Published in Transactions on Machine Learning Research (07 2202)
One-step (AI 0.375s)
One-step (Ar = 1.5s)
One-step (Ar = 3s)
10
10
FNO12S- 29-Bere
U-d 64
U-Netan
10"2
U-F Net
MSE
U-F UI-P2Nttanddette
T-F NNeaat
10
U-F F.F2
10m
10
10
1604
1064
6656
Num.
Train Trajectories
One-step errors obtained on the parameter conditioning experiments of the Navier-Stokes equation. Results are shown for selected architectures, different number of training trajectories, and different time windows: At 0.375 (left), At 1.5s (middle), and At = 3s (right). Results are averaged over 208 different unseen evaluation buoyancy force values between 0.2 and 0.5.
Figure 5:
Probing parameter conditioning. We probe parameter conditioning the velocity function formu- lation of the Navicr-Stokes equation. We test FNO and U-(F)Net variants, experiments for Dilated ResNet are loo compute-expensive due to their long runtimes, see Table 1. For training, we used a dalaset with higher temporal resolution of Át = 0.375 and get equal number of trajectories from uniformly sampling 832 different external buoyancy force values, f = f) in Equation 5, in the range 0.2 f S 0.5, using input fields at one timestep.
We conditioned our models to predict for different time windows in the range 0.375 At < 20 and different strengths of the y-component of the external buoyancy force f. Due to the unbalanced nature of the dataset size at different At, we reweighed the sampling frequency in our dataloader to try to maintain parity. We provided conditioning information in the form of an embedding vector which can be added to each or subset of residual blocks (Ho et al., 2020). Both, At and f, are continuous valued scalar parameters, and thus can be encoded into a vector representation by using sinusoidal embeddings as common Transformers (Vaswani et al., 2017). We added the embedding vector the feature maps after the first convolution/FNo layer in respective down- and up-sampling blocks. To be more precise, for each feature map replicated the respective embedding value along z- and y-coordinates. For Fourier layers, this results in adding the embedding vector, the Fourier branch, and the residual connection together. We also compare an alternative conditioning approach for U-Nets termed AdaGN (Nichol & Dhariwal, 2021), based on affine transformation of group normalization layers via projections of our embeddings in Figure 19 of Appondix B.5.
Figure 5 shows results obtained by various mod- els averaged over 208 different unseen evalua- tion force values. U-Net based methods perform best. In contrast to the unconditioned experi- ments, substituting lower blocks by FNO blocks didn't yield better generalization capabilities. In general, we observe that conditioning is more dif- ficult for FNO layers, most strikingly seen in the performance curves of FNOs. We however do not discard the possibility that for FNO layers, al- ternative parameter embedding and conditioning methods might be required. Nevertheless, our re- sults also coincide with the findings of Lu el al. (2022), which state that FNO like architectures secmed to be extremely sensitive to noise, and failed to predict solutions for even small amounts of added Gaussian noise.
10
FNO128 Smodee16.16
U-Net
U-Fi L.FiNetmdd
MSE
intrsqristion)
Pereapulaiion
ineppoation
10-3
0.1
0.2
0.3
0.4
0.5
0.6
Force Magnitude
Figure 6: Inter- and extrapolation performance of differ- ent models tested on buoyancy force values in the range 0.1 S f S 0.6 performing 5 steps rollout at At = 0.375
9

---

Published in Transactions on Machine Learning Research (07 (07/202) 2023)
A 0.375s
Ar. 0.75
1:55
(a) Buoyancy force= 0.22
(b) Buoyancy force= 0.48
Figure 7: Scalar field predictions obtained for the paramcter conditioning experiments using the best per- forming U-Nebmod model. Predicted and ground truth fields are shown for different buoyancy force values and different time windows. Model imputs are the same for different time window prediction tasks. More conditioning experiments can be found in Appendix B.5.
In Figure 6, we show performance of different models tested on buoyancy force values in the range 0.1 S f < 0.6. The curves indicate that the difficulty of the tasks increases for larger buoyancy force values, but U-Net based PDE surrogates show better interpolation and extrapolation abilitios. In Figure 7, we display example scalar fields obtained for the parameter conditioning experiments of the Navier-Stokes equations using the best performing U-Netmod model. Predicted and ground truth fields are shown for different values of the absolute buoyancy force and different time-scale values.
5 Conclusion
We presented a comprehensive comparison between various (Dilated) RosNet, FNO, and U-Net based ap- proaches on fluid mechanic problems, paving basis towards strong baselines the development of neural PDE surrogates. For U-Nets, we transferred recent architectural improvements from computer vision, most notably from object segmentation and generative modeling. We found that the original U-Net architecture of Ronneberger et al. (2015) already functions powerful neural PDE surrogate, and e.g. outperforms FNOs on the presented tasks, which challenges prevailing claims in the literature, e.g. Li et al. (2020a). Combined with recent architectural improvements, we achieved significant performance improvements of C-Net based architectures. Secondly, we reported the first side by side comparison of neural PDE surrogates on vorticity-strean vs. velocity function formulation fluid dynamics problems. Overall, velocity function formulations seem be more challenging to model. Next, we used our insights on design considerations, and introduced U-FNets, i.e., modern U-Nets that augmented with FNO downsampling layers. Those archi- tectures further improve performance without major degradation of computational cost. As such, U-FNets are appealing and well-motivated alternative in the vast design space of neural PDE surrogates containing the best of both worlds. Finally, we tested parameter conditioning techniques to train generalizable neural PDE surrogates. In doing so, we were able to show promising U-Net and FNO based architectures which generalize to different PDE parameters as well as different time-scales within a single surrogate model. How- ever, similar the findings of Lu et al. (2022) FNO layers seem have negative effects when generalizing different time-scales and PDE parameters. Finally, we hope that our codebase can be a starting point for further investigations neural PDE surrogates.
10

---

Published in Transactions on Machine Learning Research (07/2023)
Limitations & Future Work. This work focuses the "image-to-image" modeling aspect of PDE sur- rogate modeling, more precisely on the understanding of complex multi-scale spatio-temporal phenomena. That said, in this work we did not elaborate on important aspects of neural PDE surrogates such as stability over long rollouts, preservation of invariants, or generalization over sampling regularities, over domain topolo- gies and geometries, and over boundary conditions. We see many of these aspects as future work. Moreover, in this work we focused on modeling Navier-Stokes equations directly, rather then in the Reynolds-averaged Navier Stokes (RANS) form (Tennekes et al., 1972), which is very common when describing turbulent flows. Finally, this work focuses comprehensive comparison on various convolution based methods, whereas future work could extend towards Vision Transformers (Dosovitskíy et al., 2020), comparing their abilities to model and generalize across spatio-temporal information, especially with respect to data availability,
Acknowledgements
We thank Alok Singh, Ratnesh Madaan, and Shuhang Chen for their comments early versions of this paper. We are also thankful to Zongyi Li for suggesting extra settings for better comparisons of our FNO baselines.
References
David J Acheson. Elementary fluid dynamics, 1991.
Troy Arcomano, Istvan Szunyogh, Jaideep Pathak, Alexander Wikner, Brian R Hunt, and Edward Ott. A machine learning-based global atmospheric forecast model. Geophysical Research Letters, 47(9): e2020GL087776, 2020.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXio preprint arXiv:1607.064 2016.
Bassam Bamich. Discovering transforms: tutorial on circulant matricos, circular convolution, and the discrete lourier transform. arXiv preprint wrXiv:885500522 2018
Yohai Bar-Sinai, Stephan Hoyer, Jason Hickey, and Michael P Brenner. Learning data-driven discretizations for partial differential equations. Proceedings of the National Academy of Sciences, 116(31):15344-31)) 153499 2019.
Saakaar Bhatnagar, Yaser Afshar, Shaowu Pan, Karthik Duraisamy, and Shailendra Kaushik. Prediction of acrodynamic flow fields using convolutional neural networks. Computational Mechanics, (2(2):525-5445 2019.
Kaifeng Bi, Lingxi Xie, Hengheng Zhang, Xin Chen, Xiaotao Gu, and Qi Tian. Pangu-weather: A 3d high-resolution model for fast and accurate global weather forecast. arXiv preprint arXiv:2211.02556, 2022.
Johannes Brandstetter, Rianne van den Berg, Max Welling, and Jayesh K. Gupta. Clifford neural layers for PDE modeling. arXiv preprint arXix:2209.0420 2022a.
Johannes Brandstetter, Max Welling, and Daniel E Worrall. Lie point symmetry data augmentation for neural pde solvers. arXiv preprint 2202.00007744 2022b.
Johannes Brandstetter, Daniel Worrall, and Max Welling. Message passing neural pde solvers. arXiv preprint wwXih:2202.03376 2022c.
Michael M Bronstein, Joan Bruna, Taco Cohen, and Petar Veličković. Geometric deep learning: Grids, groups, graphs, geodesics, and gauges. Xiviv preprint orXix:2144. 13478, 2021.
Gengxiang Chen, Yingguang Li, Qinglu Meng, Jing Zhou, Xiaozhong Hao, et al. Residual fourier neural operator for thermochemical curing of composites. Xixio preprint 1rXi:22111.10662 2021.
11

---

Published in Transactions on Machine Learning Research (07/2023)
Wi- Chen and Nils Thuerey. Towards high-aceuracy deep learning inference of compressible turbulent flows over aerofoils. arviv preprinl arXiv:2109 090.009.0....00 2021.
Tianping Chen and Hong Chen. Universal approximation to nonlinear operators by neural networks with arbitrary activation functions and its application to dynamical systems. IEEE Transactions on Neural Networks, 6(4):911 917, 1995.
James W Cooley and John W Tukey, An algorithm for the machine calculation of complex fourier series. Mathematics of computation, 19(90):297-300 301, 1965.
Jean-Baptiste Cordonnier, Andreas Loukas, and Martin Jaggi. On the relationship between self-attention and convolutional layers. arXiv preprint 1et::191100584, 2019.
George Cybenko. Approximation by superpositions of a sigmoidal function. Mathematics of control, signals and systems, 2(4))308-3L3 314, 1989.
Jifeng Dai, Haozhi Qi, Yuwen Xiong, Yi Li, Guodong Zhang, Han Hu, and Yichen Wei. Deformable convo- lutional networks. In IEEE International Conference on Computer Vision (ICCV), pp. 764 773, 2017.
Alexcy Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Un- terthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, al. An image is worth 16x16 words: Transformers for image recognition scale. arXiv preprint arxin:2010. 11929, 2020.
Thomas Frerix, Dmitrii Kochkov, Jamie Smith, Daniel Cremers, Michael Bremner, and Stephan Hoyer. Variational data assimilation with a learned inverse observation operator. In International Conference on Muchine Learning, pp. 3449 3458. PMLR, 2021.
Kunihiko Fukushima and Sei Miyake. Neocognitron: A self-organizing neural network model for a mechanism of visual pattern recognition. In Competition and cooperation in neural nets, pp. 267-285. Springer, 1982.
Victor Garcia Satorras, Zeynep Akata, and Max Welling, Combining generative and discriminative models for hybrid inference. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché Buc, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems 32, pp. 13802 13812. Curran Associates, Inc., 2019.
Thomas B Gatski. Review of incompressible fluid flow computations using the vorticity-velocity formulation. Applied Numerical Mathematics, 7(3):227-239, 1991.
Nicholas Geneva and Nicholas Zabaras. Modeling the dynamics of pde systems with physice-consttainee deep regtrssvgeee networks. Journal of Computational Physics, 403:109056, 2020.
Daniel Greenfeld, Meirav Galun, Ronen Basri, Irad Yavnch, and Ron Kimmel. Learning to optimize multigrid PDE solvers. In Proceedings of the 36th International Conference on Machine Learning, 1CML 2019, 9-15 June 2019, Long Beach, Caléfornia, USA, pp. 1415-2423, 2019.
Steven Guan, Ko-Tsung Hsu, and Parag Chitnis. Fourier neural operator net works::: A fast and general solver for the photoacoustic wave equation. arXiv preprint arXiv:2108. 09874, 2021.
John Guibas, Morteza Mardani, Zongyi Li, Andrew Tao, Anima Anandkumar, and Bryan Catanzaro. Adap- tive fourier neural operators: Efficient token mixers for transformers. arXiv preprint ar arXiv:2111.13587, 2021.
Xiaoxiao Guo, Wei Li, and Francesco Torio. Convolutional neural norks for steady flow approximation. Proceedings of the 22nd ACM SIGKDD international conference on knowledge discovery and data mining, pp. 481 490, 2016.
Etienne Guyon, Jean- Pierre Hulin, Luc Petit, Catalin Mitescu, al. Physical hydrodgnamics. Oxford university press, 2001.
Jean-Pierre
12

---

Published in Transactions on Machine Learning Research (07/2023)
Jiequn Han, Armulf Jentzen, and Weinan E. Solving high-dimensional partial differential equations using deep learning. Proceedings of the National Academy of Sciences, 115(34):8505- 1010,,51001, 2018.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770 778, 2016a.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Identity mappings in deep residual networks. In European Conference on Computer Vision (ECCV), pp. 630-645. Springer, 2016b.
Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (GELUs). arXiv preprint arXi:: 084468666.88448 2016.
Jonathan Ho, Ajay Jain, and Picter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:6840- 6851, 2020.
Philipp Holl, Vladlen Koltun, Kiwon Um, and Nils Thuerey, phillow: A differentiable pde solving framework for deep learning via physical simulations. In NewrIPS Workshop, volume 2, 2020.
Kurt Hornik, Maxwell Stinchcombe, and Halbert White. Multilayer feedforward networks are universal approximators. Newral Networks, 2(5):359-33) 1989.
Jun-Ting Hsieh, Shengjia Zhao, Stephan Fismann, Lucia Mirabella, and Stefano Ermon. Learning neural PDE solvers with convergence guarantees. arXiv preprint arXiv:1906.0 01200, 2019.
Rakhoon Hwang, Jae Yong Lee, Jin Young Shin, and Hyung Ju Hwang. Solving pde-constrained control problems using operator learning. In AAAI Conference on Artificial Intelligence, volume 36, pp. 4504 4512, 2022.
Sergey loffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In International Conference Machine Learning (ICML), pp- 448-456. PML., 2015.
Xiaowei Jin, Shengze Cai, Hui Li, and George Em Karniadakis. Nsfnets (navier-stokes flow nets): Physics- informed neural works for the incompressible navier-stokes equations. Journal of Compulational Physics, 426:109951, 2021.
lan T Jolliffe and David B Stephenson. Forecast verification: a practitioner's guide in atmospheric science. John Wiley & Sons, 2012.
Ryan Keisler. Forecasting global weather with graph neural networks. arXiv preprint arXin:2202.01575, 2022.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXio preprint ar 6980, 2014.
Milan Klöwer, Tom Kimpson, Alistair White, and Mosè Giordano. millakk/SSeeds-Westherr/ v0.2.1, July 2022. URL https://doi. org/10 10. 5281/zenodo. 6788067.
Dmitrii Kochkov, Jamie A Smith, Ayya Alieva, Qing Wang, Michael P Brenner, and Stephan Hoyer. Machine learning accelerated computational fluid dynamics. Proceedings of the National Academy of Sciences, 118 Saaming--ccelerateedd
(21):e2101784118, 2021.
Pijush K Kundu, Ira M Cohen, and David R Dowling. Fluid mechanics. Academic press, 2015.
Remi Lam, Alvaro Sanchez- Gonzzlez Matthew Willson, Peter Wirnsberger, Meire Fortunato, Alexander Pritzel, Suman Ravuri, Timo Ewalds, Ferran Alet, Zach Eaton-Rosen, et al. Graphcast: Learning skillful medium-range global weather forecasting. arXiv preprint arXiv:2212.1279; 2022.
Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient- learning applied docu- ment. recognition. Proceedings of the IEFE, 86(11):2278 2324, 1998.
13

---

Published in Transactions on Machine Learning Research (07/2023)
Zijie Li, Kazem Mcidani, and Amir Barati Farimani. Transformer for partial differential equations' operator learning. arxiv preprint arXiv:2205.1380 13671, 2022a.
Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Fourier neural operator for parametric partial differential equations. arXiv preprint Arx:::2010. 08895, 2020a.
Zongyi Li, Nikola Kovachki, Kamyar Azizzadencsheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Neural operator: Graph kernel network for partial differential equations. arXiv preprint arXiv:2003.03. 03485, 2020b.
Zongyi Li, Nikola Kovachki, Kamyar Arizadeneshelli Burigede Liu, Kaushik Bhattacharya, Andrew Stu- art, and Anima Anandkumar. Markov neural operators for learning chaotic systems. arXiv preprint arXiv:2106.0688 2021a.
Zongyi Li, Hongkai Zheng, Nikola Kovachki, David Jin, Haoxuan Chen, Burigede Liu, Kamyar Azizzade- nesheli, and Anima Anandkumar. Plysics-informed neural operator for learning partial differential equa- tions. Xiv preprint or orXiv:211100370 2021b.
Zongyi Li, Daniel Zhengyu Huang, Burigede Liu, and Anima Anandkumar. Fourier neural operator with learned deformations for pdes on general geometries. arXie preprint arXiv:2207.05200 2022b.
Marten Lienen and Stephan Günnemann. Learning the dynamics of physical systems from sparse observations with finite element networks. arXiv preprint (tl::203308852. 2022.
Burigede Liu, Nikola Kovachki, Zongyi Li, Kamyar Azizzadenesheli, Anima Anandkumar, Andrew M Stuart, and Kaushik Bhattacharya. learning-bass multiscale method and its application to inelastic impact problems. Journal of the Mechanics and Physics of Solids, 158:104668, 2022.
Πya Loshchilov and Frank Hutter. SGDR: Stochastic gradient descent with warm restarts. orXiv preprint arXiv:1608.0380 2016.
Nya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Conference on Learning Represendations (ICLR), 2019. URL attps://openrevier http://www..eeeeeeeeeee///e//aaa het//forumm/dd-BBoR..CC..
Winfried Lötzsch, Simon Ohler, and Johannes S Otterbach. Learning the solution operator of boundary value problems using graph neural networks. arXiv preprint arXin:2206. 14092, 2022.
Lu Lu, Pengzhan Jin, and George Em Karniadakis, DeepONet: Learning nonlinear operators for identi- fying differential equations based on the universal approximation thcorem of operators. arXiv preprint UrXiv:1910.03190 2019.
Lu Lu, Pengzhan Jin, Guofei Pang. Zhongqiang Zhang, and George Em Karniadakis. Learning nonlinear operators via deeponet based on the universal approximation theorem of operators. Nature Machine Intelligence, 3(3):218 229, 2021.
Lu Lu, Xuhui Meng, Shengze Cai, Zhiping Mao, Somdatta. Goswami, Zhongqiang Zhang, and George Em Karniadakis. A comprehensive and fair comparison of two neural operators (with practical extensions) based fair data. Compuler Methods in Applied Mechanics and Engineering, 393: 14778, 2022.
Hao Ma, Yuxuan Zhang, Nils Thuerey, Xiangyu Hu, and Oskar Haidn. Physics-driven learning of the steady navier-stokes equations using deep convolutional neural networks. arXiv preprint arxix:216. 093001, 2021.
Stéphane Mallat. A wavelet tour of signal processing. Flsevier, 1999.
Romit Maulik, Vishwas Rao, Jiali Wang, Gianmarco Mengaldo, Emil Constantinescu, Bethany Lusch, Prasanna Balaprakash, Ian Foster, and Rao Kotamarthi, Efficient high-dimensional variational data assim- ilation with machine-learned reduced-ordes models. Geoscientific Model Development, (6(8))3433-344 2022.
14

---

Published in Transactions on Machine Learning Research (07 (07/2023)
Andreas Mayr, Sebastian Lehner, Arno Mayrhofer, Christoph Kloss, Sepp Hochreiter, and Johannes Brand- stetter. Boundary graph neural networks for 3d simulations. orXiv preprint or XrXiv:2066.1 11299, 2021.
Ben Mildenhall, Pratul P Srinivasan, Matthew Tancik, Jonathan T Barron, Ravi Ramamoorthi, and Ren Ng. NcRF: Representing scenes as neural radiance fields for view synthesis. Communications of the ACM, 65(1):99 106, 2021.
Bruce Roy Munson, Theodore Hisao Okiishi, Wade W Huebsch, and Alric P Rothmayer. Fluid mechanics. Wiley Singapore, 2013.
Tung Nguyen, Johannes Brandstetter, Ashish Kapoor, Jayesh K Gupta, and Aditya Grover. ClimaX: A foundation model for weather and climate. arXiv preprint arXiv:2301. 100444,01..0343, 2023.
Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. In International Conference on Machine Learning, pp. 8162-811 8171. PMLR, 2021.
P.J. Olver. Symmetry groups of differential equations. In Applications of Lie Groups Differential Equations, pp. 77 185. Springer, 1986.
Jaideep Pathak, Shashank Subramanian, Peter Harrington, Sanjeev Raja, Ashesh Chattopadhyay, Morteza Mardani, Thorsten Kurth, David Hall, Zongyi Li, Kamyar Azizzadenesheli, Pedram Hassanzadeh, Karthik Kashinath, and Animashree Anandkumar. Fourcastnet: A global data-driven high-resolution weather model using adaptive fourier neural operators. arXiv preprint arxiv: 22X112202..1224, 2022.
Chao Peng, Xiangyu Zhang, Gang Yu, Guiming Luo, and Jian Sun. Large kernel matters matters-improve somantic segmentation by global convolutional network. In IEEE Conference on Computer Vision and Pallern Recognition (CVPR), pp. 4353 4361, 2017.
Ethan Perez, Florian Strub, Harm De Vries, Vincent Dumoulin, and Aaron Courville. FiLM: Visual reasoning with general conditioning layer. In AAAI Conference Artificial Intelligence, volume 32, 2018
Tobias Pfaff, Meire Fortunato, Alvaro Sanchez-Gonzalle and Peter W Battaglia. Learning mesh-based simulation with graph networks. arXiu preprint or arXiv:2010.03409, 2020.
Silvia L Pintea, Nergis Tömen, Stanley F Goes, Marco Loog, and Jan C van Gemert. Resolution learning in deep convolutional networks using scale-space theory. IEEE Transactions on Image Processing, 30: 8342-8353, 2021.
Timothy Praditia, Matthias Karlbauer, Sebastian Otle, Sergey Oladyshkin, Martin V Butz, and Wolf- gang Nowak. Finite volume neural network: Modeling subsurface contaminant transport. arXiv preprint arXiv:2104.06011 2021.
Md Ashiqur Rahman, Manuel A Florez, Anima Anandkumar, Zachary E Ross, and Kamyar Azizzadeneshelli Generative adversarial neural operators. arXiv preprint arXiv:2205.080 08011, 2022a.
Md Ashiqur Rahman, Zachary Ross, and Kamyar Avizzadenesheli. U-no: U-shaped neural operators. arXiv preprint ar xrXio:2204.11121 2022b.
Maziar Raissi, Paris Perdikaris, and George E Karniadakis. Physics-informed neural networks: A deep learn- ing framework for solving forward and inverse problems involving nonlinear partial differential equations. Journal of Compulational physics, 378:686-700 7076 2019.
Maziar Raissi, Alireza Yazdani, and George Em Karniadakis. Hidden fluid mechanics: Learning velocity and pressure fields from flow visualizations. Science, 36766481) 2020.
Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, and Ilya Sutskever. Zero-shot text-to-image generation. In International Conference on Machine Learning, pp. 8821 8831. PMLR, 2021.
15

---

Published in Transactions on Machine Learning Research (07/2023)
Yongming Rao, Wenliang Zhao, Zheng Zhu, Jiwen Lu, and Jie Zhou. Global filter networks for image classification. Advances in Newral Information. Processing Systems, 34, 2021.
Stephan Rasp and Nils Thuerey. Data-driven medium-range weather prediction with resnet pretrained on climate simulations: A new model for weatherbench. Journal of Advances in Modeling Earth Systems, 13 (20:0000000000 2021.
David W Romero, Robert-Jan Bruintjes, Jakub M Tomezak, Erik J Bekkers, Mark Hoogendoorn, and Jan C van Gemert. FlexConv: Continuous kernel convolutions with differentiable kernel sizes. arXiv preprint arXiv:21 10. 08.0009, 2021a.
David W Romero, Anna Kuzina, Erik Bekkers, Jakub M Tomezak, and Mark Hoogendoorn. CKConv: Continuous kernel convolution for sequential data. arXiv preprint #xXi:2102.02611 2021b.
David W Romero, David M Knigge, Albert Gu, Erik J Bokkers, Efstratios Gavves, Jakub M Tomczak, and Mark Hoogendoorn. Towards general purpose CNN for long range dependencies in ND. arXiv preprinl arXiv:2206.032: 03398, 2022.
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-Net: Convolutional networks for biomedical image segmentation. In International Conference Medical image computing and computerrassisted interven- lion, pp. 234 241. Springer, 2015.
David Ruhe, Jayesh K Gupta, Steven De Keninck, Max Welling, and Johannes Brandstetter. Geometric clifford algebra networks. arXiv preprint arXiv:2302.0659, 2023.
Alvaro Manchez-Gonzaler Jonathan Godwin, Tobias Pfaff, Rex Ying, Jure Leskovee, and Peter W. Battaglia. Learning to simulate complex physics with graph networks. ar preprint orXiv:2002.09405, 2020.
Kristof T Schütt, Huziel E Sauceda, P-J Kindermans, Alexandre Tkatchenko, and K-R Müller. SchNet a doep learning architecture for molecules and materials. The Journal of Chemical Physics, 148(24):241722, 2018.
Wenlei Shi, Xinquan Huang, Xiaotian Gao, Xinran Wei, Jia Zhang, Jiang Bian, Mao Yang, and Tie- Yan Liu. Lordnet: Learning to solve parametric partial differential equations without simulated data. arXiv preprint ar arXiv:2206.00. 19, 2022.
Martin Simonovsky and Nikos Komodakis. Dynamic edge-conditioned filters in convolutional neural networks graphs. In IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pp. 2017.
Karen Simonyan and Andrew Zisserman. Very deep convolutional networks for large-scale image recognition. arXiv preprint Xrc::11409 1556, 2014.
Justin Sirignan and Konstantinos Spiliopoulos. Dgm: A deep learning algorithm for solving partial differ- ential equations. Journal of computational physics, 375:1330-1364, 2018.
Casper Kaae Sonderby, Lasse Espeholt, Jonathan Heek, Mostafa Dehghani, Avital Oliver, Tim Salimans, Shreya Agrawal, Jason Hickey, and Nal Kalchbrenner. Metnet: A neural weather model for precipitation forecasting. arXiv preprint arXiv:2003. 12140, 2020.
Kimberly Stachenfeld, Drummond Fielding, Dmitrii Kochkov, Miles Cranmer, Tobias Pfaff, Jonathan Godwin, Can Cui, Shirley Ho, Peter Battaglia, and Alvaro Gonzzlzz Learned coarse models for efficient turbulence simulation. arXiv preprint arXiv:21 15275, 2021.
Christian Szegedy, Wei Liu, Yangqing Jia, Pierre Sermanct, Scott Reed, Dragomir Anguclov, Dumitru Erhan, Vincent Vanhoucke, and Andrew Rabinovich. Going deeper with convolutions. In IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pp. 1-9, 2015.
16

---

Published in Transactions on Machine Learning Research (07/2023)
Makoto Takamoto, Timothy Praditia, Raphael Leiteritz, Dan MacKinlay, Francesco Alesiani, Dirk Pfttiger, and Mathias Niepert. PDEBench: An Extensive Benchmark for Scientific Machine Learning. In Newnal Information Processing Systems (NewrIPS) Track on Datasets and Benchmarks, 2022. URL https://
coi.org/100 10. 18419//arrss229990
Roger Temam. Navier-Stokes equations: theory and numerical analysis, volume 343, American Mathematical Soc., 2001.
Hendrik Tennekes, John Leask Lumley, Jonh L Lumley, et al. A first course in turbulence. MIT press, 1972.
Kiwon Um, Robert Brand, Yun Raymond Fei, Philipp Holl, and Nils Thuerey. Solver-in-the-lopp- Learn- ing from differentiable physics to interact with itcrative pde-solvers. Advances in Neural Information Processing Systems, 33:6111-61222 2020.
Charles Van Loan. Computational frameworks for the fast Fourier transform. SIAM, 1992.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in Neural Information Processing Systems, 30, 2017.
Cornelis Boudewijn Vreugdenhil. Numerical methods for shallow-water flow, volume 13. Springer Science & Business Media, 1994.
Nils Wandel, Michael Weinmann, and Reinhard Klein. Learning incompressible fluid dynamics from scratch- towards fast, differentiable fluid models that gencralize. arXiv preprint arXiv:2006. 0e02200000766,, 2020.
Nils Wandel, Michael Weinmann, Michael Neidlin, and Reinhard Klein. Spline-PINN: Approaching PDEs without data using fast, phystes-informed hermite-spline CNNs. In AAAI Conference on Artificial Intel- ligence, volume 36, pp. $529-8538, 8538, 2022.
Rui Wang, Karthik Kashinath, Mustafa Mustafa, Adrian Albert, and Rose Yu. Towards physics-informer deep learning for turbulent flow prediction. In Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Dala Mining, pp. 1457-1466, 2020a.
Rui Wang, Robin Walters, and Rose Yu. Incorporating symmetry into deep dynamics models for improved generalization. arxiv preprint arXiw:2002.0300 2020b.
Shenlong Wang, Simon Suo, Wei-Chiu Ma, Andrei Pokrovsky, and Raquel Urtasun. Deep parametric con- linuous convolutional neural networks. In IEEE Conference on Computer Vision and Pallern Recognition (CVPR), pp. 1189-259,, 2018
Gege Wen, Zongyi Li, Kamyar Axizzaenesheli,, Anima Anandkumar, and Sally M Benson. U-fno- an enhanced fourier neural operator- based deep-learning model for multiphase flow. Advances in Water Resowrces, 163:104180, 2022.
Jonathan Weyn, Dale Durran, and Rich Caruana. Improving data-driven global weather prediction using deep convolutional neural networks on a cubed sphere. Journal of Advances in Modeling Earth Systems, 1200):2000000000000, 2020.
Jonathan A Weyn, Dale R Durran, Rich Caruana, and Nathaniel Cresswell-Clay Sub-seasonal forecasting with a large ensemble of deep-learning weather prediction models. Journal of Advances in Modeling Earth Systems, 13(7)=e2021 2021.
Daniel Worrall and Max Welling. Deep scale-spaces: Equivariance over scale. Advances in Neural Information Processing Systems, 32, 2019.
Tailin Wu, Takashi Maruyama, and Jure Leskovee. Learning accelerate partial differential equations via latent global evolution. arXio preprint ar arXi::2206.0765) 2022.
17

---

Published in Transactions on Machine Learning Research (07 (07/2023)
Wenxuan Wu, Zhongang Qi, and Li Fuxin. PointConv: Deep convolutional networks 3d point clouds. In IEFE/CVF Conference on Compuler Vision and Pollern Recognition (CVPR), pp- 9621-9630, 2019.
Yuxin Wu and Kaiming He. Group normalization. In European Conference on Computer Vision (ECCV), pp. 3-19, 2018.
Yan Yang, Angela F Gao, Jorge C Castellanos, Zachary Ross, Kamyar Axizzadenesheli, and Robert W Clayton. Seismic wave propagation and inversion with neural operators. The Seismic Record, 1(3):126 134, 2021.
Sergey Zagoruyko and Nikos Komodakis. Wide residual networks. In British Machine Vision Conference 2016. British Machine Vision Association, 2016.
Yinhao Zhu and Nicholas Zabaras. Bayesian deep convolutional encoder-docoder networks for surrogate modeling and uncertainty quantification. Journal of Compulalional Physics, 366:415-447, 2018.
Kirill Zubov, Zoe McCarthy, Yingbo Ma, Francesco Calisto, Valerio Pagliarino, Simone Azeglio, Luca Bot- tero, Emmanuel Luján, Valentin Sulzer, Ashutosh Bharambe, al. NeuralPDE: Automating physics- informed neural networks (PINNs) with error approximations. ar preprint arXiv:2107.090 0944443 2021.
18

---

Published in Transactions on Machine Learning Research (07 (07/2023)
Contents
1 Introduction
1
2 Preliminaries
2
3 PDE Surrogates
5
3.1 Operator learning
6
4 Experiments
7
5 Conclusion
10
A Related work
20
B Experiments
20
B.1 Experimental details
20
B.2 Additional model details
B.2.1 ResNet
21
B.2.2 Dilated ResNet
21
B.2.3 FNO
21
B.2.4 U-Net
21
B.2.5 Parameter Conditioning
22
B.2.6 Spatial-spectral parameter conditioning for Fourier layers
................................
22
B.3 Shallow water equations.
23
B.4 Navicr-Stokes equations.
34
B.5 Parameter conditioning.
40
19

---

Published in Transactions on Machine Learning Research (07 (07/2023)
A
Related work
Neural PDE modcling is appearing in many flavors. Various works are numerical-neural hybrid approaches where the computation graph of the solver is preserved and heuristically-chose parameters are left for the neural network to predict (Bar-Sinai et al., 2019; Kochkov et al., 2021; Greenfeld et al., 2019; Hsich et al., 2019; Praditia et al., 2021; Um et al., 2020; Garcia Satorras et al., 2019). The works of Gonozlzz et al. (2020); Pfaff et al. (2020); Mayr et al. (2021) are of similar flavor where neural network predictions are input to the time-update of node positions in graphs and meshes. Fully neural network based approaches can be roughly split into two parts. First, mothods that focus on the approximation of the solution function of the underlying PDE (Sirignano & Spiliopoulos, 2018; Han et al., 2018; Raissi al., 2019; Jin el. al., 2021;
Raissi et al., 2020; Zubov et al., 2021). And second, methods that focus on the surrogate learning of solution operators. CNN-basc models
among the first PDE surrogates (Guo et al., 2016; Bhatnagar et al., 2019;
Zhu & Zabaras, 2018). Operator learning models were popularized via Fourier Neural Operators (Li et al., 2020a) and FNO-based applications and refinements (Li et al., 2021b; Rahman et al., 2022b; Rao et al., 2021; Guibas et al., 2021; Li et al., 2021a; Rahman et al., 2022a; Pathak et al., 2022; Wen et al., 2022; Liu et al., 2022; Yang et al., 2021; Guan et al., 2021; Hwang et al., 2022; Chen et al., 2021; Li et al., 2022b), as well as via DeepONets (Tuu et al., 2019; 2021; 2022). Other directions include the modeling of PDE solution operators via latent space models, transformers, and graph neural networks (GNNs) (Wu et al., 2022; Li et al., 2022a; Brandstetter et al., 2022c; Lötzsch et al., 2022; Lienen & Günnemann, 2022). The ever persisting chicken-egg problem (Brandstetter et al., 2022b; Shi et al., 2022) of how to obtain high quality ground truth training data for neural PDE surrogates is approached either via clever data augmentation (Brandstetter et al., 2022b), via equivariant neural solvers (Wang et al., 2020b), or via "data-free" learning paradigms (Geneva & Zabaras, 2020; Wandel et al., 2020; 2022; Shi et al., 2022). Pratical applications of neural PDE surrogates can especially be found in weather forecasting (Pathak et al., 2022; Guibas et al., 2021; Keisler, 2022; Rasp & Thuercy, 2021; Weyn et al., 2020; 2021; Arcomano et al., 2020; Sonderby ct al., 2020; Frerix et al., 2021; Maulik el al., 2022; Lam al., 2022; Bi et al., 2022; Nguyen et al., 2023), and modeling of fluid dynamics (Ma et al., 2021; Stachenfeld et al, 2021; Wang et al., 2020a; Brandstetter et al., 2022a; Ruhe et al., 2023).
B Experiments
This appendix supports Section of the main paper.
B.1 Experimental details
Loss functions and metrics. We report the summed MSE (SMSE) loss defined as:
N₂
Naalde
LSMSE =
Σ Σ [|u((,,t,) á. t,),
(8)
N2
yE2 j=1
where is the target, à the model output, Naclds comprises scalar fields well individual vector field components, and N, is the total number of spatial points. Equation 8 is used for training with N, = 1, and further allows us to define our two main metrics:
One-step loss where N₂ 1 and Nnelds comprises all scalar and vector components.
Rollout loss where N; = 5 and Nielss comprises all scalar and vector components.
Alternatively, we train with the summed relative MSE (RMSE) loss as introduced in Li et al. (2020a):
1
N. Nialda
[u,yy, t2) L (v,tt))1
LSMSE
(9)
Ny
MEZ
j=1
==1
20

---

Published in Transactions on Machine Learning Research (07 (07/2023)
Training and model selection. We optimized models using the AdamW optimizer (Kingma & Ba, 2014;
Loshchilov & Hutter, 2019) with the best. learning rates of [10-1,2 2 20--) and weight decay of 10-5 for 50 epochs and minimized the summed mean squared error (SMSE) which is outlined in Equation 8. We used cosine anncaling as learning rate scheduler (Loshchilov & Hutter, 2016) with a linear warmup. For baseline ResNet models, we optimized number of layers, number of channels, and normalization procedures. For the reported results we used group normalization (Wu He, 2018) with group which equivalent to Layer norm (Ba et al., 2016) (except for final normalization layer in U-Nets where we use groups). We further tested different activation functions. For baseline FNO models, optimized number of layers, number of channels, and number of Fourier modes. Larger numbers of layers channels did not improve the performances for both ResNet and FNO models. For U-Net like architectures, especially for U-Netatt: we specifically needed to optimize the maxinum learning rale to be lower (10-1). We further optimized for different number of hidden layers, and initialization and normalization schemes. For the reported results, we usod pre-activations (He et al., 2016b) and layer normalization (Ba et al., 2016). We used an effective batch size of 32 for training.
Computational resources. All exporiments used 4x16 GB NVIDIA V100 machines for training. Average training times varied between 2 and 140h, depending task and number of trajectories. Parameter conditioning runs were the most expensive ones.
Runtime comparison. We warmup the benchmark for 10 iterations and report average runtimes over 100 runs single 16 GB NVIDIA V100 machine with input batch size of 8.
B.2
Additional model details
B.2.1
ResNet
We use two embedding and two output layers with kernel sizes of 1 1.
B.2.2
Dilated ResNet
The implemented Dilated ResNet models consist of 4 residual blocks where each block individually consists of dilated CNN layers with dilation rates of [1. 2, 4, 8, 4, 1]. We implement Dilated ResNets with and without group normalization layers applied to each layer in the respective dilation blocks.
B.2.3 FNO
We use FNOs consisting of {4,8) FNO layers, where {8, 16, 16,32} modes are multiplied in the Fourier space, and (64,128) channels are used. We use two embedding and two output layers with kernel sizes of 1 x as suggested in Li et al. (2020a). The number of non-zero Fourier modes, the number of FNO layers and the nunber of channels are hyperparameter, where we report results for different values in each of the experiment. We use GeLU activation functions, and no normalization scheme. Normalization schemes and residual connections did not improve performance, as already reported Brandstetter et al. (2022a).
B.2.4 U-Net
We use embedding and one output layers with kernel sizes of 3 x 3. To allow fair comparison to FNO (and ResNet) architectures, we ablated architectures also for kernel sizes of 1 x 1 for embedding and output layers.
U-Net2o1s- We use channel multiplicrs of 22, The network consists overall of 4 downsampling, one bottleneck and 4 upsampling layers. We further use batch normalization (Toffe & Szegedy, 2015), and no bias weights. The implementation is based on the PDEbench repository of Takamoto et al. (2022). Compared to the implementation of Takamoto et al. (2022), we use GeLU activations instcad of tanh activations since we observe significant better performances. For the sake of completeness, we also report tanh results terming the models U-Nel
21

---

Published in Transactions on Machine Learning Research (07 (07/2023)
U-Netbase We use channel multipliers of (2, 2, 2). We replaced batch normalization (loffe & Szegedy, 2015) with group normalization (Wu & He, 2018) with number of groups equal 1 to be consistent with other architectures. Additionally, compared to the U-Net2015 version, we use bias weights but no bottleneck layer.
U-Netmod- We use channel multipliers of (1, 4), and residual connections in each down- and upsam- pling block. We use pre-normalization and pre-activations (He et al., 2016b). Additionally, we zero-initialize the second Conv layer in each residual block.
att Adding attention to all down- and upsampling blocks made training unstable, and would have required an extensive hyperparameter search. We therefore only use attention in the middle blocks after downsampling. We further only use a single attention head along with a residual connection by bypassing attention.
B.2.5 Parameter Conditioning
Embedding. We use sinusoidal embedding as proposed in Vaswani et al. (2017) for positional encoding of scalar values, such as prediction time window and force strength:
Emb(2, d) = [[cos cos 1000022/d; : sin 100002#/d
x
x
for 0 0 < <d,
(10)
where z is the embedded quantity and d is the output embedding dimension.
Projection. We use a two-layer feed-forward network to project each of the embeddings to higher dimension (4x hidden channels), and add them together before passing them to each block via another linear layer.
Conditioning. We explore two different mechanisms for conditioning, originally proposed in the image modeling literature. Simple "Addition" as proposed by Ho et al. (2020) which easily be extended to FNO layers and " "AdaGN" proposed by Nichol & Dhariwal (2021) which requires normalization layers to be applicable and was therefore restricted to U-Net based architectures in our experiments.
Addition: A single Linear layer is used to scale the dimensions appropriately to match the dimen- sions of the conditioned block. Conditions are added to the first Conv layer's output, followed by normalization and a second Conv layer. This conditioning is applied to all blocks in the network.
AdaGN. Instead of directly adding the Linear projection of the embedding to the respective blocks, the projection y is split into [y., yo] to scale and shift the normalized output h before passing to the second Conv layer in each block, similar to is done in FiLM (Perez et al., 2018):
h' ys @ GroupNorm(h) + yb
(11)
where @ denotes the pointwise product. The conditioning is applied to all blocks. Since FNO architectures were implemented without group norm, AdaGN was not applicable to those.
In Appendix B.5, we ablate " "Addition" and "AdaGN" embeddings for U-Net architectures.
B.2.6 Spatial-spectral parameter conditioning for Fourier layers
Since FNO like architectures are usually implemented without normalization schemes, only "Addition" is applicable as conditioning strategy. Adding the conditioning at the end of each FNO layer, i.e. in the spatial domain, omits that the conditioning information is accessible in the Fourier domain too. This was somewhat unsatisfying,
explored straightforward mechanism to apply conditioning to the Fourier branch as well. We implement an alternative FreqLinear layer to project the embeddings into Fourier space too. Each Fourier mode is first multiplied with the embedding, then weights are mode-wise multiplied and the inverse Fourier transform is performed. Adding conditioning both in the Fourier and the spatial domain soems work best. We term this alternative embodding "Spatial-Spectral embodding. In Appendix B.5, we ablate " "Addition" and "Spatial-Spectraa" embeddings for FNO like architectures including UFNets with FNO blocks in the downsampling path.
22

---

Published in Transactions on Machine Learning Research (07 (07/2023)
B.3 Shallow water equations.
The shallow water equalions are solved on regular grid with periodic boundary conditions as described in Section 4 of the main paper. The inputs to the shallow water experiments are respective fields at the previous timestops. Pressure and vorticity fields are normalized for training. Example rollout trajectories are displayed in Figure for the velocity function formulation and in Figure 9 for the vorticity stream function Figures 10,11,
formulation. We outline further details on the results on the shallow water experiments and Tables 2,3. Additionally, we show results for 1-day predictions in Figure 12 and Table 7. We further ablate different encoding/ decoding choices for U-Net like architectures in Figures 13,14. Finally, compare different the specs of different FNO, UNO, and U-FNet architectures in Table 4.
Ground
Prediction
Ahs. emnerr
(a) Scalar pressure field
Ground
Prediction
Abs, CiTir
(b) Vector field z-component
Ground
Abs. onoer Prodietion
(c) Vector field y-component
Figure 8: Shallow water 2-day predictions, velocity function form. Example rollouts of the scalar pressure and the vector wind field of the shallow water experiments are shown, obtained by a U-FINetmodea6 PDE surrogate model (top), and compared to the ground truth (bottom). Predictions are obtained for a time window At 48 h. The respective model input fields comprise two timesteps, we only show the first those (Jeft-most ground truth column).
23

---

Published in Transactions on Machine Learning Research (07/2023)
Ground
Prediction
Abs. einer
(a) Scalar pressure field
Ground
Ales. crror Prodictiio
(b) Scalar vorticity field
Figure 9: Shallow water 2-day predictions, vorticity stream function form. Example rollouts of the scalar pressure and the scalar vorticity field of the shallow water experiments are shown, obtained by U- F2Netmodes16.16 PDE surrogate model (top), and compared to the ground truth (bottom). Predictions are obtained for a time window At = 48 h. The respective model input fields comprise two timesteps, we only show the first those (left-most ground truth column).
24

---

Published in Transactions on Machine Learning Research (07/2023)
ResNNt
ResNet226
DilResNet 28
DilResNet norm
FNOI 25-012888aaa 28-8
FNOI 28-410
FNO64-4moder
FNO96-4modess
PNO128-4modoo
UNO64
UNO128
U-Not2o1564
U-Netgo16 128
U-Netbae64
U-Netuuae 128
G-Notaeed64
U-Notai164
U-F I.FINetmodos
Neffooooo
U-F2Nctmaddss..4
C-F2Notmadde11..
U-F2Notmodess.
U-F2Net moods.01
Shallow water velocity, rollout
Shallow water velocity, one-step
10
100
109
10
MSB
10
10
10
10
10
5600
448
5600
Num. Train Trajectorics
Figure 10: Shallow water 2-day predictions. Rollout and one-step of various architectures the shallow water equations are reported. Results obtained for predictions of 2-day time windows for the velocity function formulation and are averaged over three different random seeds. Note the logarithmic scale of the y-axes.
25

---

Published in Transactions on Machine Learning Research (07 (07/202) /022)
R2s
ResNet256
DilltesNet128
DilRessNe112-nonn
FNO128-8
FNO 5
FNO 128-4mo
FN06444modes32
Amodee4modee32
2X0128-4mode332
UNO64
UNQ128
U-Netaor364
U-Netz013 128
U-Netaaar64 64
U-Netbawe 128
U-Netate64
U-F U-FI INetwodess
U-FI 1o
U-F2Ner
U-F2Not
U-F2Nctmodo10..1
Shallow water vorlicily, rollout.
Shallow water vorticity, one-step
109
10
10
NNSE
10
10
10
10-3
10
1071
448
5600
448
5600
Num.
Train Trajectories
Figure 11: Shallow water 2-day predictions. Rollout and one-step errors of various architectures on the shallow water equations. Results are obtained for predictions of 2-day time windows for the vorticity stream function formulation and are averaged over three different random seeds. Note the logarithmic scale of the y-axes.
Shallow water velocity, rollout
Shallow water velocity, one-step
309
RosNat122
ResNet256
10
DitResNet129
10
DifReeNett 28-m
FNO128-8..ddod
UN064
MSB
10
UNO128
U-Nel 64
U-Nettuee12 128
U-Neto
10
U-Net.a:61
U-FINetanddeil
U-P2Nermudutts
10
10
U-P2Nesaudusee
236
2800
236
2800
U- C-F2Neesstttts.cc
Num. Train Trajeclories
Figure 12: Shallow water 1-day predictions. Rollout and one-step errors of various architectures on the shallow water equations are reported. Results are obtained for predictions of 1-day time windows for the velocity function formulation and averaged over three different random seeds. Note the logarithmic scale of the y-axes.
26

---

Published in Transactions on Machine Learning Research (07 (07/202) 2023)
Shallow water vorticity, rollout
Shallow water vorticity, one-step
10
10
U-Nesman66--B 3xc3
U-N
U.Netart 3od
U-Net.n
10
U- P.Nettoddas Ixl
MSE
U-FINet..... 3-33
T-P U.FIINet
U-P2Nelmossan 4-3a
modeess a-1al
U.P2Nctmosee::44
10
448
5800
448
5600
Num. Train Trajectorics
Figure 13: Shallow water 2-day predictions. Ablation results of different encoding and decoding choices for various U-Net architectures are reported. 1 x and 3 3 kernels are compared for both encoding and decoding. Rollout and one-step errors are obtained on the shallow water equations for 2-day predictions for the vorticity stream function formulation and are averaged over three different random seeds.
Shallow water velocity, rollout
Shallow water velocity, one-step
U. N.temon 3bc3
U-Nettann66111
64-
U.Ne 64-1xl
U-F
U-P U-PiNet modess Ixl
MSE
U-FINol
U-F1 Ne 1sl
U. P..PPNetmmsssss.. 4 Bxll
10
448
5600
448
6600
Num. Train Trajectories
Figure 14: Shallow water 2-day predictions. Ablation results of different encoding and decoding choices for various U-Net architectures are reported. 1 x and 3 x 3 kernels are compared for both encoding and decoding. Rollout and one-step errors are obtained on the shallow water equations for 2-day predictions for the velocity function formulation and are averaged over three different random seeds.
27

---

Published in Transactions on Machine Learning Research (07/2023)
Shallow water 2-day predictions, velocity function formulation. Rollout and one-step errors of various architectures on the shallow waler equations are reported. Summee mean-squared errors (SMSE) are obtained for 2-day predictions for the velocity function formulation and are averaged over three different random seeds. If results are displayed without standard deviation, the obtained standard deviation lower than the four digit precision minimum. The best model of each model class is highlighted.
Table 2:
SMSE
METHOD
Trajs.
onestep
rollout
RosNet122
448
0.7787 + 0,0049
5.7408 : 0. 0.0223
ResNell18
5600
0.5465 +: 0.0130
5.4946 1: 0.2009
Res
148
0.4751 + 0.0005
1.0954 + 0.0181
NesNet256
5600
0.5294 + 0.0364
5.4155 0..0:00000 0.0657
DilRes 128
448
0.3800 + 0.0042
2.4860 + 0.0260
NetRes
5600
0.0:29 + 0,0014
0.5476 = 0.0109
DilRtesNet128-mun norm
448
0.1723 + 0.0026
1.3377 + 0.0144
DilltesNe128-nomm
5600
0.0262 + 0.0006
0.3770 + 0.0081
FNO128 Smodest
448
1.0322 1 0.0055
3,8635 1: 0.0080
FNO128-S Smodd
5600
0.2023 + 0,0023
0.85-19 + 0124
FNO128-8
0,3681 + 0,0042
1.7088 0.0096
FNO128-8 moddmis
5600
0.0236 + 0.0001
0.0878 + 0.0007
FNO128- 180128-mmodosi
448
0,3802 + 0.0021
1.85-12 + 0 .0066
FNO128-46 128--modssi
5600
0.0397 + 0.0002
0.1601 + O. 00100
FNO64- 7N064-4madsa2
448
0.3750 + 0.0012
2.039K + 0. 0050
F- Amoddad2
5600
0.0225 + 0.0007
0,1015 + 0044
PN096-4mateen
0.2794 0.0053
1.7888 0.0253
PN096-4modsse
5600
0.0102 + 0.0002
0.0399 + 0.0008
FNO128-4m 4msdeedd
448
0.2492 + 0,0040
1.4460 + 0.0226
FNO128
4modesa2
5600
0.0060 + 0.0001
0.0213 1 0.0003
UNO61
148
0.8131 + 0,0048
3.6621 # 0.0103
UN064
5600
0.0319 + 0.0003
0.1208 + 0.0010
UNO128
448
0.6328 + 0.0041
3.02/0 + 0064
UNO128
5600
0.0096 + 0.0001
0,0282 1 O. 0.0022
U- U-Nofbaac64 64
448
0.1693 + 0.0026
1,2224 + 0.0108
U U-Nelto..6l
5600
0.0128 +: 0.0016
0.1026 1 0.0161
U.Net pase 128
448
0.1076 4: 0.0008
0.9096 + 0.0067
U-Net 128
5600
0,0054
0.0139 1 0. 0001
U- 144
448
0.1033 + 0.0001
0.8847 + 0.0109
U-Netmea61
5600
0.0031 + 0.0001
0.0165 + 0 00455
Ne
0.1013
0.0028
0.8681
U-Notam.6441x
5600
0.0031 + 0.0003
0.0389 + 0.0057
U-Nelast
448
0.0954 1: 0.0014
0.8158 1 0.0318
U-Net 64
5600
0.0060 + 0.0005
0.0684 + 0.0206
U-Net
448
0.0819 + 0.0060
0.6419 1 0.1307
U-Nel 64-1x1
5600
0.0051 1 0.0001
0.0724 1 0.0037
U-FINetmssost
448
0.0797 + 0.0010
0.4553 # O 0057
5600
0.001 + 0.0002
0.0081 0.0012
U-FiNNtmmddsi
448
0.0717 + 0.0013
0.3988 + 0.0115
5600
0.0018 1 0.0001
0.0060 1 O. 0.0002
U-FINet Ixl
448
0.0743 + 0.0032
0.4185 + 0.0259
U-F T-FiNetmodess- -Ix
5600
0.0039 + 0.0018
0.0185 1 0. 0.0079
U-F U-FINelmodent 1xl
448
0.0699 1 0.0007
0.3923 1 0. 0.0001
5600
0.0027 1 0.0004
0.0089 0.0017
U-F2Netmodww.
448
0.0843 1 0,0001
0.4931 0.0031
U-F2Netamodo8..44
5600
0.0027 + 0.0012
0.01 + 0.0050
U-F2Netmodwin.
448
0.0765 + 0.0019
0.4508 + 0.0116
5600
0.0017 + 0.0004
0.0052 + 0.0013
U-F2Net
448
0.0793 + 0.0001
0.4553 + O. 0.0021
U.
5600
0.0013
0.0044 1 O. 0.0002
0.0906
0.0042
0.5596 0.0317
F2Net modo116.11
5600
0.0012 + 0.0001
0.0037 + 0. 0.0002
U Ee
448
0.0793 1 0.0016
0.4587 1 0.0064
U. F2N
1xl
5600
0.0015 1 0.0001
0.0067 1 0.0005
U-F2N U-F2Netmodetts Mondeets x-.. -xx
148
0.0728 + 0.0012
0.1238 1 0. 0.0019
moder16 e-
xlx1 1x1
5600
0.0016 + 0.0003
0.0054 1 0. 00009
U-F2 U-F2Ne matddii
448
0.0769 + 0.0052
0.4517 + 0 03-18
U-F2N modeeti
5600
0.0021 + 0.0010
0.0078 0.0035
U-F D-F3Net
448
0.0920 + 0,0011
0.5318 1 0. 0.0083
U-F3Net
5600
0.0018
0.0062 + O .0001
U-F3N U-F3Nefmmdeedee efmodeeete...
148
0,0812
0.1623 1 0. 0.0021
U-F3Nrt mf0dd.....
5600
0.0016
0.0048
28

---

Published in Transactions on Machine Learning Research (07/2023)
Table 3: Shallow water 2-day predictions, vorticity stream function formulation. Rollout and one-step errors of various architectures on the shallow water equations are reported. Summed mean-squared errors (SMSE) are obtained for 2-day predictions for the vorticity stream function formulation and are averaged over three different random scods. If results are displayed without standard doviation, the obtained standard deviation is lower than the five digit precision minimum. The best model of each model class is highlighted.
SMSE
METROD
Trajs.
oncestop
rollout
ResNet128
448
0.189
0.00232
1,80900+ 0.04636
ResNet128
5600
0.221 82 = 0. 03
2.04572 0.050
ResNet256
418
0.132 08 + 0.0033
1.710 + 0.001 12
ResNe1256
5600
0.183 80 1 0.021 08
1.906 15 + 0.07268
DilResNNe128
448
0.077 97 0.003
0.598 1. 0.02661
5000
0.01047
0.000
0.14051 0.003
448
0.041 05 0.000
0.345 + 03
DilRes DelResNet12S-nnn Ne128-norm
5600
0.006 19 + 000
0.095 + 0.001 10
PNO128- Smnddes
448
0.185 63 0.00278
0.892 44 + 0.011 09
FNO128- Swrddoed
5600
0.039 17 0.001 01
0.213 + 0.00281
FNO128.6 Smodes16
0.100 54
0.57404
FNO128-4 28-8 10118888m00011 Smodk16
5600
0.005 44 0.00001
81 + 0.00005
FNO 28-4modde16
448
0.086 01 + 0.000
0.511 27 + 4.00041
FNO1 28-1mmeee11
5600
0.008 36 + 0. 03
0.047 23 + 0.00005
FNO6111
7NO62-Amedes22 Amoden522
418
0.089 23 0.00233
0.553 = 0.011
FNO64-4msden Amodes82
5600
0.005 64 0.000
FNO96-1
418
0.082 08 + 0.001 32
0.196 11 0.008 51
FNO96-4medenn
5600
0.002 58 + 0.00003
0.015
0.000
FNO128-42 4modes52
448
0.074 83 0.000
0.455 1 0.00224
FNO128
4modes32
5600
0.001 47
0.008 47
UNO64
448
0.21437
1.068 14
UN064
5600
0.007 51 + 0.00001
0.039 18 + 0.00002
UNO128
448
0. 161666 32 0.001 40
0.848 + 0.001 90
UNO128
56000
0.00207
0.008 58
U-Nelhsee64
0.031 22 0.0000
0.28432 16
U.Nelloa.64
5600
0.002 03 0.000
0.02277 1. 0.000
U-Netawe12 128
448
0.025 05 0.000 35
0.227 93 + 0.00297
U-Netawe 128
5600
0.000
99 + 0. 0.00004
0.01425 + 0. 80
04
448
0.023 03 0.000 20
0.217 13 + 0.002 54
5600
0.000
0.00003
0.01031 0.000
U-Netmaa61-Ix
448
0.023
83 0..00.0000
0.218 + 0.000
Uo Txl
5600
0.000 60
0.01063 + 0.000
U-Nota 64
448
0.022 69 0.000
0.196 43 1 0.005 66
C-
5600
0.001
08 0.000 55
0.01994 + 0.016 17
U-Notaan4-Is
448
0.021
33 + 83
0.17791 + 0.021
5600
0.000
0.00361 0.00005
U-FINet
448
0.018
27 0.000 20
0. 0.9 + 0.00297
U.FINel gooees
5600
0.000 41
0.00291 1 0.
U-FINetmedss6
448
0.01882
0.00007
0.131 56 + 0.000 74
5600
0.000
0.00008
0.003 42 + 0.000
U-FINetmudus Ix1
448
0.017 71 0.00001
0.127 + 0.001
U-FINetmodoe Ix1
5600
0.000 41 0.00002
0.00336 + 0.000 11
D-FINel Ixc
448
0.018 82 1 0.000
0.13287 1 0.001 27
U.FINet po0cs11111 Ixl
5600
0.000
0.000
0.002 1 0.000
U.F2Net
448
0.01751
0. 19
0.12200 + 0.001
5600
0.000 39 0.0000
0.002 0.000
U-F2Netmmodos, -Ix1
418
0.016 56 0.000
0.115 + 0.003
e- -ll
5600
0.000 1 0.000
0.002 1 01
U.P2Ncsmooss16
448
0.018 90 0.000 42
0.13388 1 0.003 63
5000
0.000
0.00001
0.001 0.00002
U-F2Netmmobo16.
448
0.017 05 5.000
0.12052 + 0.004
U.P2Ne Ixl
5600
0.000
33
0.001 83
U-F2Netmuse88a
448
0.016
0.000
0.11781 + 0.00233
5600
0.000
32 0.00001
0.001 + 0.00005
U-F2Net
418
0.035 66 6.00
0. 25 100000000 0.028
U-F2Nel m5
5600
0.000 32
0.001 GT + 0.00001
U-F2Nel att
448
0.017
31 0.00031
0.121 50 1 0.001 19
U. F.F22elaat mm.desss6..
5600
0.000 66
0.003
U-FaNe Ne
0.022 58 0.
0.147 0.00131
U-FSNel 2e0oooeee
5600
0.000 40
0.00237 1 0.00001
U- U-F3Nel mlodeeood
448
0.022 03 0.000
0.15885 1 55
5600
0.000 35
0.00001
0.00208 + 0.00005
29

---

Published in Transactions on Machine Learning Research (07/2023)
Table 4: Comparison of parameter count, runtime, and memory requirement of various FNO, UNO, and U- FNet architectures. Subscript numbers indicate the used number of Fourier modes. For U-FNet experiments subscript numbers indicate the number of Fourier modes in the lowest, second-lowest, and third lowest block.
Runtime
Mem. [MB]
METHOD
Channels Bees.aaers//Boor
Params.
Fwd.
Fwd.+
132 size
Peak usage
FNO128- INO128-Bmedeat
128
33.7 M
0.057
0.162
134
2161
FNO FNO128-85 28-8mede216
128
8
134
0.059
0.171
537
2953
FNO128-4modee Amedenie
128
07.2
0.031
0.089
268
FNO61- 10661-4mmde322
64
4
67.1 M
0.016
0.050
268
1204
PNO96-4modess3
FNO128
7NO128-44modeoo
128
4
268 M
0.036
0.818
1100
3420
110
0.070
0.134
440
1925
UNO128
128
440
0.160
0.341
1800
5513
U-FINebmodes
154
U-FI Neftndeeee
64
9
185 M
0.084
0.208
743
4037
0.213
652
U-F2N NeF22eemmd &
61
9
193 M
0.085
0.216
772
1016
U-F2N ebmode116,8
64
9
224
0.086
U.211
897
4149
U-F2Net, df moddelsl1.
64
344 M
0.090
0.232
1400
4496
U-F3N ef mode0e888...
64
9
198
0.086
0.221
658
4332
U.F3N m.odessts..
30

---

Published in Transactions on Machine Learning Research (07 (07/202) 2023)
Table 5: Shallow water 2-day predictions, velocity function formulation. Rollout and one-step errors of various architectures the shallow water equations reported. L2 training objective of Li et al. (2020a) is used. Summed mean-squared errors (SMSE) are obtained for 2-day predictions for the vorticity stream function formulation and are averaged over three different random seeds. results are displayed without standard deviation, the obtained standard deviation is lower than the five digit precision minimum. The best model of each model class highlighted.
SMSE
METHOD
Trajs.
onestep
rollout
DilResNet128
448
0.312 + 0.004 90
2.115 = 0.036
DieResNett188
5600
0.042330
0.000 75
0.542 0.012
448
0.148 49 0.002 38
1.197 04 0.008 19
DilResNet12S-noom
5600
0.026 + 0.000 73
0.382 68 + 0.007 60
FNO128 1001288-4mddsl// 1odes16
448
0.339 75 + 0.000
1.673 92 + 0.003 99
FNO128 20112-44.0 1 160128-4modes16
5600
0.040 + 0.000 49
0.163 17 + 0.002
FNO64-4modes
448
0.302 18 + 0.002 98
1.678 19 + 0.015
FNO64-4 modess22
5600
0.023 + 0.001 02
0.107 56 + 0.006 50
PNO96-4modessi
448
0.234 + 0.005 60
1.351 64 + 0.025 51
FNO96-4 #V09664modrss13
5600
0.010
0.000
0.040 0.001 50
U-Netbasc 64
448
0.147 26 0.002 30
1.115 39 0.005
U-Nethaso 64
5600
0.009 88 + 0.000 11
0.067 11 + 0.001 57
U-Netiase 128
448
0.095 07 + 0.000 22
0.84498 + 0.002 95
U- U-Netisase 128
5600
0.004 + 0.000 08
0.032008 + 0.000 44
U-Net201566
448
0.169 45 + 0.001 13
1.264 75 + 0.003 96
U-Net201564
5600
0.012 79 + 0.000 09
0.084 20 + 0.001 29
U-Net2015 128
448
0.114 96 + 0.000 23
0.985 53 + 0.001 19
U-Net2013 128
5600
0.005 41 + 0.000 06
0.039 10 + 0.001 39
5- 64
448
0.475 99 0.002 45
2.760 14 0.004 88
ta 64
5600
0.029 97 + 0.001 37
0.19094 0.006 18
15-tanh U-Netgo 15-tanh 128
448
0.333 82 + 0.004 30
2.187 05 + 0.025 79
U- U-Netgo15-tanb 128
5600
0.015 05 + 0.000 32
0.097 31 + 0.001 29
ooNet 64
448
0.088 51 + 0.001 42
0.808 42 + 0.011
mod 64
5600
0.002 25 + 0.000
0.021 95 + 0.001
U-F2 m
448
0.061 98 + 0.000
0.371 33 + 0.001 92
U-F2Netmodesste
5600
0.001 23 + 02
0.003 76 + 0.00004
31

---

Published in Transactions on Machine Learning Research (07/2023)
Table 6: Shallow water 2-day predictions, vorticity stream function formulation. Rollout and one-step errors of various architectures on the shallow water equations are reported. L.2 training objective of Li et al. (2020a) is used. Summed mean-squared errors (SMSE) are obtained for 2-day predictions for the vorlicity stream function formulation and are averaged over three different random seeds. If results displayed without standard deviation, the obtained standard deviation is lower than the five digit precision minimum. The best model of each model class is highlighted.
SMSE
METHON
Trajs.
onestep
rollout
DelRosNet125
448
0.069 14 + 0.004 30
0.541 08 + 0.034 62
DelRteNNet128
5600
0.010 + 0.000 19
0.141 19 + 0.005
120-oor
448
0.036 62 + 0.000 21
0.314 18 + 0.002
DilResNee12SSnormm
5600
0.006 59 = 0.000 11
0.095 90 = 0.000
PNO128-8mocde
448
0.178 17 + 0.004 62
0.864 18 + 0.019 15
FNO128 1700128-mmodess
5600
0.039 73 + 0.000 99
0.21397 0.002
FNO1 28- 8 16des16
448
0.093 02 + 0.003 25
0.543 + 0.016 52
FNO12 28 8 modes16
5600
0.005 10 = 0.000 02
0.027 47 0.000 04
FNO64-4 4 modes33
448
0.08056- + 0.001 65
0.505 15 + 0.008 60
FNO64-4 40644-4mode3322 modes32
5600
0.005 63 + 0.000 01
0.037 19 + 09
FNO96-44 modes12
448
0.074 + 0.001 33
0.458 38 + 0.007 61
FNO96- ANO96-4modes32
5600
0.002 48 + 0.000 04
0.01481 + 0.000
FNO128 modes33
448
0.067 41 0.000 71
0.418 39 0.003 05
FNO128- 4modes32
5600
0.001 33 0.000 01
0.007 42 0.
UNO64
448
0.192 + 0.000 10
0.973 45 + 0.001 86
UNO64
5600
0.007 29 = 0.000 01
0.037 24 + 0.000 10
U-Netbase 64
448
0.030 15 = 0.000 12
0.258 04 + 0.001 29
U-Netbase 64
5600
0.001 99 + 0.000 02
0.021 40 + 0.000 17
U- N-eebas 128
448
0.021 85 + 0.000 20
0.206 + 0.001 45
U Netbasse 128
5600
0.000 84 + 0.000 02
0.010 + 0.000 49
U-Net201564
448
0.038 41 + 0.000 52
0.307 50 = 0.003:27
2015
5600
0.002 66 0.000 02
0.027 44 0.000 14
U-Netgo15 128
448
0.026 89 + 0.000 08
0.239 41 0.001 32
U-Net2015 128
5600
0.001 11 + 0.000
0.012 36 + 0.000
U-Netgo15-tamb 64
448
0.086 29 = 0.003 04
0.581 40 = 0.010 45
U-Net2015-tanni 64
5600
0.005 49 + 0.000
0.049 60 + 0.000 40
U-Netgo1s 128
448
0.059 54 + 0.000 36
0.449 91 + 0.001 58
U-N 128
5600
0.002 65
0.025 40 + 0.000
U-Netmod 64
448
0.021 55 + 0.000 44
0.201 22 + 0.004
U-Netmod 64
5600
0.000 61 + 0.000 01
0.010 22 + 0.00004
U-F2N Neoddeees
448
0.017 78
0.128 92 0.000
U-F J-F22etmoddee16,
5600
0.000 36
0.000 01
0.001 92 0.000
32

---

Published in Transactions on Machine Learning Research (07/202) 2023)
Table 7: Shallow water 1-day predictions, velocity function formulation. Rollout and one-step errors of various architectures on the shallow water equations are reported. Summed mean-squared errors (SMSE) are obtained for 1-day predictions for the velocity function formulation and are averaged over three different random seeds. If results are displayed without standard deviation, the obtained standard deviation is lower than the four digit precision minimum. The best model of cach model class is highlighted.
SMSE
МЕTHON
Trajs.
onestep
rollout
ResNet
256
0.3152 + 0.0039
2.7775 + 0.0122
1esNet118
2800
0.1329 + 0.0010
1.7711 = 0.0273
ResNet256
256
0.1827 + 0.0018
1.8015 = 0.0037
ResNet256
2800
0.1294
0.0122
1.7957 = 0.0947
DilResNet-128
256
0.1596
0.0097
1.4477 = 0.0754
DilResNet-122
2800
0.0161 0.0005
0.1900 = 0.0061
DilResNet-128-norm
256
0.0749 + 0.0026
0.6621 = 0.0301
DilResNNet-18-nommn
2800
0.0102 + 0.0002
0.1369 = 0.0009
PNN0128-8
256
0.1561
0.9394 = 0.0081
FNO128
PNO128-8mode16
2800
0.0122 + 0.0001
0.0505 = 0. 0.0004
UNO64
256
0.2732 + 0.0009
2.0338 = 0.0071
UNO64
2800
0.0135
0.0001
0.0587 = 0.0004
UNO128
256
0.2300 = 0.0039
1.7308 = 0.0173
UNO128
2800
0.0046
0.0158 + 0.0001
U-Nctbase 64
256
0.0569
0.0008
0.4457 = 0.0038
U-Nel
64
2800
0.0043
0.0316 + 0.0007
U-Netbase 128
256
0.0428 + 0.0020
0.3761 = 0.0189
U-Netbase 128
2800
0.0023 + 0.0002
0.0229 = 0.0022
U-Netmoa664
64
256
0.0272 + 0.0004
0.2678 = 0.0017
m-Ne 64
2800
0.0015 + 0.0002
0.0243 + 0.0065
mo 64
256
0.0239 + 0.0008
0.2089 = 0.0291
U-Netmoa66 64
2800
0.0014
0.0008
0.0144 = 0.01
U-F1 1odde11
256
0.0257
0.0005
0. 1638 0.0032
U-t
2800
0.0009 + 0.0004
0.0041 = 0.0019
me
256
0.0253
0.0016
0.1710 + 0.0124
U-F2Net
2800
0.0006 + 0.0002
0.0023 + 0.0009
U- U-F2Netmodess..
256
0.0248 + 0.0006
0.1659 = 0.0031
U-F2Net,nodsst
2800
0.0005
0.0022 = 0.0003
U.
256
0.0315 + 0.0007
()2398 + 0.0068
U-F2Netmudes11111
2800
0.0003
0.0013
256
0.0238
0.0016
1545 + 0.0126
U- F.t2Neeattnnodee166
2800
0.0070 + 0.0093
0.0553 = 0.0841
33

---

Published in Transactions on Machine Learning Research (07/2023)
B.4 Navier-Stokes equations.
2D Navier-Stokes data is obtained on grid with spatial resolution of 128 128 (Az = 0.25, Ay = 0.25). and temporal resolution of At = 1.5s, a viscosity parameter of // = 0.01, and buoyancy factor of (0, 0.5)*. The equation is solved on a closed domain with Dirichlet boundary conditions (v = 0) for the velocity, and Neumann boundaries 0s ör = 0 for the scalar field. We run the simulation for 21s and sample every 1.5s. Trajectories contain scalar and vector fields at different. time points. The inputs to the Navier-Stokes experiments are respective fields at the previous timesteps. Exemplary rollout trajectories are displayed in Figure 16. We outline further details on the results in Figure 17 and Table 8. Additionally, ablate different encoding/ Seallingg/eeene decoding choices for U-Net based architectures in Figures 18.
34

---

Published in Transactions on Machine Learning Research (07 2023)
Ground truth
Prediction
A error
02
(a) Scalar field
Ground truth
Prediction
Abs. error
(b) Vector field a-component
Ground truth
Prediction
Abs, error
4.70
4.00
(c) Vector field y-component
Figure 15: Navier-Stole Stokes, velocity function form. Example rollouts of the scalar and vector velocity field of the Navicr-Stokes experiments are shown, obtained by a F2096-4modess22.22 PDE surrogate model (middle), and compared to the ground truth (lop). Predictions obtained for time window Al 1.5s. The respective model input fields comprise four timesteps, only show the last of those (left-most ground truth column).
35

---

Published in Transactions on Machine Learning Research (07 /023)
Ground truth
Prediction
A error
08
(a) Scalar field
Ground truth
Prediction
Abs. error
(b) Vector field a-component
Ground truth
Prediction
Abs. error
06
04
05
(c) Vector field y-component
Figure 16: Navicr-Stokes, velocity function form. Example rollouts of the scalar and vector velocity field of the Navier-Stokes experiments are shown, obtained by a U-FINetmodeste PDE surrogate model (lop), and compared to the ground truth (bottom). Predictions are obtained for a time window At 1.5s. The respective model input fields comprise four timesteps, we only show the last of those (left-most ground truth column).
36

---

Published in Transactions on Machine Learning Research (07 (07/2023)
ResNet128
ResNet256
DilatedResNee1255
DiiatedttesNee112mmorrn UN064
FNO128
PNO128-Bmoders Bmoders
(N0128-8moddedh
FNO1 28-2 28-4m0de120
PNO64-4mededs
PN096-4modesS2
UNO128
U-Nefgo1n
U-Netao:s 128
U-
U-Nethase
U-Netmeat64
U-Netar:64
C-P1 Notmodert
U-F1Netmodect6
U-F2Netmodess,
U-F2Netmodee16.
D-F2Netmodwott
U-F2Netmader16.
C-FLNetats,mmeett
Navicr-Stokes velocity, rollout
Navicr-Stokes velocity, one-step
109
10
MSE
10a1
10
10
2080
5200
2080
5200
Num. Train Trajectories
Figure 17: Navier-Stokes velocity function form. Rollout and one-step errors of various architectures on the Navier-Stokes equations reported, obtained for predictions of 1.5s, and are averaged over three different random seeds. Note the logarithmic scale of the y-axes.
Rollout
One-step
U-Net.... 6d-3xc3
2.4
U-Netmos 64-1x1
U.Not 64-3663
U-Net, 64-1x1
10
U-P U-PiNetmotess 3363
MSB
U.PiNetmotus
10
U-3'1Net U-d'INoaaodudee--a
U-P
U-F2N 3x8
10
10
U-1 U- 41x1
U--22Ne 16.
2080
5200
Num. Train Trajectories
2080
5200
Figure 18: Navier-Stokes velocity function form. Ablation results of different encoding and decoding choices for various U-Net architectures are reported. 1 x and 3 3 kernels are compared for both encoding and decoding. Rollout and one-step errors are obtained on the Navier-Stokes equations in velocity function form and are averaged over three different random seeds.
37

---

Published in Transactions on Machine Learning Research (07/2023)
Table 8: Navier-Stokes, velocity function formulation. Rollout and one-step errors various architectures on the Navier-Stokes equations are reported. Summed mean-squared errors (SMSE) are obtained and are averaged over three different random seeds. If results are displayed without standard deviation, the obtained standard deviation is lower than the five digit precision minimum. The best model of each model class is highlighted.
SMSE
METHOD
Trajs.
onestep
rollout
R8sNet128
2080
0.01860
0.000 19
0.14002 0.003 44
ResNet128
5200
0.017 0.000 17
0.14230 0.005 98
ResNet256
2080
0.01675
0.000 32
0.143 11 + 0.009
ResNet256
5200
0.017 25 0.000 52
0.137 17 0.002 84
DilResNet-128
2080
0.002 14 + 0.000 07
0.014 60 + 0.000 14
DilResNeet-128
5200
0.001 02 0.000 01
0.007 48 0.000 04
DilRee-eet-2S--uori
2080
0.001 67 0.000 01
0.011 48 1 0.000 01
5200
0.000 88 0.000 06
0.006 0.000 43
FNO128-8mood
2080
0.014 00 + 0.000 36
0.056 96 + 0.000
FNO128- (1:0128--mmcoo 28d8m0
5200
0.008 79 0.000 12
0.038 36 0.000
FNO128-8 modesi6
2080
0.008 90 + 0.000 50
0.0 + 0.002 05
FNO128-8 moddes16
5200
0.005 10 0.000
0.025 76 0.000
FNO128- modesi6
2080
0.008 31
0.037 69 + 0.000 17
FNO128- modesid
5200
0.005 57
0.027 17 0.000 13
FNO64-
AX064-4modes32 modes32
2080
0.008 99 + 0.000 14
0.039 94 + 0.000 45
FN064-4
4-4de3322
5200
0.005 76 0.000 09
0.027 87 0.000 19
TNO96-4mood modes32
2080
0.008 43 + 0.000 19
0.036 11 T 0.000 4-1
FN096-4 #X096-4modess32 modes33
5200
0.005 07 0.000 06
0.024 14 0.000 13
UNO64
2080
0.022 00 + 0.000 91
0.083 91 + 0.001 79
UNO64
5200
0.008 37 0.000 14
0.040 10 0.000 74
UNO128
2080
0.019 33 + 0.000 03
0.073 14 + 0.000 35
UNO128
5200
0.006 71 0.000 02
0.033 75 0.000 21
U-Net2015-tam 61
2080
0.006 51 + 0.000
0.033 + 0.000 01
U-Net2015- 64
5200
0.003 59
0.000
0.02051 0.000
U-Net2o15-tanh 128
2080
0.004 56 0.000 07
0.025 27 + 0.000 43
U-Nel2015-tanb 128
5200
0.002 75
0.01671
0.000 09
2080
0.003 56 + 0.000 02
0.020 + 0.000 13
64
5200
0.002 26
0.01386 0.000 04
U-Net201a 128
2080
0.002 64 + 0.000 05
0.015 + 0.000 30
U- U-Netao15 128
5200
0.001 65 0.000 02
0.01076 0.000 07
N- 64
2080
0.003 44 = 0,000 06
0.019 10 + 0.000 17
U-Netraa 64
5200
0.001 97 0.000 03
0.011 97 0.000 08
U- U-e 128
2080
0.002 35 = 0,000 01
0.013 83 + 0.000 03
U- U-Net base 128
5200
0.001 42 0.000 01
0.008
0.000 13
U-Netmood61 64
2080
0.001 66 1 0.000 01
0.010 53 1 0.000 06
U-Netmod 64
5200
0.000 88 0.000
01
0.006 21 0.000 08
64
2080
0.001 63 1 0.000 03
0.010 48 1 0.000 22
64
5200
0.000 90 0.000 03
0.006
0.000 18
U-F1 U-FINet
2080
0.001 57 0.000 02
0.010 10 + 0.000 15
U-FI Netmodeas
5200
0.000 92 0.000 04
0.006 39 0.000 28
U-F1 NFF medeest1
2080
0.001 55 0.000 02
0.010 12 + 0.000 19
U-F1 Netmodes16
5200
0.000 88 0.000 01
0.006 21 0.000 01
H-F22eetnodesd
2080
0.001 87 0.000
0.011 67 0.000 08
5200
0.001 08 = 0.000 07
0.007 32 = 0.000 53
U-F2 U-F2Net modees16,
2080
0.002
0.01261 0.000 03
5200
0.001 13 0.000 02
0.007 = 0.000 11
U- U-F2Netmodostt,
2080
0.001 97 + 0.000 03
0.01227 + 0.000 08
U-F2Netmoddess.
5200
0.001 13 0.000 03
0.007 66 + 0.000 20
U-F2Netmoodes1,.1e
2080
0.002 29 0.000 03
0.013 88 + 0.000
5200
0.001 18
0.007 90 0.000 04
U- U-F2Netate F.F2eeatt.mddest...
2080
0.001 99 0.000 12
0.012 45 0.000
U- F W.FFeeettt.maa. moddmeees..
5200
0.001 07 0.000 09
0.007 32 + 5.000
38

---

Published in Transactions on Machine Learning Research (07 (07/202) 2022)
Table 9: Navier-Stokes, velocity function formulation. L2 training objective of Li et al. (2020a) is used. Rollout and one-step errors of various architectures on the Navier-Stokes equations are reported. Summed mean-squared errors (SMSE) are obtained and are averaged over three different random seeds. If results are displayed without standard deviation, the obtained standard deviation lower than the five digit precision minimum. The best model of each model class highlighted.
SMSE
METHOD
Trajs.
onestep
rollout
DilResNet
2080
0.002 06 0.000 01
0.013 96 0.000 05
DilResNett123
5200
0.001 02 + 0.000 03
0.007 49 + 0.000 12
121ResNee118-norm
2080
0.001 47 0.000
0.010 11 0.000
DilResNet12S-non
5200
0.000 82
0.006 04
FN0128-8
2080
0.012 49
0.051 87
FNO128 170128-88 modes8
5200
0.008 05
0.035 91
FNO1 28 8 16
2080
0.008 23
0.038 65
FNO128-8 1
5200
0.004 84
0.024 85
PNO64-4modees
2080
0.007 81 + 0.000 14
0.035 94 + 0.000
FNO64-6 modrss3
5200
0.005 17 0.000
0.025 56 0.000
FNO96 modes32
2080
0.007 45 + 0.000 15
0.032 42 0.000
FNO96- modess3
5200
0.004 70 + 0.000 06
0.022 27 0.000 11
FNO128 modess32
2080
0.007 10 + 0.000
0.030 80
FNO128 modes32
5200
0.004 43
0.020 76
UN064
2080
0.017 07
0.070 89
UNO64
5200
0.007 25
0.035 96
U-Netpase 6
2080
0.003 11 0.000 01
0.017 74 + 0.000
base 64
5200
0.002 02 + 0.000 01
0.012 22
U-Nebbase 128
2080
0.002 12 + 0.000 02
0.012 + 0.
U- N-Netiase 128
5200
0.001 37
0.008 68 0.000 04
U-Net201564
2080
0.003 24 0.000 07
0.01861 0.000 44
U-Net201564
5200
0.002 23 + 0.000 01
0.013 76 + 0.000 15
U-Netgo15 128
2080
0.002 38 + 0.000 02
0.014 46 + 0.000
U U-Net2o15 128
5200
0.001 61 + 0.000
0.010 47 + 0.000
U-Netmod64 64
2080
0.001 51
0.009 75 + 0.000
U- 64
5200
0.000 91 + 0.000
0.00637 + 0.000
39

---

Published in Transactions on Machine Learning Research (07 (07/2023)
One-step (At = 0.375s)
Oue-step (At = 0.75
10
U-Netmog
U-Netar 64-Add
U-Nelar 64-AdaGNN
U- L-FINehaaodeee--e
I AdarNN
10
10
U- P2N cfmodde30, AdaGG
U- PTN mhart.mcdells Add
U- PTNN clatt deatt
L- F2N wlatt aede
One-step (Ar 1.58)
One-step (Ar = 3s)
10"22
MSE
One-step (At 68)
ΣA
10
10-3
1604
6656
1064
6056
Num. Train Trajectories
Figure 19: Navier-Stokes parameter conditioning experiments. Ablation results for different parameter con- ditioning methods. " "Addition") " "(Add) and " "AdaGN" are compared one-step errors, reported for different time windows and averaged over 208 unseen values of the buoyancy force term For low time windows, AdaGN seems to be beneficial.
B.5 Parameter conditioning.
We use the same spatial resolutions and boundary conditions as doscribed in Section B.4. The inputs to the Navier-Stokes parameter conditioning experiments are respective fields at the previous timestep. Exemplary rollout trajectories predicted by a single surrogate model for different buoyancy force values are displayed in Figures 20,21,22, We outline further details the results in Table 10. Additionally, ablate different parameter conditioning choices in Figure 19, namely "Addition" versus " "AdaGN" for U-Net blocks, and "Adittion" vs "Spatial-Spectral" for Fourier blocks. The default choice is "Addition..
40

---

Published in Transactions on Machine Learning Research (07 (07/2023)
Ground truth
Prediction
Abs. error
(a) Scalar field
Ground truth
Prediction
Abs. emor
(b) Vector field r-component
Ground truth
Prediction
Abs. error
(c) Vector field y-component.
Figure 20: Parameter conditioning for Navier-Stokes equations, f = 0.21. Example rollouts of the scalar and vector velocity field of the Navier-Stokes experiments shown, obtained by U-Netmod PDE surrogate model (top), and compared to the ground truth (bottom). Predictions are obtained for time window At = 1.5s and a buoyancy force term of f = 0.21. Model inputs are respective fields at the last timestep (left-most ground truth column).
41

---

Published in Transactions on Machine Learning Research (07 (07/)023)
Ground truth
Prediction
Abs. error
0.4
(a) Scalar field
Ground truth
Prediction
Abs. emor
(b) Vector field r-component
Ground truth
Prediction
Abs. error
(c) Vector field y-component.
Figure 21: Parameter conditioning for Navier-Stoker equations, f 0.33. Example rollouts of the scalar and vector velocity field of the Navier-Stokes experiments shown, obtained by U-Netmod PDE surrogate model (lop), and compared to the ground truth (bottom). Predictions are obtained at time window At = 1.5s and a buoyancy force term of f = 0.33. Model inputs are respective fields at the last timestep (left-most ground truth column).
42

---

Published in Transactions on Machine Learning Research (07/202) 2023)
Ground truth
Prediction
Abs. error
95
95
(a) Scalar field
Groond truth
Prediction
Abs. ciror
14
(b) Vector field z-component
Ground truth
Prediction
05
Abs. cmoE
(c) Vector field y-component.
Figure 22: Parameter conditioning for Navier-Stokes equations, f 0.48. Example rollouts of the scalar and vector velocity field of the Navier-Stokes experiments shown, obtained by U-Netmod PDE surrogate model (top), and compared to the ground truth (bottom). Predictions are obtained at time window At 1.5s and a buoyancy force term of f = 0.48. Model inputs are respective fields at the last timestep (left-most ground truth column).
43

---

Published in Transactions on Machine Learning Research (07/2023) 2023)
Table 10: Parameter conditioning on the Navier-Stokes equation, velocity function formulation. Summed mean-squared errors of various architectures are reported for different number of training trajectories, and different time windows. Conditioning results at different time windows are averaged over 208 unseen values of the buoyancy force term. The best model of cach model class is highlighted. Different parameter conditioning choices are ablated, namely " "Addition" versus "AdaGN" for U-Net blocks, and " "Addition" vs "Spatial- Spectral" for Fourier blocks. The default choice is "Addition".
SMSE
METHOD
Trajs.
0.375
0.75
1.5s
3.0s
6.0
FNO128 1X0112modes16 16
1664
0.005 17
0.006 93
0.011 73
0.026 66
0.064 23
FNO128 28 modes16
6656
0.003 88
0.004 83
0.007 40
0.015 44
0.041 77
FNO128 28 16 modee SpaSpec
1664
0.004 62
0.006 67
0.014 04
0.038 34
0.076 48
FNO128 1o SpaSpec
6656
0.003 48
0.004 55
0.008 01
0.019 73
0.052 08
U-Netmod 64
1664
0.000 72
0.001 11
0.002 16
0.006 22
0.028 05
Uodetmod 64
6656
0.000
0.000 61
0.001 15
0.003 25
0.015 60
U-Netmod A4-AdAGN
1664
0.000 59
0.000 90
0.001 77
0.005 68
0.029 89
U-Netmod A4-Ad
6656
0.000
0.000 50
0.001 00
0.003 00
0.01632
U-Net 64
1664
0.000 82
0.001 22
0.002 34
0.006 30
0.028 46
U-Netan 64
6656
0.000 46
0.000 69
0.001 30
0.003 56
0.016 75
u-Net A4-AdaGN
1664
65
0.001 01
0.001 95
0.005 97
0.031 60
U-Net ats 6 64-AdaGN
6656
0.000 35
0.000 55
0.001 06
0.003 10
0.017 03
U-FINet
1664
0.001 38
0.001 80
0.002 75
0.006 45
0.024 97
6656
0.000 54
0.000 78
0.001 33
0.003 38
0.015 04
U-F INetmodesse -SSaSSpe
1664
0.000 76
0.001 16
0.002 42
0.007 70
0.034 95
U-F LFFLNetmodes1e SSpSSpec
6656
0.000 40
0.000 61
0.001 16
0.003 39
0.017 54
U-FiNetsst modess6
1664
0.001 42
0.002 05
0.003 33
0.007 62
0.027 95
U-F1Nelate modes16
6656
0.000
0.000 86
0.001 46
0.003 58
0.015 79
1-FINetare modest SpaSpec
1664
0.00095
0.001
0.002
0.007 67
0.031 93
FFNNetatt, 16- SpaSpec
6656
0.000 46
0.000 72
0.001 32
0.003 72
0.017 86
U-F W.FFeetatt,modeert AdaGN
1664
0.001
0.002 00
0.003 11
0.007 48
0.028 79
- AAANN
6656
6.
0.000 86
0.001 37
0.003 47
0.015 96
L-FFNeehmodeele-AAAACCC
1664
0.001 61
0.002 12
0.003
0.007 58
0.027 69
TLFiNetmoddele-AAaGG
6656
0.000 55
0.000 77
0.001 32
0.003 38
0.015 64
m-
1664
0.001 78
0.002 40
0.003
0.008 50
0.033 46
T-F2Nelmoddese,,
6656
9.
0.001 16
0.001 68
0.003 76
0.015 64
U-F2Net Spaappe
1664
0.000 71
0.001 18
0.002 33
0.007 26
0.032 87
U-F2 L SSpaeeeeeeeeSSSpepeeeeeeee
6656
0.000 39
0.000 64
0.001 22
0.003 68
0.018 62
U-F2Netmodes16, AdaGN AdaGN
1664
0.001 86
0.002 89
0.00482
0.013 07
0.047 67
U-F2Netmodesse. AANNN
6656
0.000
0.001 24
0.001 85
0.004 18
0.018 09
U-F2Netatt. 10de
1664
0.00197
0.00291
0.00464
0.011 29
0.038 10
10des
6656
0.001 16
0.001 53
0.002 25
0.004 90
0.018 88
U-F L.F2Netaat.modeeee...
Spaappe SpaSpec
1664
8.000
0.001 35
0.002 78
0.008 43
0.034 90
Ne Spec
6656
0.000 44
0.000 71
0.001 37
0.004 04
0.018 96
U- L-F2Netatt, modeet, AdaGN
1664
0.001 96
0.002 81
29
0.010 19
0.036 73
U-F2Nelatt
1odes16.6 modosell..
AdaGN
6656
9.000
0.001 21
0.001 79
0.004 14
0.01705
44


Reviews:
[{"review_id": "Owwix4muUd", "reviewer": "Reviewer_rik3", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "This submission contains mostly introductory content, and I cannot see much innovations. Section 2 is an introductory section. The content before 3.1 introduces three architectures. The only new architecture is Fourier U-Net. Section 3.1 is also like an introductory section. \n\nThe experiment section does not seem to give definite conclusions. While this section provide rich experiment results, there is not much insights. For example, what can be reasons for different performances. \n\nOne question: has any of these benchmark problems been used by previous work? If so, then it is nice to include results from previous work. ", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This submission makes a conceptual comparison between Fourier transform and convolution. It also conducted a series of experiments to compare learning models including Fourier Neural Operators (FNO), U-Net, ResNet, and Fourier U-Net. ", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "No", "audience_explanation": "", "requested_changes": "If this submission aims to improve over previous models, the submission needs to include sufficient innovations. If this paper plan to provide a benchmark of different methods, then it needs to include more methods and provide better analysis of different methods. ", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "YRVVVNzRPR", "reviewer": "Reviewer_RNE6", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**\n- The contribution of the benchmark is a great one for the community, which tends to be either too specific or too broad in its comparisons.\n- The clarity in general is very good, and one can readily use this paper to dive in the research field.\n\n**Weaknesses**\n- The experiments section, while quite detailed, is unstructured in my opinion, which is a pity given the appeal of the work. As a consequence, the paper currently reads as a documentation for a benchmark rather than a paper that has detailed empirical findings and insights other than some big ones (that, for instance, FNO is perhaps not as generalizable as a standard U-Net).\n  - I am missing clear, structured, and systematic experiments, where 'one variable is changed at a time' such that  gradual conclusions and insights are derived, that lead to the overall conclusion. Currently, the writing of the experiments has a more free-form flow, which is in fact quite comfortable, but hard to 'map'.\n  - Since this is an empirical study, I am missing clear ablations that highlight specific aspects of the comparisons, the baselines, the data. For instance, do we need more or fewer layers with PDE-modelling? Is this somehow related to the complexity of the PDE, how chaotic it is (eg in terms of Raynolds number)? Do we need bigger or smaller latent features for U-Nets? Is regularization necessary, since the domain appearance changes over time? Are frequency modelling with FNO blocks more important than depth in Figure 4? (and many more questions like that) Currently, there are experiments like that, but they are not clearly organized. The experiment regarding conditioning on parameters is in a good direction.\n  - Not all empirical findings are sufficiently explained. Why are dilated convolutions work so much better for Navier-Stokes than shallow water (compared to other methods)? Perhaps Navier-Stokes requires larger 'context', for which convolutional filters of varying size are important. That said, and considering the claim that dilated convolutions are a good next step, in my experience in computer vision, this type of adaptability does not really pay of all that much (although it can be helpful). Rather, even having multiple filters per layer of different sizes/scales is good enough, simpler to implement, and likely more sample efficient (like in Inception architectures). Perhaps it makes sense to include a baseline with multi-size filters as well.\n  - Since this is an empirical study using state-of-the-art computer vision architectures, it is rather odd to leave transformers out. Truth be told, transformers are mentioned in the limitations of the paper. However, given its claims, the paper will not be complete without a transformer baseline, even a simple one.\n  -  Since the paper is about 'Generalized PDE-Modelling', I believe that using only a singe 'PDE dataset' and variants, is not sufficient. Surely, Navier-Stokes is the standard PDE to solve but I believe it is not enough to claim generality. Put otherwise, how certain are the authors that the reported findings would generalize to different PDEs? And can a third reader realistically derive these type of insights from just one dataset.\n- The clarity at places should be improved. Some examples in (almost) page order.\n  - 'Local and global information is therefore treated in a more distributed fashion than in FNO like architectures.' $\\rightarrow$ what does distributed mean here? I suppose spatial distribution (because one can also think of frequency distribution)?\n  - 'Common practice is to use rather small filters' $\\rightarrow$ This is common practice in computer vision, the idea being that with smaller filters one can build deeper non-linear networks, and the cascade of filters anyways leads to implicitly larger receptive fields. Is this the same for PDEs, and if yes why? I can also imagine that in PDEs the low-level patterns and textures are simpler compared to natural images, thus larger filters might be sufficient (in that shallower networks would be good enough).\n  - 'We hypothesize that replacing convolutions with FNO layers which set Fourier modes above cut-off frequencies to zero might be advantageous, especially in the lower parts of the downsampling blocks of U-Net architectures' $\\rightarrow$ While intuitively I think I understand what the authors mean (that near the bottleneck layers the learned features reflect more global/high-level patters, for which likely higher-frequencies associated with low-level textural patterns are not important), it is good if the paper motivates the hypothesis.\n  - Can you explain why you include the 'scalar vorticity stream function form' for the Navier-Stokes? It might be used more in the PDE world, however, it is not clear what is the difference from a machine learning perspective. Does it lead to different types of solutions for the PDE?\n  - Can you clarify what do you mean here 'The chosen data sets to do so consist of solution pairs u,u′ ∈ U where the pair itself is from the same solution space U , but different pairs {u,u′}1 and {u,u′}2 are from different solution spaces U1 and U2 characterized by different force terms.' Intuitively, I understand what you want to say, but I find the text unnecessarily complex.\n  - End of page 7: 'ater' $\\rightarrow$ 'water'\n  - 'In general, we observe that conditioning is more difficult for FNO layers, most strikingly seen in the performance curves of FNOs.' $\\rightarrow$ Isn't this sort of expected, since one combines information from different domains/geometries. FNO encodes frequency information, while the conditioning variable would (I suppose) be better seen as the spatial/time domain.\n - Colors in the bar charts are not easy to distinguish.\n - By rollout and one-step loss, do you use the true next state of the PDE when making the prediction (like in curriculum learning, only for the test time unrolling here)? Or you assume no access to the background simulator and the model relies only on its own predictions? If the latter is the case, I would expect an ablation also on the error accumulation and the parameters that are influenced by (eg the type of PDE etc). If the former is the case, why is this setting chosen and is it realistic to expect to have a simulator running in the background? It would still make sense when examining real observations (not simulated data), however, the study focuses on synthetic simulated data.\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper focuses on the modelling of partial differential equations (PDE) using neural network surrogates, trained to predict the next state of the PDE $t$ time steps ahead. Specifically, the paper is a experimental review of several popular baselines either from computer vision (like U-Net) or from the neural PDE modelling literature (like FNO). For this, the paper proposes a new benchmark framework for future apple-to-apple comparisons in the literature, focusing on Navier-Stokes and derivate problems. The writing is generally good, and the text is certainly clear for the most part of the paper. The structure of the paper is ok, but it could be improved at places (details below). Certainly, this is a valuable effort and one that is needed for the field of neural PDE modelling to move forward.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "No", "audience_explanation": "", "requested_changes": "Summarizing the weakness points, the requested changes are:\n\n- Consider extending your datasets beyond Navier-Stokes and variants. Since the paper is about 'Generalized PDE-Modelling', having 2 or even better 3 different PDEs is reasonable.\n\n- Improve the structure of the experiments section.\n  - Please, organize them in separate ablation studies in subsections.\n  - Each ablation should be clearly separable and offer clear and convincing conclusions and insights. In combination, the studies should lead to the final conclusion of the paper, since this is an empirical study.\n  - The experimental settings should change in a coherent and motivated fashion. Rather than changing the $\\Delta t$ or other hyperparameters in an *ad hoc* fashion, use the same parameters, use a range of parameters, or motivate clearly your current use of parameters.\n\n- Add transformers or an equivalent family of architectures, if anything, to give a fair baseline for future comparisons.\n\n- Improve clarity in the parts of the text and the visualizations referenced in the weakness section.\n\n", "broader_impact_concerns": "Since this paper has the ambition to be a widely-used benchmark, I do think it makes sense to put some effort on ablations, on better structure, and clear empirical findings. To this end, while I understand that most of my suggestions would not change the paper organically (as in, the conclusions would not change), I believe they are important for the kind of paper the authors aspire to.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "BgGJZGw8Y9", "reviewer": "Reviewer_WYLT", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The paper contains a rigorous experimental comparison of U-Nets and FNOs for PDE modelling. It is important for the field that such comparisons are done, because papers by authors proposing new methods are not always unbiassed and rigorous. Another strength is that the new architecture is apparently quite effective. \n\nThe paper does not contain a big new idea, but in my view not every paper needs to have one.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "PDEs are ubiquitous in science and engineering, but solving them can be expensive. Neural surrogates have gained increasing interest over recent years for their ability to speed up computations. It is important for these networks to be able to recognize complex multi-scale spatio-temporal patterns. This paper benchmarks two kinds of approaches, namely the Fourier Neural Operator, U-Nets and ResNets. They also import recent improvements to the U-Net architecture from the generative modelling literature. \n\nThe results show that already a plain U-Net performs better than FNOs on the tested tasks. The recent improvements to U-Net further improve results. The authors show that their best methods can generalize to different PDE parameters and time scales, which is generally considered challenging. \n\nFinally, the authors contribute their code, which should help facilitate research and benchmarking in this research area.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "No changes requested", "broader_impact_concerns": "No concerns specific to this paper", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The reviews were somewhat lukewarm, with the authors acknowledging that there is no single "big idea" in the paper - the main contribution is broad experimental comparison of reasonable baselines and documentation of various empirical observations.  Experiments fall short of providing a deeper understanding of architecture and data-complexity as a function of dynamics complexity (e.g., if a system is near chaotic, do we need deeper/larger architectures and larger datasets); the analysis is restricted to a only a few experimental settings; and wildly popular architectures like Transformers are missing from the analysis.

As the same time, TMLR evaluation criteria emphasizes accurate+clear evidence of claims, together with relevance to some TMLR audience. The paper satisfies these criteria, within the narrow scope of the claims it makes. 

Minor revisions: Please address the final reviewer remarks on lack of connections with previous work, with more comprehensive coverage in related works. Consider reporting a transformer baseline for completeness, though this is not a requirement.

Decision: Accept (poster)