Towards Stability of Autoregressive Neural Operators
Authors: Michael McCabe
Peter Harrington
Shashank Subramanian
Jed Brown
Venue: Accepted by TMLR ()

Abstract: Neural operators have proven to be a promising approach for modeling spatiotemporal systems in the physical sciences. However, training these models for large systems can be quite challenging as they incur significant computational and memory expense---these systems are often forced to rely on autoregressive time-stepping of the neural network to predict future temporal states. While this is effective in managing costs, it can lead to uncontrolled error growth over time and eventual instability. We analyze the sources of this autoregressive error growth using prototypical neural operator models for physical systems and explore ways to mitigate it. We introduce architectural and application-specific improvements that allow for careful control of instability-inducing operations within these models without inflating the compute/memory expense. We present results on several scientific systems that include Navier-Stokes fluid flow, rotating shallow water, and a high-resolution global weather forecasting system. We demonstrate that applying our design principles to neural operators leads to significantly lower errors for long-term forecasts as well as longer time horizons without qualitative signs of divergence compared to the original models for these systems. We open-source our code for reproducibility.

Published in Transactions on Machine Learning Research (11/2023)
Towards Stability of Autoregressive Neural Operators
Michael McCabe'
michael www.ee/mmoc.aee@oooooooooo.oooo
Deportment of Computer Science University of Colorado Boulder
Peter Harrington'
pharrington@lbl.go
Laurence Berkeley National Labovatory
Shashank Subramanian'
shasharisabrrmmmiiinnBl..gg
Laurence Berkeley National Laboratory
Jed Brown
jed
Department of Computer Science
University of Colorado Boulder
Reviewed on OpenReview: https: // openreview. net/ forun? id- RFJUUtKYOG
Abstract
Neural operators have proven to a promising approach for modeling spatiotemporal sys- tems the physical sciences. However, training these models for large systems can be quite challenging as they incur significant computational and memory expense these systems are often forced rely on autoregressive tine-stepping of the neural network to predict future temporal states. While this is effective in managing costs, it can lead to uncontrolled error growth over time and eventual instability. We analyze the sources of this autoregressive er- ror growth using prototypical neural operator models for physical systems and explore ways to mitigate it. We introduce architectural and application-spedt for careful control instability-inducing operations within these models without inflating the compute memory expense. We present results on several scientific systems that include Navier-Stokes fluid flow, rotating shallow water, and high-resolution global weather fore- casting system. We demonstrate that applying our design principles to neural operators leads to significantly lower errors for long-term forecasts as well as longer time horizons without qualitative signs of divergence compared to the original models for these systems. We open-source our code for reproducibility.
improvements that allow
1 Introduction
There is an increasing interest in using neural networks to simulate scientific systems described by partial differential equations (PDEs), with a specific focus on neural operators. These are neural network models that learn the mapping betwoen two function spaces that represent the inputs and solutions to a PDE through training using finite collection of input-solution pairs of data. (Li et al., 2021b; Lu et al., 2021, and references therein). These models have demonstrated great success in simulating PDE systems across a broad range of scientific disciplines. With the success of large neural network models such foundation models (Bommasani et al., 2021; Devlin et al., 2018) in natural language processing (NLP) and computer vision (CV), recent efforts have atlempted to scale these neural operators for grand challenge problems in scientifie computing, particularly in geophysical systems like weather and climate. These methods have demonstrated impressive performance in predicting crucial physical variables and are comparable to traditional numerical methods over short time horizons under certain metrics, while enjoying orders magnitude improvement in computational
"This work conducted in part while author intern Lawrence Berkeley National Laboratory.
Equal contribution.

---

Published in Transactions on Machine Learning Research (11/2023)
expense in operational (inference) setting. For instance, in weather forecasting, many models (Pathak et al., 2022; Bi et al., 2022; Lam et al., 2022) are able to achieve (or even beat) the deterministic accuracy of numerical methods with x faster timestto-solutinn
Scaling neural operators to these domains, however, presents unique challenges. First, even at modestly high resolutions (such as a 25 km grid over the earth for geophysical systems), a single physical state variable, represented as an image channel, consists of more than a million grid points (pixels). This is about 20x larger than standard CV image datascts such ImageNet (Deng et al., 2009). For some systems, there may be hundreds of such state variables represented at each grid point (e.g., Hersbach et al., 2020). The high dimensionality places strict limitations on architectures that scale poorly with resolution and/or channel dimensions.
To account for scaling difficulties while making spatiotemporal predictions, most models predict temporal soquences of physical states an autoregressive manner. Prior efforts across a wide set of spatiotemporal forecasting problems have observed an inherent trade-off in autoregressive lorecasting: smaller lime-steps typically result in an easier one-step task, but total error grows with an increasing number of steps even over fixed time window (Bi et al., 2022; Tran et al., 2023; Li et al., 2021a; Stachenfeld et al., 2022; Krishnapriyan et al., 2022). Methods to reduce this autoregressive error growth to enable either longer forecasts or smaller step sizes have included using multiple models trained at different timescales to reduce the total number of steps (Bi et al., 2022), learned step sizes (Tran et al., 2023), dissipative priors (Li et al., 2021a), spectral disretization in time (Li et al., 2021b), noise injection (Stachenfeld et al., 2022), and multi-step rollouts during training (Lam et al., 2022; Pathak et al., 2022). These approaches have shown promise but many drastically increase training cost. (multi-step, multi-model, spectral-in-time), add additional hyperparameters that must be tuned (noise injection), are only applicable in certain cases (dissipative priors).
In this paper, we critically analyze sources of this autoregressive error growth. We highlight geophysical applications as carth systems are both large in scale and often demand long-range forecasts of physical state variables for example, forecasts of the atmospheric state (wind, temperature, precipitation, etc.) are commonly run for several years to understand the atmosphere's global behavior, particularly with a changing climate. Autoregressive error corruption renders such forecasts inadequate for any scientific analyses and it is critical to develop methods to mitigate this. We focus on models derived from prototypical neural operators such as variants of the Fourier Neural Operator (FNO) (Li et al., 2021b) which have demonstrated promising results on various scientific applications that include earth systems (Pathak et al., 2022). These models utilize spectral transformations to implement global operations in a compute-efficient manner for example, the FNO implements global convolutions through fast Fourier transforms (FFTss. We suggest novel changes to their model architectures by borrowing ideas from classical numerical methods such as pseudospectral schemes for solving nonlinear time-dependent PDEss as well as application-specife modifications that address instabilities that arise due to the structure of the problem.
Specifically, our contributions are as follows:
1. Connecting sources of instability and numerical analysis. We draw parallels between insta- bilities that arise in autoregressive spatiotelpo forecasting of physical stale fields and slandard numerical analysis through the lens of pseudospectral numerical schemes and demonstrate that au- toregressive spatiotemporal models show signs of aliasing and numerical instability, contributing to nonlinear error growth and divergence (much like in pscudospectral mumerical methods).
2. Controlling error growth through model and application-speeific innovations. We intro- duce several modifications to our architectures to control the above error growth without inflating our computational pipputaationalmmmorry /mmmory expense. These include (1) frequency-domain spectral normalization to control the sensitivity of spectral convolutions, (2) deptawisse-epprrable spectral convolutions for ef- ficient parameter usage, (3) adapting the classical Double Fourier Sphere (DFS) method to represent functions defined on a spherical geometry as a torus for geophysical applications, and (4) data-driven spectral convolution.
3. Boosting long-range autoregrossive performance. We demonstrate the efficacy of our pro- posals by applying them in off- off-the-shelf fashion to standard neural operators such as the FNO and its variants. Our experiments across multiple physical systems including Navier-Stokes benchmark
2

---

Published in Transactions on Machine Learning Research (11/2023)
fluid simulations, shallow water equations on rotating sphere, and weather forecasting demonstrate improved stability and long-term accuracy. The last system is especially challenging problem that models the high resolution spatiotemporal patterns of atmospheric variables on planetary scale with broad scientific and socictal implications. In this setting, even for particularly volatile fields such as surface wind, we show that our improvements enable up to 800% longer instability-free forecasts, paving the way for long-range predictions using neural operators. Models for smoother problems like shallow water using our modifications show no signs of instability at any point in our experiments.
Outline. We first provide some background on the connection between pseudospectral methods and our observed sources of instability in neural operator models in $2. We then explore these sources of error in the context of neural operators in $3, introduce our innovations to overcome these in $4, and finally demonstrate the impact of our innovations across different physical systems in $5.
2 Background
2.1 Fourier Pseudospectral Method
The standard approach for introducing neural operators explores the connection between neural operator methods and integral methods for the solution of linear PDEs (Li et al., 2021b). However, these methods also have strong mechanistic similarity to pscudospectral methods used for solving time-dependent nonlinear PDEs (Orszag, 1970; Eliasen et al., 1970) use this alternative perspective to motivate their construction as the connection exposes sources of numerical instability in these models.
Spectral mothods approximate a function u(x,t) t) defined over space and time by a lincar combination of K basis functions (x, = -E U=(t))x(z)) (x) where ΦR is the kth element of the set of basis functions and (t) is the corresponding time dependent coefficient. Bases are chosen to allow for efficient numerical algorithms while offering rapidly decreasing truncation error T []x[u]²(( (t) = >i- U2(t) as K increases. A function that can be represented with no truncation error for finite K called bandlimäted. In this work, unless explicitly stated otherwise, we will be referring to the Fourier basis (x()) = e We note that this choice of basis implicitly assumes periodic boundary conditions which will be discussed later in the context of earth systems. Pseudospectral methods differ from pure spectral methods in that they augment the spectral representation with representation in the spatial domain sampled at N points (z1....,NN) 1996) such that now both our spectral and spatial representations consist of discrete sets of elements. This property, along with the existence of fast spectral transforms, enables efficient computation of nonlinear terms.
EN) (Fornberg,
Relying on discrete transforms introduces complications that must be accounted for in pseudospectral meth- ods. Instead of computing true basis coefficients (1) Sa u(x,, (z))(()) dx, we instead compute the finite approximation (t) Em (,, Different basis functions may be indistinguishable when sampled at a finite number of points and so the estimated coefficients become the sum of the true coefficients of these indistinguishable basis functions. This phenomenon is called aliasing and the unresolvable modes are called aliases. For the Discrete Fourier Transform (DFT), modes above the Nyquist frequency of N/2 are aliases of lower frequency modes (Smith, 2007). The discrete aliasing error for the DFT can be defined as A [u[ - (E) (t)) Aliasing is often introduced by the application of nonlinear transformations applied the spatial domain through a process described below.
To explore this, we examine the following one-dimensional nonlinear advection equation applied to time- varying scalar field E (T) defined on the periodic domain [0, 1]:
dula,t) u(x tu(a,t) 01
(1)
Or
Assuming that u has bandlimit K N 2 we can exactly represent u by the Fourier basis coeffiecients âs computed by the DFT. In frequency space, applying linear differential operators can be done via element- wise multiplication and applying amounts to reducing the equation to system of ODFs defined by the
3

---

Published in Transactions on Machine Learning Research (11/2023)
convolution:
Oāg (1) = (kk 1) Ot
Du (k: 1) [k], Ox
Ak E 2K}-
(2)
Computing these derivatives is where the pseudospectral method has significant advantage over spectral methods. The direct convolution in the frequency domain has cost of O(N2), but by the dual of the Convolution Theorem (Smith, 2007), we could also compute this by transforming back into the spatial domain (O(N log N)) and subsequently performing elementwise multiplication (O(N)). Once the time derivatives are computed, one can use any off-the-shelf ODE integrator to march the equation forward via the method of lines (Boyd, 2013).
While the transform approach is significantly more efficient, computing the nonlinear term in the spatial domain introduces alinsing. The bandlimit of the newly computed time derivative is up to twice that of the original function (for the full derivation, see Appendix A.1). Without increasing the spatial sampling rate, these newly generated modes will alias back onto modes below the Nyquist limit. This is true not only for polynomials as seen here, but also many commonly used nonlinear functions in deep learning, which can seen informally through the Weirstrass Approximation Theorem since our continuous nonlinearities can be arbilrarily closely approximately by infinite series of polynomials over a finile interval (Hunter Nachtergaele, 2001). Aliasing is major contributor instability for pseudospectral methods and, as we will see in later sections, in the use of autoregressive neural operators well.
2.2 Connection to Fourier Neural Operators
The pseudospectral method can be summarized as using fast transforms alternate between the spectral domain, where we quickly compute high-accuracy spatial derivatives via convolution, and the spatial domain where we can compute nonlinearities with similar efficiency. The resulting value is then added back to the original state according to an ODE integration rule. We note similar pattern in Fourier Neural Operators (FNO).
Consider one FNO block applicd to problem with one spatial dimension. Let u(() E NxD( denote the hidden state at layer & E (O....,,, L} at each discretization point such that w(0) u. We use N, K, D to denote the sizes of the spatial, frequency, and channel dimensions respectively. The block is parameterized by ( E x and R(,) C xXXD((x)x and computed as:
D
((+1) 11 Vic
wl0 (6) Uic +
(FNO
(3)
D
(FNO = Σ
R/00 (Ful kee
(4)
where F denotes the Discrete Fourier Transform and h denotes arbitary nonlinear activation function. domains respectively. In subsequent sections we will
Note that we use i, k to index the refer to the first term inside the noulinearity of Equation 3 as a pointwise linear operation (or equivalently Ixl convolution) and the second term, which is expanded in Equation 4, as the spectral convolution.
For the purposes of our later discussion on aliasing mitigation, the key conncction between the FNO and pseudospectral methods is the fact that both methods repeatedly alternate between the spectral domain where they compute large-kernel convolutions and the spatial domain where they apply nonlinearities. In both cases, the nonlincarity in the spatial domain potentially introduces aliasing while the spectral trans- formation offers opportunities for efficiently mitigating this aliasing. This connection allows us to explore solutions from the pseudospectral literature for stability problems faced by the FNO. Deeper connections in terms of functional form are derived for interested readers in Appendix A.2.
Section 3.1 empirically explores the consequences of this shared aliasing behavior while Section 4.2 proposes mitigation strategies.
4

---

Published in Transactions on Machine Learning Research (11/2023)
2.3 Related Work
Aliasing in Deep Learning. While aliasing in deep neural networks has previously been studied in classification and generative models Zhang (2019); Zou et al. (2020); Vasconcelos et al. (2021); Ribeiro & Schön (2021); Karras et al. (2021), there has been a surprisingly small amount of exploration in the context of spatiotemporal forecasting. Zhang (2019) previously showed that aliasing allows convolutional architectures to learn non-shift equivariant features. Zou et al. (2020), Ribeiro & Schön (2021), and Vasconcelos et al. (2021) demonstrate in multiple settings that while low-pass filtering is valuable to improving model accuracy, CNNs typically do not allocate capacity to do so via normal training procedures. In the generative space, Karras et al. (2021) observed significant improvements to consistency in GAN-generated data by employing anti-aliasing measures. Within the context of neural operators, Fanaskov & Oscledets (2022) showed that aliasing caused by ReLU nonlinearities can lead to significant error. Their fully spectral solution avoids aliasing, but scales poorly due to O(K²D) complexity in frequency mode mixing.
Fourier Neural Operators. Fourier neural operators and their variants (Li ct al., 2021b) have been a popular and successful class of neural networks used for solving PDEs. They have been shown to be efficient universal approximators for the solution of PDE systems (Kovachki et al., 2021; De Ryck & Mishra, 2022) and have been successfully applied to scientific computing problems in a variety of challenging domains (Pathak et al., 2022; Wen et al., 2022; Li et al., 2022; Guibas et al., 2021; Guan et al., 2021; Yin et al., 2022; Witte et al., 2022; Li et al., 2021a; Tran et al., 2023; Fanaskov & Oseledets, 2022), particularly those related to fluid dynamics and earth systems. Our focus in this paper is on stabilizing and improving the autoregressive use of neural operators for spatiotemporal problems (also referred to as Markov Neural Operators in Li et al., 2021a), which is a necessity for applying these methods to large-scale systems. The Markov Neural Operator of Li et al. (2021a) is shown to reproduce statistical properties of the attractor in dissipative chaotic dynamical systems, allowing long-term autoregressive forecasting via use of Sobolev losses and dissipative data augmentation. Here, we also aim to achieve stable trajectories with autoregressive FNOs, but target systems with external, non-constant forcings and focus on architecture-specffic sources of error. While we showcase the FNO architecture and its variants due to their connection with pscudospectral PDE mothods, many of our contributions can be applied to any neural network used for autoregressive temporal forecasting.
Spherical Architectures. One motivation for addressing stability in a scalable manner is to apply these tools to carth systems. In these scttings, we must also account for spherical geometry. We address this in a lightweight manner in FNO via use of the DFS method, but other neural archilectures designed for spherical geometry (Cohen et al., 2018; Esteves et al., 2017: Defferrard et al., 2020) are relevant to our discussions, and we directly compare our model against some of these in Section 5.1. We note that in order to to make the cost of handling spherical geometry minimal, our DFS approach does not seek to obtain exact SO(3) equivariance as found in the SO(3) or spherical convolutions of of Cohen et al. (2018) or Esteves et al. (2017).
3
Where Autoregressive Neural Operators Fail
3.1 Aliasing and Unbounded
Error Growth
We begin by demonstrating experimentally where this autoregressive growth occurs in the spectrum. For this, we use a simplified version of the Navier-Stokes problem with Kolmogorov forcing from Kochkov et al. (2021) this system is regularly used for testing neural operator models. Usually, the objective to predict a single timestep significantly larger than the timestep of the integrator used to generate the data. The large timestep is the key to the speed-up offered by neural operator methods over conventional numerical mothods. However, our goal in using this system is to demonstrate the growth of autoregressive error on nontrivial task in which prior work has established the feasibility of approximating dynamics over the given interval, so we instead predict the result of 300x smaller step to produce what should be an even casier problem than the original test case. Architectural and training details for figures related to this experiment can be found in Appendix C.I.
Figure 1 shows that the vanilla FNO must sacrifice resolution for stability. While the gap is marginal, for the initial steps, the FNO with no spectral truncation outperforms its truncated relatives, but the FNO without spectral-in-time discretization diverges even on this relatively simple task within small number of
5

---

Published in Transactions on Machine Learning Research (11/2023)
RMSE over Time
Averaged 1D Power Spectrum
Step
Step
14
Model
10a
ReFNNO
FNO
Magnitude
12
Max Prequency
Model
Nyquist
10"1
ReFNO
10
3 Nyquist
FNO
1 Nyquist
Ground Truth
RMSE
Step 20
Step 30
6
10a
Magnitude
10-1
2
10-3
20
40
60
80
100
10
20
30
0
10
20
30
Step
Frequency
Frequency
Figure 1: Navier-Stokes Error and Spectra: (Left) Conventional spectral neural operators such the FNO gain short-term accuracy by using more modes at the cost of long run stability while restructured ("Re" prefixed) approach is able to obtain the best of both worlds (Right) The instability in the full spectrum models starts in the high frequencies and grows. This is consistent with aliasing behavior.
steps. The right- hand side which depicts the spectra of un-truncated models begins to give us some insight into what is happening here. For these full spectrum models, we observe uncontrolled error growth in the high frequencies. This behavior consistent with the accumulation of aliasing error. Recall from Section 2.1, that for N samples with truncation limit the discrete aliasing error defined as:
A t) = Σ Σ (t))
(5)
k=-K K 1Ø0
The conjugate symmetry of the spectrum of real-valued functions implies that the magnitude of newly generated frequencies immediately above the bandlimit fold back onto frequencies immediately below the truncation limit causing the observed pileup near the Nyquist frequency
These issues are not limited to simple or small-scale tasks. Figure shows autoregressive rollouts from FourCastNet (FCN, Pathak et al., 2022), model basod on the Adaptive Fourier Neural Operator (Guibas et al., 2021) used for high-resolution autoregressive global weather forecasting, and also accumulates spectral pathologies. Here occur near the aliases of the zero frequency. The differences in the spectral pathologies between isotropic and multi-resoluiion architectures are explored briefly in A.3 though in limited sense and we believe it is an interesting topic for future work. As errors from aliasing or other sources inaccuracy accumulate during autoregressive time-stepping, the inputs passed to a model will gradually drift away from the training distribution into spaces where the numerical behavior of the trained model is difficult characterize. Thus in addition mitigating aliasing, in Section 4.1, we explore approaches to model less sensitive to small movements away from the training distribution in scaleable manner.
observe similar buildup, though in slightly different fashion the spikes now
3.2 Geometry and the Fourier Transform
We highlight goophysical applications as they have seen some of the largest scale uses of neural operators. Our goals include being capable of scaling our contributions the magnitude of earth systems. The FNO does not
"This often called "spectral blocking" (Boyd, 2013) the truncation limit effectively blocks the movement, of energy higher wave numbers causing pileup the Nyquist frequency.

---

Published in Transactions on Machine Learning Research (11/2023)
Step 10
Step 50
Step 100
Model
102
ReDFCN
Magnitade
FCN
UNet
10"
Ground Truth
10-2
200
400
600
600
400
600
Frequency
Frequency
Frequency
Figure 2: Averaged 1D Spectrum for FourCastNet (FCN): FCN begins demonstrating spikes in the spectrum corresponding to aliases of the zero frequency while our restructured version (ReDFCN) has consistent behavior over time. Convolutional basclines start off strong but eventually develop noisy spectra. from poor geometric handling.
inherently scale to these problem sizes, but the AFNO, more parameter-offiiient variant, successfully used in NourCastNet (FCN), one of the first neural network models to demonstrate strong performance higl-resolution global weather forecasting. However, the 2D FFT used in this model for spectral convolution proves to be a problem in the spherical setting. It implicitly considers the globe as periodic both in the zonal (east-west) direction and in the meridional (north-south) directions. This representation leads to an artificial discontinuity and allows for "local" operators to exchange information across the poles. Moreover, representing discontinuous functions in the Fourier basis can lead spurious maxima through the Gibbs phenomenon (Boyd, 2013). Some downstream effects of this pathology are depicted in Figure 4, showing model Jacobians "looking" pole-to-pole (panel b) which can lead to development of unphysical artifacts at the poles (panel c). These can grow and propagate to the rest of the domain, destabilizing model forecasts.
4 Mitigating Known Sources of Error
In this section, we list our architectural improvements which directly address the aforementioned issues relating to aliasing, sensitivity, and geometry, well as other orthogonal innovations that improve efficiency and scalability of the underlying architecture.
4.1 Efficiently Reducing Model Sensitivity
From the lens pseudospectral methods, the spectral operations in the FNO are seen as a simple spectral- domain convolution (Rippel al., 2015) rather than as the realization of an integral equation solver. This viewpoint gives the flexibility to borrow tools developed for general convolutional networks and adapt them to the spectral convolution setting which is what we do here.
Depthwise Separable Spectral Convolution. Popularized by the Xception architecture (Chollet, 2017) and MobileNet (Howard et al., 2017), depthwise separable convolutions decouple spatial convolution from channel mixing to trade expressive power at a given width for memory savings via a parameter count reduction. In most cases, this savings is then spent to expand the width of a layer. As the rapid parameter growth rate of the FNO is one of the leading obstacles scaling, creating depthwise separable variant of spectral-domain convolution is natural approach to improve scalability.
Recall from Equation 4 that the FNO employs dense spectral convolutions with the form:
(9
(F
Σ
(t) Rect [Fu( kkk
(6)
c==1

---

Published in Transactions on Machine Learning Research (11/2023)
where R( E C We can define depthwise separable convolution in the spectral domain by
instead parameterizing the convolution with decoupled channel mixing matrix V R x and spectral filter p(() E ( and computing:
d(n) vi The 1F
(DS(FF kcc Σ
(7)
Note that the channel mixing parameterized by V can be performed in the spatial domain for small efficiency savings. This reparancterization reduces the parameter growth rate from N x D x Dout to (N x D + (Din x Dowl) which in modern architectures can reduce parameter counts by hundreds or close to a thousand times. It is possible to reduce this even further for larger problems through the use of meta-networks similar to those used for continuous position emboddings (Liu et al., 2022a) generate convolutional filter weights, though at the cost of additional operations.
Spectral Normalization. In Figure 1, see that once the data becomes sufficiently corrupted, error growth becomes exponential. This behavior may be partly explained by examining the distribution of singular values in weight matrices. Empirically, we find that singular values can be very large, up to 50 for some layers. This sensitivity concern is a well-known issue (Liu et al., 2020; Van Amersfoort et al., 2020; Rosca et al., 2021) in other tasks, particularly uncertainty quantification. One mitigation approach is spectral normalization (Miyato et al., 2018), which ensures that behavior both in and out-of-distribuutin during training is bounded by rescaling weights by the spectral norm (or equivalently the maximum singular value) of the weight matrix. This ensures that the largest singular value of the weight matrix is less than one.
Circulant matrices, which represent discrete circular convolution in the spatial domain, are well known to be diagonalized by the Discrete Fourier Transform with cigenvalues equal to the Fourier cocfficients of the convolutional filter (Smith, 2007). In the case of spectral convolution, the filter is already parameterized in frequency space, so we can adapt the spectral normalization constraint exactly and differentiably by applying a squashing function to the polar parameterization of the filter coefficients. Representing the filter coefficients (1) in polar form r(a, 6) = am2000 it is sufficient to constrain |a| S 1 which can be achieved by constructing the spectral filter coefficients as:
r(a = (a)c
((a))"
(8)
where is the sigmoid nonlinearity. This ensures that the spectral norm of the convolution operator is bounded from above by one. We use this procedure for spectral domain spatial convolution and traditional spectral normalization for pointwise linear operations.
4.2 Aliasing and Learning to Filter
One of the most prevalent approaches for eliminating aliasing in pseudospectral methods is oversampling in the spatial domain or, equivalently, filtering a fixed number of the highest frequencies in the Fourier domain (Orszag, 1971). The polynomial nonlinearities frequently encountered nonlinear PDEs produce new Fourier modes predictable manner, so it is often possible define explicit oversampling or low-pass filtering strategies which exactly climinate aliasing depending on the order of the polynomial. These schemes can significantly increase the cost. of the simulation, but are often necessary for accurate, stable computation.
The dromann behavior of the non-polynomial nonlinear functions used in deep learning is signif- icantly more complicated and can vary depending on input values - for analytic functions, this can be directly observed through analysis of the Taylor expansion around different. points. In these cases, a fixed truncation strategy may be either computationally wasteful and overly diminish frequency-domate resolu- tion or may be insufficient depending on the problem. The spectral convolution of the FNO can learn problem-speciffc truncation strategies by setting filter coefficients at or very close to zero; however, this is not enough to perform low-pass filtering in the conventional FNO block. As we can in panel (a) of Figure
8

---

Published in Transactions on Machine Learning Research (11/2023)
PNO
b) Post-Filtered
NL Spectral
FFT
Couvolution
elf)
LayerNorm
Pointwise Linear
PW Linear
Pointwise Linear
Spectral Convolution
h
Pointwise
PW Linear
Linear
a
NL Spectral
[z]
al Jetto +
Convolution
h
IFFT
Figure 3: a) In the FNO, the aliasing-inducing noulinearity (h) is applied both paths, including the path that skips the spectral convolution. This makes it difficult for the model learn deal with these newly generated high frequency modes as there always a path that is unfiltered. b) By ensuring that all data passing through the nonlinearity also passes through the spectral convolution, opportunity to learn to filter any new frequency modes. c) Since the aliasing patterns of non-polynomial nonlinearities are data dependent, we further augment this approach with nonlinear spectral convolution which learns data-dependent offsets (|x), for learned filter coefficient components (a, 0) via mode-wise MLP.
allow the network the
3, the standard FNO structure includes pointwise lincar operations which bypass the convolution but still pass through the nonlinearity. Due this, learning a simple low-pass filter would actually require exact cancellation of modes by the fixed nonlinearity.
We propose to restructure the FNO block that nonlincaritics are always followed by learnable filters capable of directly filtering newly generaled high frequencies. Furthermore, since the aliasing behavior of the nonlinearities used in deep learning can vary depending on the input, we additionally replace the static convolutional filter of the FNO with one which is a dynamic function of the input data. This proposed block is presented in panels (b) and (c) of Figure 3. Dynamically generated filters would be infeasible the dense convolution case where we would need to predict D weights for each mode, but in depthwise separable spectral convolution, the convolutional filter has the same dimensionality as the input data, allowing generate filters by simple mode-wise MLP in the frequency domain. These dynamic filters are implemented as data-driven olfsets to learned biases in both the phase (0) and magnilude (a) components of the filter. The data-driven and learned components are then combined and the magnitude component is squashed by a sigmoid activation as described in Section 4.1. We note that while we are now using nonlinear function to generate the filters, the application of the filters remains a linear operation and does not produce new modes.
4.3 Improved Domain Representation
One simple approach which correct the artificial discontinuity induced by the 2D FFT while requiring almost no architectural modifications is the Double Fourier Sphere (DFS) method (Merilees, 1973; Orszag, 1974). Given an equireclangular grid defined by uniform spacing in spherical coordinates À E [0, 2T), E [0, T] (longitude, colatitude). DFS converts the original zonally periodic domain into a new domain defined on a torus. This is done by concatenating the original data with its half-phase glide reflection (Martin, 1996), a transformation consisting of reflection followed by a translation of that reflection such that a column the north-south direction of the concatenated representation follows line of longitude around the sphere.

---

Published in Transactions on Machine Learning Research (11/2023)
The mathematical formulation can be found in Appendix A.3. A sample visualization of this process for mospherii wind snapshot is given in Figure 4(a).
3D Geonelry
Representation
Withont DFS
Withons DFS
Sphere
pole query point affected by south pole
Kiottcc energy arrfacts in polar regions
With DFS
WiRL prs
Torus
instead looks acites north pole
net
(a) Sphere- Torus Transform for DFS
(b) Saliency Map for North Pole Predictions
[c] Kinetic Energy at 6O Steps (15 Days)
Figure 4: (a) The sphere-to-ton transformation of Merilees (1973) for the Double Fourier Sphere method. (b) Multi-sample averaged absolute value of derivative of predicted v10 value
north pole pixel with respect to all spatial locations from FourCastNet trained with and without the DFS transformation. The original formulation leads to the prediction at the North pole showing high sensitivily to input values at south pole. (c) Predicted surface kinetic energy after 60 autoregressive steps for FCN models trained with and without DFS. Without DFS, the polar discontinuity can contribute to rapid divergence near the poles. The Double Fourier Sphere transform is less widely used than alternative bases such as the spherical harmon- ies (Atkinson & Han, 2012), but offers advantages in both computational and implementation complexity it amounts to structured padding followed by 2D FFT. Furthermore, the DFS mapping has been proven to be continuous for certain function classes (Mildenberger & Quellmalz, 2022), has shown use in fast low-rank approximation (Townsend et al., 2015), and numerical experiments using the basis have obtained similar accuracy to full spherical harmonic models (Spotz et al., 1998; Cheong, 2000; Park et al., 2013; Fornberg, 2006). In Figure 4, we can see that this relatively simple operation sufficient to eliminate the nonphysical cross-polar behavior in FostCastNet as the saliency map (Simonyan et al., 2014) no longer shows cross-polar communication and subsequently the polar growth is eliminated.
While DFS resolves the artificial discontinuity introduced by the 2D FFT, naive impleteentations introduce two distinct issues which we introduce further refinements to address below:
Shaping with Gooretry-Aware Filters. Sphorical harmonics with the conventional triangular trunca- tion are isotropic the sense that their resolution does not vary over the sphere. On the other hand, the DFS transform applied a uniform grid has significantly higher resolution near the poles than the equa- tor which can lead to spatial distortion. This distortion, however, can be limited by filtering the frequency domain according to Spotz et al. (1998), who found that applying low-pass filters on aatitudo-y-lltitune basis can mimic the isotropy of spherical harmonics and greatly improve the stability of DFS-based numeri- cal simulations. To apply zonal wavenumber k at is the same as wavelength of zonal wavenumber an at the equator. Thus, to get uniform zonal resolution, if one has bandlimit K at the equator, one must apply a low-pass filter to remove frequencies i > sin along each latitude band.
dependen filters for a given colatitude 0, we note that the wavelength of
Hybrid Convolutions. Conventional CNNs struggle in spherical settings in large part due the spatial distortion introduced by the curvature. On an equirectangular grid defined in spherical coordinates X, 9, the majority of that distortion comes from the variance in the spacing of X coordinates across colatitudes 0. This limits the use of conventional CNNs to cubed-sphere grids like those used in Weyn et al. (2020). On the other hand, 0 coordinates are uniformly spaced which we can exploit with a factorized hybrid convolution.
Factored spectral convolution was previous employed in Tran al. (2023). Here, we take the abstraction a step further with a hybrid convolution using spectral convolution in the zonal direction followed by spatial convolution in the meridional direction. As each latitude has different. radius, allow each spectral zonal
10

---

Published in Transactions on Machine Learning Research (11/2023)
convolution to learn unique filter coefficients, though restrict the number of non-zero coefficients per latitude according to the bandlimiting discussed above. We note that using different filters, particularly given that we are not enforcing smoothness constraints across latitudes, does introduce aliasing in the meridional spectrum. However, we do not find this to be an issue in practice.
The largest benefit of the factorized hybrid convolution efficiency. As the meridional convolution is spatial, we can avoid the full sphere-to-torus DFS transformation by simply using transform-inspiree padding along the meridional direction. This circumvents the primary drawback of the full DFS method, which doubles the size of the input data when it appends the glide reflection.
5 Experiments
5.1
Rotating Nonlinear Shallow Water Equations on the Sphere
h
RMSE
10
RMSE
RMSE
200
5
100
100
200
300
100
200
300
0
100
200
300
Step
Step
Step
ReDFNO
DeepSphere
UNet
SNO
S2CNN
Climatology
FNO
ReSNO
Figure 5: RMSE over time for SWE problem. The ReDFNO demonstrates significantly stable behavior than both classical spherical architectures and other baselines. While the structural improvements were designed for DFS-based models, they also demonstrate significant stability improvements when applied to spherical harmonic models (ReSNO).
The rotated shallow water equations (SWE) on sphere are a classic test problem for dynamical cores be used in large-scale weather and climate models as they capture number of similar phenomena but are better understood and operate at more practical scale (Williamson et al., 1992). Here we use the forced hyperviscous equations in two dimensions:
Ø1 Ōu(x, u(xxt) Vau(x, -Vllz,t)-Vaahaal gVah(a, -VV 2Ω u(x,t)
(9)
(h(x,1) HV. u(x, Vn (Ma,, (),((, ()) (a, F(x,t)
(10)
Øt
where v is the hyperdiffusion coefficient, Ω is the Coriolis parameter, is the velocity field, H is the mean height, and h denotes deviation from the mean height. F is daily/scasonall) varying forcing with periods of 24 and 1008 simulation "hour" respectively. The simulations were performed using the spin-weighted spherical harmonic spectral method Dedalus (Burns et al., 2020) with 500 simulation hours of burn-in where the next three simulation years (3024 hours), were collected for the data set. We use initial conditions for training data and 3 and initial conditions for validation and test data, respectively. Full simulation settings can be found in Appendix C.2.
We use the shallow water equations to examine the impact of our proposed changes to the FNO model and compare our method to more popular spherical deep learning models such as the Spherical CNN (Esteves et al., 2017) and DeepSphere (Defferrard et. al., 2020) to illustrate that geometry alone is insufficient to address the observed stability issues. Additionally, we run a comparison against what we call a Spherical Neural Operator both with our modifications ("Restructured SNO', RoSNO) and without them (SNO). The SNO is an FNO implemented using SHT instead of 2D FFT. As our core model uses the Double Fourier Sphere representation, refer to it the Restructured Double Fourier Neural Operator (ReDFNO). We
11

---

Published in Transactions on Machine Learning Research (11/2023)
further include UNet-structured (Ronneberger et al., 2015) CNN constructed from ConvNext blocks (Liu et al., 2022b) to provide system-agnostic baseline. All models are implemented in PyTorch (Paszke et al., 2019) and we use TS2Kit (Mitchel et al., 2022) for spherical harmonic transforms. All models were trained for four hours or 40 epochs on 4x NVIDIA A100 GPUs.
Spectrum Evolution
Resulting Spatial Artifacts
Step 1
Step 5
S2CNN (t=200)
DeepSphere (t=20)
10a
Magnitude
10
10-7
Step 20
Step 200
ReDFNO (t=200)
GT (t-200)
Magnitude
10
10-7
Frequency
Frequency
RRDDFNO
S2CNN
DS
GT
Figure 6: Visualizing failure modes: (left) Both the S2CNN and DeepSphere models have UNet-like structure and so develop spikes in the spectrum corresponding different resolutions. ReDFNO maintains consistent spectrum over many timesteps. (right) These spectral artifacts manifest spatially with S2CNN becoming very blocky and DeepSphere developing high amplitude, high froquency patterns near the poles.
We present our results in Figure 5. The gap between the ReDFNO and ReSNO is striking and can be seen as unexpected given the spherical harmonics do not introduce the same level of distortion as the DFS representation. However, for all spherical harmonic models, TS2KIT uses the Driscoll-Healy (Driscoll & Healy, 1994) SHT. For bandlimit K, Driscoll- Healy assumes input grid of size 2K 2K while we have 2K 4K grid. As result, the transform truncates half the zonal bandwidth.
Results.
The ReDFNO strongly outperforms all non-stabilized models.
Figure 6 is perhaps more informative in the context of the analysis performed in this paper as we can see how the models begin to diverge for architoctures with appropriate spherical handling. Both DecpSphere and the S2CNN exhibit similar patterns to Figure 2. Both use UNet-style construction without low-pass filtering on downsampling passes, though the Sriscoll-Healy transform used in TS2KIT adds implicit filtering to the S2CNN which could contribute to its more stable behavior. DeepSphere blows up quickly while for the S2CNN, see spikes develop the aliases of the zero frequency then grow slowly. This manifests as block artifacts in the full resolution prediction. On the other hand, given our approach is designed around spectral principles, we see minimal spectral artifacts for our stabilized model and in fact, the model can run almost indefinitely without spectral distortion. Over sufficiently long roll-outs, we hit the limit of predictability and the error converges to that of the persistence forecast, as occurs conventional atmospheric forecasting, We include an ablation of our different model architecture innovations and their effect on long-term stability Appendix Table Numerical rather than graphical results can be found in Appendix Table 2.
5.2 ERA5
In this section, we focus on a real-world application and one of loday's grand challenge problems global weather forecasting. We use the public dataset ERA5 (Hersbach et al., 2020), provided by ECMWF (Euro- pean Center for Medium-Range Weather Forccasting), which consists of hourly prodictions of several crucial atmospheric variables (such as wind velocities and geopotential heights) at spatial resolution of 0.25 (that corresponds to 720 x 1440 lat-lon grid; 25 km spatial resolution) from years 1979 to present day. ERA5
12

---

Published in Transactions on Machine Learning Research (11/2023)
25 Steps
30 Steps
100 Sceps
400 Steps
u10
FCN
8-
RMSSE.
6-
Ours
07-
20
40
60
100
GT
Step
ReDFCN FCN
-.. UNet
Clinastoloyy Pensstennc
12m
FON
RMSE
Ohrs
2-
0-
20
20
60
100
GT
Step
Figure 7: Comparing FCN with our restructured version and ground truth (GT): (left) RMSE by step for u10m and t2m. (right) Visualizing the forecasts over long time horizons. We can see that while base FCN diverges quickly after its targeted forecast horizon. Our stabilized version produces detailed predictions out to 400 steps (100 days).
is a reanalysis dataset that combines observational data with state-of-the-art numerical simulation methods (through Bayesian data assimilation (Kalnay, 2003)). This dataset is widely accepted to be representative of the best available estimate of the atmospheric state of the Earth at any time point. Following Cast- Net (Pathak et al., 2022), we locus predicting (forecasting) 20 atmospheric variables al 25 km spatial resolution and 6-hourly time intervals (see Appendix C.3.3 for a list of variables used). Our training and inference pipeline identical to FourCastNee (and other similar works) with our revised model training for hours across 16 nodes with 4x NVIDIA A100 GPUs per node. Each training data point is a tensor of size x 720 x 1440 (20 channels represent spatial profiles of different atmospheric variables) and the model is trained predict timestep + 8t from initial time point f, where öt 6 hours. During inference, the model autoregressively predicts further timesteps (similar to SWE).
We apply our model (and input representation) changes to the FourCastNet architecture and demonstrate that our innovations can significantly improve the long-term stability of FourCast.Net on several atmospheric variables, which is crucial for sub-scasonal to seasonal time scales (beyond two weeks), and more importantly, for these models to be useful for understanding and analyzing Earth's climate.
Results. Figure 7 shows our results on the u10m and 12m variables. Results for the remaining 18 variables be found in Appendix C.3.3. The results consistently demonstrate that changes are able to match the predictive accuracy of FourCastNet while enabling significantly longer prediction horizons. Here, we highlight two fields-for u10m, while model error saturates around persistence, the rollout to 400 steps (100 days) shows remarkable level of detail given that the base model shows signs of divergence at around 50 steps. This an 800% improvement in prediction horizon; for t2m field, while our stabilized model does not diverge, we can see one of its lingering weaknesses we observe unphysical behavior around coasts and mountain ranges (discontinuities due to Gibb's phenomenon. This behavior appears around step 100 and by step 400 has impacted significant area. Nonetheless, the model has not diverged at 400 steps.
13

---

Published in Transactions on Machine Learning Research (11/2023)
6 Discussion and Conclusion
We showed that we can extend the forecast horizon for autoregressive neural operators significantly at min- imal additional cost. For sufficiently smooth fields, like those in the shallow water system and Kolmogorov- forced Savier-Stokess our design principals allow us to take spectral neural operators that were previously unstable and extend their forecast horizons indefinitely. For more complex and large-scale systems like ERA5, we are still able to demonstrate an 800% improvement in forecast horizon, demonstrating a notable advance towards one of the goals of deep learning- weather forecast- ability to extend weather-scale models to seasonal or even climate-length prediction.
Limitations. The ERA5 example demonstrates some of the challenges that still exist in fully extending to climate modeling. The imperfect spectral matching in the multi-resolution architecture implies that work remains to fully understand that particular pathology and rectify it. Further, ERA5 has sharp discontinuities along coasts and mountain ranges and we find that with one-step-ahead training, our current methods still fail to learn sufficiently diffusive dynamics to avoid energy accumulation in these areas indefinitely without significant spectral truncation despite significant improvements those fronts.
Societal Impact. Given the diffieulty of evaluating long-run behavior, one of the key risks inherent to the development of data-driven models is that their success on short-term metrics could be used by bad actors to discount first-principles projections despite the fact that these models behave unpredictably away from the training manifold. Our work is concerned with reducing that risk by identifying and addressing the issues that cause current models to fail.
While work remains, these results represent a significant improvement in our understanding of autoregressive neural operators for spatiotemporal forecasting. We provided analyses detailing where these models fail and identified principled, low cost architectural changes that can mitigate these failure points. Our experiments reinforced these amalyses and demonstrated the success of our innovations, while simultaneously supporting the argument that the failure modes identified apply to more families of nonlinear forecasting models than specifically studied here.
References
Kendall Atkinson and Weimin Han. Spherical hormonics and approzimations on the unit sphere: an intro- duction, volume 2044. Springer Science & Business Media, 2012.
Kaifeng Bi, Lingxi Xie, Hengheng Zhang, Xin Chen, Xiaotao Gu, and Qi Tian. Pangu-Weather A 3D High--esoolution Model for Fast and Accurate Global Weather Forecast, 2022. URL https://arxiw.org abs/2211.02555
Rishi Bommasani, Drew A Hudson, Ehsan Adeli, Russ Altman, Simran Arora, Sydney von Arx, Michael S Bernstein, Jeannette Bohg, Antoine Bosselut, Emma Brunskill, et al. On the opportunities and risks of foundation models. arXiv preprint arXim:2108.077 07258, 2021.
J.P. Boyd. Chebyshev and Fourier Spectral Methods: Second Revised Edition. Dover Books on Mathe- matics. Dover Publications, 2013. ISBN 9780486141923. URL https://books google. com/books?id= b4TCAgAAQBAJ.
Keaton J. Burns, Geoffrey M. Vasil, Jeffrey S. Oishi, Daniel Lecoanet, and Benjamin P. Brown. Dedalus: A flexible framework for numerical simulations with spectral methods. Physical Review Research, 2(2): 023068, April 2020. doi: 10.1103/ 2.w.eeeeeeeeearr....22222000 2223006800445555//721111 200.50///..... 10388.
Hyeong- Bin Cheong. Double Fourier Series on a Sphere: Applications to Elliptic and Vorticity Equations.
Journal of Computational Physics, 157(1) 349, 2000. ISSN 0021-9991. doi: https:///doioor//10.0000
1999.1999.6385. URL
François Chollet. Xception: Deep Learning with Depthwise Separable Convolutions. In 2017 IEEE Confer- ence on Computer Vision and Pattern Recognition (CVPR) pp. 1800-1807, 2017. doi: 10.1109/CCPP
2017.195.
14

---

Published in Transactions on Machine Learning Research (11/2023)
Taco S. Cohen, Mario Geiger, Jonas Köhler, and Max Welling. Spherical CNNs. In International Conference on Learning Represenlations, 2018. URL. https://openrevieww nett/forummidd-RkbbbZZbbb
Yann N. Dauphin, Angela Fan, Michael Auli, and David Grangier. Language modeling with gated con- volutional networks. In Doina Precup and Yee Whye Teh (eds.), Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pp. 933 941. mlr, www.ppressssooee http//////aapp/////// html.
PMLR, 06-11 Aug 2017. URL https://proceeddii
Tim De Ryck and Siddhartha Mishra. Generic bounds on the approximation error for physics-informed (and) operator learning. arXiv preprint orXiv:2205.11998, 2022.
Aaron Defazio and Konstantin Mishchenko. Searning-rate-ive learning by D-adaptation. arXiv preprinl arXiv:2301.07793, 2023.
Michaël Defferrard, Martino Milani, Frédérick Gusset, and Nathanaël Perraudin. DeepSphere: a graph- based spherical CNN. In International Conference on Learning Representations (ICLR), 2020. URL
https://openreview http:ppwwww..eee.eeeeeeoraraaaaa
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. ImageNet: A large-scale hierarchical image database. In 2009 IEEE Conference on Computer Vision and Pattern Recognition, pp. 255, 2009. doi: 10.109//CVPR.2099.33
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. orXiv preprinl ar 1rXiu:1810.04800 2018.
J.R. Driscoll and D.M. Healy, Computing Fourier Transforms and Convolutions the 2-Sphere. Advances in Applied Mathematics, 15(2):202-255 1994. ISSN 0196-8858. doi: httpps://ooi.org///i 1006 htt.//aam.............. 1904.
URL https: http:/wwwwcccc
Stefan Flfwing, Fiji Uchibe, and Kenji Doya. Sigmoid- linear units for neural network function approximation in reinforcement learning, 2017.
Erik Eliasen, Bennert Machenhauer, and Erik Rasmussen. On numerical method for integration of the hydrodynamicaal equations with a spectral represenlation of the horizontal fields. Kobenhavns Universitet, Institut for Teoretisk Meteorologi, 1970.
Carlos Esleves, Christine Allen-Blanchetle, Ameesh Makadia, and Kostas Daniilidis. Learning SO(3) Equiv- ariant Representations With Spherical Cuns. CoRR, 2017. URL
06721.
V. Fanaskov and I. Oseledets. Spectral neural operators, 2022. URL
10573. 10573.
Bengt Fornberg. A Practical Guide to Pseudospectral Methods. Cambridge Monographs Applied and Computational Mathematics. Cambridge University Press, 1996. doi: 10.1077/CR097805116666
Bengt Fornberg.
Pseudospectral Approach for Polar and Spherical Geometries. doi: 10.1137/091600
16:1071-1081, 7 2006. ISSN 1064-8275. URL https: siam. 1ttp///11..11779166611111111
http://dbiioog/11.11/////0000
Steven Guan, Ko-Tsung Hsu, and Parag V Chitnis. Fourier neural operator networks: A fast and general solver for the photoacoustic wave equation. arXiv preprint arXic:2108. 093774 2021.
John Guibas, Morteza Mardani, Zongyi Li, Andrew Tao, Anima Anandkumar, and Bryan Catanzaro. Effi- cient Token Mixing for Transformers via Adaptive Fourier Neural Operators. In International Conference on Learning Representations, 2021.
Hans Hersbach, Bill Bell, Paul Berrisford, Shoji Hirahara, Andrés Horányi, Joaquín Muñoz-Sabaler, Julien Nicolas, Carole Peubey, Raluca Radu, Dinand Schepers, Adrian Simmons, Cornel Soci, Saleh Abdalla, Xavier Abellan, Gianpaolo Balsamo, Peter Bochtold, Gionata Biavati, Jean Bidlot, Massimo Bonavita, Giovanna De Chiara, Per Dahlgren, Dick Dee, Michail Diamantakis, Rossana Dragani, Johannes Flem- ming, Richard Forbes, Manuel Fuentes, Alan Geer, Leo Haimberger, Sean Healy, Robin J. Hogan, Elías
15

---

ptt.iiee//rrccee/////ppooo20o2o. Published in Transactions on Machine Learning Research (11/2023)
H6lm, Marta Janisková, Sarah Keeley, Patrick Laloyaux, Philippe Lopez, Cristina Lupu, Gabor Radnoti, Patricia de Rosnay, Tryna Rozum, Freja Vamborg, Sebastien Villaume, and Jean-Noöl Thépaut. The ERA5 global reanalysis. Quarteriy Journal of the Royal Meteorological Society, 146(730):1999 2049, 2020. ISSN 1477-870X.
Andrew G. Howard, Menglong Zhu, Bo Chen, Dmitry Kalenichenko, Weijun Wang, Tobias Weyand, Marco Andreetto, and Hartwig Adam. MobileNets: Efficient Convolutional Neural Networks for Mobile Vision Applications, 2017. URL https://arxiv. http:://axxii.orr/aabs.... wtwgabs/////a/ 0t///1704.040 04861.
J.K. Hunter and B. Nachtergaele. Applied Analysis. World Scientific, 2001. ISBN 9789810241919. URL https://books. google cow/bookssiddoooooeeeiiieee
Eugenia Kalnay. Atmospheric modeling, data assimilation and predictability. Cambridge University Press, 2003.
Tero Karras, Miika Aittala, Samuli Laine, Erik Härkönen, Janne Hellsten, Jaakko Lehtinen, and Timo Aila. Alias-Free Generative Adversarial Networks. In Proc. NewrIPS, 2021.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization, 2017.
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(21): 2021. doi: 10.1078/pna..210174411. URL pttts::/www.ppaa..oo org/doi/ abs/10.1073/pmas. http://mmaa..11111111113 2101784118..
Nikola Kovachki, Samuel Lanthaler, and Siddhartha Mishra. On universal approximation and error bounds for Fourier Neural Operators. Journal of Machine Learning Research, 22:Art No, 2021.
Aditi S. Krishnapriyan, Alejandro F. Queiruga, N. Benjamin Erichson, and Michael W. Mahoney. Learning continuous models for continuous physics, 2022.
Remi Lam, Alvaro Sanchez- Sanchez-Gonzallee Gonzalez, Matthew Willson, Peter Wirnsberger, Meire Fortunato, Alexander Pritzel, Suman Ravuri, Timo Ewalds, Ferran Alet, Zach Eaton-Rosen, Weihua Hu, Alexander Mcrose, Stephan Hoyer, George Holland, Jacklyun Stott, Oriol Vinyals, Shakir Mohamed, and Peter Battaglia. GraphCast: Learning skillful medium-range mrangee-rgggeeeeeee global weather forecasting, 2022. URL https://arxiv.. org/ abs/2212.127999
Zhijic Li, Wenhui Peng, Zelong Yuan, and Jianchun Wang. Fourier neural operator approach to large eddy simulation of three-dimensional turbulence. Theoretical and Applied Mechanics Lellers, 12(6):100389, 2022. ISSN 2095-0349. doi:
URL https: /ttps://www.ce sttp///ww..ccenccrrrect...
Zongyi Li, Nikola Kovachki, Kamyar Azizzadencsheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stu- art, and Anima Anandkumar. Markov neural operators for learning chaotic systems. arXiv preprint arXiv:2106 06898, 2021a.
Zongyi Li, Nikola Borislavov Kovachki, Kamyar Azizzadeneshcli, Burigede liu, Kaushik Bhattacharya, An- drew Stuart, and Anima Anandkunar. Fourier Neural Operator for Parametric Partial Differential Equa- tions. In International Conference on Learning Representations, 2021b. URL https://openreriis...t net/ forum? 2ormm?id==oPP00QVmm.
Jeremiah Liu, Zi Lin, Shreyas Padhy, Dustin Tran, Tania Bedrax Weiss, and Balaji Lakshmi- narayanan. Simple and principled uncertainty estimation with deterministic deep learning via distance awareness. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (ods.). Advances Neural Information Processing Systems, volume 33, pp. 7498-7512. Curran Associates, Inc., 2020. URL https: s/tope//prrcceedings neurips. wttpppp..ccppppp.ppepp http:/pper.illeepprrr////// pdf.
16

---

Published in Transactions on Machine Learning Research (11/2023)
Ze Liu, Han Hu, Yutong Lin, Zhuliang Yao, Zhenda Xie, Yixuan Wei, Jia Ning, Yue Cao, Zheng Zhang, Li Dong, Furu Wei, and Baining Guo. Swin transformer v2: Scaling up capacity and resolution. In International Conference on Computer Vision and Pattern Recognition (CVPR), 2022a.
Zhuang Liu, Hanzi Mao, Chao- Yuan Wu, Christoph Feichtenhofer, Trevor Darrell, and Saining Xie. A convnet for the 2020s, 2022b.
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.
G.E. Martin. Trensformation Geometry: An Introduction to Symmetry. Undergraduate Texts in Mathe- matics. Springer New York, 1996. ISBN 9780387906362. URL https://books. www.le...oooooooooooooo com/books ?id= KW4EwONaQJgC.
Philip E. Meriloes. The pscudospectral approximation applied the shallow water equations on a sphere. htlp:// doi. 0r9/10. 00800000,6973,1 /00//00096933117393655550 URL https: /wwp///wwwwwwo///oo/////// htt.a///ooo.ae...oooooooo/////// http:////////aa//..../////////0/ 9648342.
11:13 20, 1973. ISSN 14809214. doi: 10.
Sophie Mildenberger and Michael Quellmalz. Approximation Properties the Double Fourier
Sphere Method. Journal of Fourier Analysis and Applications, 28:1-30, 2022. ISSN 15315851, doi: 10.
URL https: https://linkk wwr.inkk.sriiggr.... com/articceet1011077
500041-022-09928--
Thomas W. Mitchel, Noam Aigerman, Vladimir G. Kim, and Michael Kazhdan. Möbius Convolutions for Spherical CNNs. In ACM SIGGRAPH 2022 Conference Proceedings, SIGGRAPH '22, New York, NY, USA, 2022. Association for Computing Machinery. ISBN 9781450393379. doi: 10. 252553333329337222444 URL httpss//dolt org/ 1145/3528233 3555222233333774....
3530724.
Takeru Miyato, Toshiki Kataoka, Masanori Koyama, and Yuichi Yoshida. Spectral Normalization for Gen- crative Adversarial Networks. In International Conference on Learning Representations, 2018. URL
nttps://openreree httpr/iieeeeeeettoorurttt----BB
Steven A Orszag. Transform Method for the Calculation of Vector-Coupled Sums: Application to the
Spectral Form of the Vorticity Equation. Journal of Atmospheric Sciences, 27:890 - 895, 1970. doi:
2.0.0 URL https:// //orrnaal ametsoc org/view/ 0011190000200000_____11 08900_0800__00 __too_____o_____1_____________ xml.
jourmalleaattse/77//550000-009
Steven A. Orszag. On the climination of aliasing in finito-difference schemes by filtering high-wavenumber components. Journal of Almospheric Sciences, 28(6):1074 1074, 1971. doi: https//A0o.orr//101113/
URL https: / /ttp://jourrn jorrnalusllsss.. ametsoc org/vves//vooooraals//////////// atsc/28/6/1500- 1971_028 1111111_22_111____11111111111111 2
Steven A Orszag. Fourier Series on Spheres. Monthly Weather Review, 102:56 75, 1974. doi: 10.1175/
1520- 152.000081117//00000000655555550 mwre/ mwte//1D20//155000088889994442
URL https://journals ametsoc awetsoc. poottsoc.or/viee//jrraa.... xml.
Hoon Park, Song- You Hong, Hyeong- Bin Cheong, and Myung- Seo Koo. A Double Fourier Series (DFS) Dynamical Core in a Global Atmospheric Model with Full Physics. Monthll Weather Review, 141(9):3052
doi: 10.1175/NIWR-D--2-00000 URL https: //jj://jjolasssee anstsooo.org//iie////onaaa 3061, 2013.
1.ml.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjic Bai, and Soumith Chintala. PyTorch: An Imperative Style, Hogo-PPaformance Deep Learning Library. H. Wallach, H. Tarochelle, A. Beygelzimer, F. d' Blu F. Fox, and R. Garnett (eds.), Advances
17

---

Published in Transactions on Machine Learning Research (11/2023)
in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019. URL https:// proceediigs.neuripp... ppreeeligg.aartps..ccppp///////p
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: global data-driven high-resolution weather model using adaptive fourier neural operators. arXiv preprint arXiv: 11214, 2022.
Antônio H. Ríbeiro and Thomas B. Schön. How Convolutional Ncural Networks Deal with Aliasing. ICASSP 2021 2021 IEEE International Conference on Acoustics, Speech and Signol Processing (ICASSP), pp. 2755 2755-2759, 2021. doi: 10.1109/ICAASPPP
Oren Rippel, Jasper Snock, and Ryan P. Adams. Spectral Representations for Convolutional Neural Net- works, 2015. URL https://arxiv.org///a wwtp/a////////eeeeeeeeeeeeeeg//e 0t////1560003700
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-Net: Convolutional Networks for Biomedical Image Segmentation. In Nassir Navab, Joachim Hornegger, William M. Wells, and Alejandro F. Frangi (eds.). Medical Image Computing and Computer-Assistern Intervention MICCAI 2015, pp. 234 241, Cham, 2015. Springer International Publishing. ISBN 978-3-319-24544-4
Mihaela Rosca, Theophane Weber, Arthur Gretton, and Shakir Mohamed. A case for new neural network smoothness constraints, 2021.
Mark Sandler, Andrew Howard, Menglong Zhu, Andrey Zhmoginov, and Liang-Chieh Chen. Mobilenetv2: Inverted residuals and linear bottlenecks, 2019.
Noam Shazeer. Glu variants improve transformer, 2020.
Karen Simonyan, Andrea Vedaldi, and Andrew Zisserman. Deep inside convolutional networks: Visualising image classification models and saliency maps, 2014.
Julius O. Smith. Mathemalics of the Discrele Fourier Transform (DFT). http://www..3lkoorr/bbooos 2007. ISBN 978-0-9745607-448.
W3K Publishing,
William F. Spotz, Mark A. Taylor, and Paul N. Swarztrauber. Fast shallow-wate water equation solvers in latitude-longituude coordinates. Journal of Computational Physics, 145:432 444, 9 1998. ISSN 00219991. doi: 10.1006/JCCH..1998.66
Kim Stachenfeld, Drummond Buschman Fielding, Dmitrii Kochkov, Miles Cranmer, Tobias Pfaff, Jonathan Sonzaez- Learned simulators for turbulence. In International Conference Learning Representations, 2022. URL https::/opeereviiee fot/foorumiddmmRRBo]]]
Godwin, Can Cui, Shirley Ho, Peter Battaglia, and Alvaro Sanchez- Gonzalez.
Alex Townsend, Heather Wilber, and Grady B. Wright. Computing with functions in spherical and polar geometries I. The sphere, 2015. URL https://arxiv org/ abs//510.111....11 wttp:///..rg/aa/111.....1 08094.
Alasdair Tran, Alexander Mathews, Lexing Xie, and Cheng Soon Ong. Factorized fourier neural operators.
In The Eleventh International Conference on Learning Representations, 2023. URL https://openrevvee: net/forum?? 1it
Joost Van Amersfoort, Lewis Smith, Yee Whye Tch, and Yarin Gal. Uncortainty cstimation using a sin- gle deep deterministic neural network. In International conference on machine learning, pp. 9690-9700. PMLR, 2020.
Cristina Vasconcelos, Hugo Larochelle, Vincent Dumoulin, Rob Romijnders, Nicolas Le Roux, and Ross Goroshin. Impact of aliasing on generalization in deep convolutional networks. In Proceedings of the IRFE/CVP International Conference Computer Vision, pp. 10529000555 2021.
18

---

Published in Transactions on Machine Learning Research (11/2023)
Gege Wen, Zongyi Li, Kamyar Azizzadenesheli, Anima Anandkumar, and Sally M Benson. U-FNO- An enhanced Fourier neural operator-based deep-learning model for multiphase flow. Advonces in Water Resources, 163:104180, 2022.
Jonathan A. Weyn, Dale R. Durran, and Rich Caruana. Improving Data-Driven Global Weather Pre- diction Using Deep Convolutional Neural Networks Cubed Sphere. Journal of Advances in Modeling Earth Systems, 12(0):20200550002000,,, 2020.
doi: 10101/..10000000000022222 c2020MS002109
URL https: /agupubs. onl inelibrarry..... silaarrywlllly wiley. cwwc/oo///ooooo wwwp/ab///// 10. 00.1099/2020MSS00210
David L. Williamson, John B. Drake, James J. Hack, Rüdiger Jakob, and Panl N. Swarztrauber. A standard test set for numerical approximalion to the shallow waler equations in spherical geometry. Jowrnal of Com- putational Physics, 102(1):211-22)) 222,,, 1992. ISSN 0021-9991. doi:
80016-6. URL https: http:://wwwwwwwwewwe...m sci wwweenniieeeet......ee....ee.e.e com/ htt.://wwc..ct....ooooooooaaaaaa
Philipp A Wille, Russell Hewelt, and Ranveer Chandra. Industry-scale CO2 Flow Simulations with Model- Parallel Fourier Neural Operators. In NewrlPS 2022 Workshop Tackling Climate Change with Machine Learning, 2022. URL https: //www. climatechange httpcaaannnaaaaaaaaaeeeeeeeeeee ai/papersssseeeaaaaaaaa httppr///eerris22222.12
Xingyu Xie, Pan Zhou, Huan Li, Zhouchen Lin, and Shuicheng Yan. Adan: Adaptive nesterov momentum algorithm for faster optimizing deep models, 2023.
Ziyi Yin, Ali Siahkoohi, Mathias Louboutin, and Felix J Herrmann. Learned coupled inversion for carbon sequestration monitoring and forecasting with fourier neural operators. arXiv preprint orXiv:2203.1423 2022.
Yang You, Jing Li, Sashank Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiaodan Song, James Demmel, Kurt Keutzer, and Cho-Jui Hsieh. Large batch optimization for deep learning: Training bert in 76 minutes, 2020.
Richard Zhang, Making Convolutional Networks Shifl-Invariant Again. In ICML, 2019.
Xueyan Zou, Fanyi Xiao, Zhiding Yu, and Yong Jac Lee. Delving Deeper into Anti-aliasing in ConvNets. In BMVC, 2020.
A Extended Details
A.1 Pseudospectral Method
Here we perform a step-by-step derivation of the pseudospectral equations from Section 2.1. Recall that we are exploring the nonlinear advection problem for scalar field u E (T) defined on the periodic domain x E [0:
Ou(x,)) uxxxtt) Ot
Ou(x, 1) Ox
(11)
Pseudospectral methods compute a solution to the differential equation at a set N spatial locations (ai....,,,, > through the use of dcrivatives computed high accuracy in the spectral domain (Fornberg, 1996). Assuming that u has bandlimit K where N we can exactly represent at the nh collocation point the discrete Fourier expansion:
K
(t) = Σ
dg hon k=-K K
(12)
19

---

Published in Transactions on Machine Learning Research (11/2023)
where 2 and (t) / Ex=1 un (t)) In frequency space, computing the exact derivative amounts to elementwise multiplication which by the Convolution Theorem (Smith, 2007) is equivalent to a convolution in the spatial domain, so plugging 12 into 11, we get:
VN
Oun(t) Öt
-
Σ
K
(e(t)e
)
Σ
milm (t) han)
(13)
Expanding the product of summations reveals modes with wave number up 2K. In frequency space, we can represent the equation as a system of independent ODEs at each wave number A = -2K,.... 2K:
(ūg(t)
2K
Σ
(Te ())x-e(t)ae(t))
(14)
8= 2K
-aa+ Ou(z,1) 1)
[kc]
(15)
Ox
which now requires performing convolution in frequency space as well. Naive computation of this convolution is quadratic the original bandlimit, but this cost can be avoided with fast transform methods like the FFT. Through the dual of the Convolution Theorem, we have that û=ú methods, we perform the convolution clementwise multiplication in the spatial domaín with the limiting cost being that of the transform are computed, one can use any off-the-shelf ODE integrator to march the equation forward via the method of lines (Boyd, 2013).
((FF-a-FF-a û) and for pseudospectral
O(N log N) in the case of the FFT. Once the time derivatives
However, this introduces a complication as the bandlimit of the RHS of Equation 13 twice that of the original function. For K - the bandwidth of the time derivative now greater than the Nyquist limit. Without increasing the spatial sampling rate, these newly generated modes will alias back onto modes below the Nyquist limit introducing new source of error.
A.2 Parameterizing the FNO as a Pseudospectral Method
We begin with the form of an FNO block. Recall that our input (+)) E RN for layer E {o..... I.} at each discretization point such that v(0) = u. N, K, denote the sizes of the spatial, frequency, and channel dimensions respectively. The block is parameterized by w(O e R x and R(!) E ( x and computed as:
D
(f+1)
h Σ
(0) Uoe
(0)
(F-1
(16)
c==1
D
FNO
(Fol
Σ
R10 [F
he!
(17)
c==1
We aim to show how this block can be parameterized such a way that is equivalent to:
t (/() + 1)
(++1)
(18)
at
0u!0 (1)
Use
(t)
o F (F
(19)
Øt
PS (F kc = 12 Fu(1)
(20)
where we can assume 77 1 without loss of generality it can be absorbed into any multiplicative term. Converting between the two requires structural adjustment and then explicitly setting our parameters
20

---

Published in Transactions on Machine Learning Research (11/2023)
and activations. For the first step, move the first term (channel mixing by W) outside of the nonlinear activation. We then set W to be the identity matrix giving us:
% vi
(e+1) (2) +
( (FU
(21)
Now for R2 we assume Din = and Dmat The reason for this will become apparent below. We then set R. to be:
- 12mk
c=d= =1
Rkee!
c=
(22)
elsc
otherwise. This gives two channels - the first applies the spectral differentiation operation (20). The second is an identity. We can then define as a GLU activation (Dauphin et al., 2017; Shazeer, 2020) which splits the channels and performs elementwise multiplication. This completes the mapping by giving us the FNO-like update:
1/c
(f+1)
ie + GLU (F
(23)
D
1 (Fu ke
Σ
Rkee (Fy [ke
(24)
c'ml
A.3 Aliasing and Nonlinearity
Aliasing in discretized settings is the result individual modes being indistinguishable the given sampling points. For the Fourier basis, we know from the Shannon-Nyqustt Sampling Theorem (Smith, 2007) that for set of N evenly spaced points, we are able to exactly recover signals with bandlimit up to N/2 via the Discrete Fourier Transformation (DFT). However, in cases where the signal is not bandlimited, the coefficients returned by the DFT are incorrect as higher modes alias onto lower modes. In this case, if we denote by UR the true Fourier coefficients that we would obtain in the infinitely sampled regime, the DFT returns coefficients:
ax
(25)
n==
Un Σ UK+jN
(26)
170
If we denote by Ug the true Fourier coefficients one would obtain in the contimuous regime, we can observe that aliasing error, the distance between the true truncated function and the function returned by the DFT is defined as:
K APK[M] = Σ (U= 2x) = Σ k=-K 1=0
Uk+sNy
(27)
b=-K
In many deep learning settings, addressing aliasing has led to improvements in accuracy and generalizability (Vasconcelos et al., 2021; Ribeiro & Schön, 2021; Zhang, 2019; Karras et al., 2021), but these findings are often neglected in the design of new architectures. However, in the autoregressive forecasting setting, these types of issues become a significantly larger obstacle.
In Equation 13, we saw that the nonlinearity in the true nonlinear advection equation results in bandlimit expansion. This is in fact true of any arbitrary nonlinearity, though the results are less predictable in
21

---

Published in Transactions on Machine Learning Research (11/2023)
No Resampling
Last Step Oversample
2x Oversample Per Step
Step 10
109
10
10"
10-2
10
8
10
50
100
150
200
100
600
0
200
100
600
10
Step 50
10
10
10-1
10-3
10
50
100
150
200
100
600
200
100
600
10a
Step 150
10
10
10"
10
10 3
100
150
200
400
600
Frequency
Frequency
Frequeney
sigmoid square
gelu
tanh
Figure 8: Spectrum evolution of ERA5 data under repeated application of random projections and nonlin- carities. Left is at fixed sampling level. Center applies the final nonlinearity at higher resolution. Right oversamples by 2x before applying every nonlinearity then truncates the top half the modes to return to the original resolution.
the non-polynomial cases. For continuous nonlinearities like those typically used deep learning, this can be viewed implication of the Weirstrass Approximation Theorem which states that any continuous function interval can be arbitrarily closely approximated by linear combination of polynomials. Exactly characterizing this behavior is in open problem, but we can observe it empirically.
To do this, we perform a simple experiment. We take a sample from one of our experiment datasets, ERA5 and repeat the following procedure:
1. Project via Gaussian random matrix (equivalent to an untrained linear layer) into a larger dimen- sional subspace.
2. (Optional) Oversample data spatially via Fourier interpolation.
3. Apply nonlinearity.
4. (Optional) Downsample via Fourier truncation.
Remix via Gaussian random matrix. Renormalize.
As we can see in Figure 8, the oversampled activation (right) produces drastically different and more consis- tent spectra compared to the left and center pancls. The contrast between the left and center are particularly interesting and relevant for U-shaped architectures. On the left, we see that at fixed resolution, the square function does eventually show signs of more typical spectral blocking pattern, but repeated application of
22

---

where F is the discrete Fourier transform and O denotes elementwise multiplication. For large kernels, spectral domain convolution is significantly faster as the limiting factor is the N log N transform compared to the 2 of the spatial implementation. However, for the small kernel convolutions generally used in spatial CNNs, methods like im2col enable linear scaling for the spatial convolution. In practice using Py Torch, we' ve found the kernel size where the execution times equalize to be surprisingly small, though in many cases the convolutions must carry a much larger mumber of Published in Transactions on Machine Learning Research (11/2023)
the activation functions flattens the spectrum. The right panel indicates this effect is diminished, though it doesn't 't appear to be eliminated, by explicit anti-aliasing. An interesting aspect is that applying activations to oversampled versions of the flattened spectrums now leads to spikes at the Nyquist frequencies for each resolution. This is a drastic contrast with the anti-aliased spectrum which has no spikes. This experiment shows that while exactly characterizing the impact of aliasing difficult to analytically model, it does seem to contribute toward the issues we observe in real training scenarios.
A.4 Comparing Spatial and Spectral Convolution
Spatial and spectral convolution are connected through the Convolution Theorem. Given N uniformly-spaced samples from signals f.g: the circular discrete convolution is defined:
(f * [
Σ
fimlgin ml
(28)
Equivalently by the convolution theorem, we could define this as:
(f * [)
(29)
local response is desirable and parameters to represent the same filter.
However, in some cases, there are advantages to non-sparse weights as well. If we consider a one-dimensional kernel of size 3, the discrete Fourier transform is computed on the expanded kernel [wa, W1, O, ...,..... 0, w-1::
W kn
(30)
While the typical interpretation of this equation is the inner product of w and the trigonometric polynomial with frequency k, we can equivalently view it as the evaluation of the sum of trigonometric polynomials with frequencios n evaluated at & weighted by w,- Thus, the resulting frequency response is non-sparse but has very small number of degrees of freedom a result of the sparsity of the spatial kernel. In domains where the spectrum rapidly decays, the learning process might then favor using these degrees of freedom to learn patterns relevant to the larger magnitude low frequencies and largely ignore the significantly smaller magnitude high frequencies.
A.5 DFS transform details
Here provide additional details of the DFS method (Merilees, 1973; Orszag, 1974). For a field f defined on an equirectangular grid, the DFS representation of f is given by f m| R:
01)+,,0), ( (0,m]
=(),0)=- - A(2,0), (1,0) [0,7]
a(A+=,-0)) (1,0) [0,7] [-T,
h(*+T,-0),
= x [-m,0]
where g(A, 0) = f(A n, 0) and h(>, 0) = f(A,0) The transformed function f is now periodic in both the A and 8 directions. As demonstrated by Orszag (1974), the 2D Fourier basis equivalent a complex expo- nential expansion in the zonal direction and alternating anne/oosiee expansions in the meridional direction. Subsequently, much of the recent. work in the area focuses on improved basis representations for the DFS
23

---

Published in Transactions on Machine Learning Research (11/2023)
Table 1: Computational costs associated with changes. These are theoretically timings and do not reflect hardware realities and relative costs of different layers in the network. Dimensions are N for the number of points on each spatial axis (assuming H-W=NN, K for the kernel size of a spatial convolution, and D for the number of channels in the hidden dimension.
CHANGE
PARAMETERS
FLOPs
OLD
NEW
OLD N²
NEW
DENSE - DS CONV
D²
N + D
N' + +D) log + ND D2) NO +3D²)
DFS DFS HYBRID
2N D
D + K
2N log + N2 D?)
STATIC + DYN CONV
N
2D N +
N² *F(+DD) D
method. However, we choose to use the naive function expansion as it simplifies the implementation and allows for more natural handling of vector-valued fields in which we just flip the sign of any field directed across the axis of reflection.
We note the DFS method provides additional advantage of flexibility in patchified settings like that of the AFNO used in FourCastNet (Pathak et al., 2022). In the patchified sctting, the extremal north/south "tokens" " are no longer representative of the poles, but rather of large regions surrounding the poles. Elficient spherical harmonic transforms are oblivious to this. As an example, the asymptotically optimal Driscoll- Healy transform (Driscoll & Healy, 1994) employed by TS2KIT (Mitchel et al., 2022), which we use to implement spherical architectures in Section 5, for instance, assigns a quadrature weight of zero to polar samples. This amounts to discarding the information from the entire patch, which for NeurCast contains 10,080 non-polar samples per variable. This limitation does not exist with the DFS method as we have full control of zonal and mcridional bandwidth independently and can adjust bandwidth restrictions to account for the widest region of the patch.
B Computational Comparison
While we include timings in Table 1, we include here a table describing the FLOP/paramen comparison on a block-wise basis the strict timings are affected by memory movement and sequential kernel launches and may not represent an optimized version. In practice, as seen in Table 3, switching to DS convolution results in significant speedup despite the slight increase of FLOPs while the Hybrid convolution results a larger slowdown despite the theoretically smaller increase. The practical runtimes could further change complex support, especially complex mixed-preee pixcc-prrcision support, improves in deep learning librarics.
C Experiment Details
C.1 Navier-Stokes
C.1.1 Data
Data generated using code provided by Tran al. (2023) to match the settings described by Kochkov al. (2021). It solves the Navior-Stokes equations (nondimensionaalized with Reynolds number Re 1000) for a system with constant Kolmogorov forcing f = 4 cos 1y))==J11 .lu across initial conditions. The trajectories are generated using pseudospectral method at 2048 x 2048 resolution and is integrated forward at At .0002 by a fourth order Carpenter-Kennedy method. 32/4/4 trajectories were generated for train/ haann/aald//tsss with 9764 snapshots per trajectory for a total of (000088/89966//////
examples.
For the task described Li et al. (2021b) the sampling occurs at intervals of At as this large time-step provides challenge for numerical integration and thus performing well at such a timo-stop demonstrates a significant speed improvement for the deep learning approach. For our purposes of demonstrating autore- gressive growth, we instead sample at. Al .0035. The timestep entirely chosen due to the fact that this
24

---

Published in Transactions on Machine Learning Research (11/2023)
the interval at which the generation code saved snapshots by default. For consistency with the larger scale tasks, also predicted e values as opposed to vorticity.
C.1.2 Model Configurations and Training Details
Several architectural choices were made for consistency across models. In each network, appended grid coordinates were replaced by fully learnable position embeddings similar to those used in Pathak et al. (2022). Furthermore, RcLU activations were replaced by the continuously diffcrentiable alternative SiLU (Elfwing et al., 2017). Apart from these changes, each network was generated using the four-layer variant released by the authors of the respective comparison models with hidden dimension of 128.
The ReFNO only uses the modifications from Sections 4.1/4.2 as the domain for this problem is the 2-Torus modifications to account for spherical geometry unnecessary. As this is smaller problem, the depth- separable implementation did not provide enough of a memory savings to increase the hidden dimension as we do in for the Shallow Water example (see: C.2), so we instead augmented the spatial MLP with the bottleneck structure popularized by Sandler et al. (2019) so that run-time was comparable between models.
All models were trained using identical settings. To minimize the impact of hyperparameter tuning, we used the automated learning rate search provided by Defazio & Mishchenko (2023) code for Adan (Xie et al., 2023). During each training step, one input snapshot of predicting the output at to + At. These were optimized using mean-squared error rescaled by the mean norm of the dataset to reduce the impact of precision issues for 20 epochs at batch size of 128 per run. Error bars are the 95% confidence intervals produced by sampling over 7 individual training runs per model evaluated across all initial conditions in the test set.
provided to the model for the purposes
C.2 Rotated Shallow Water
C.2.1 Data
We generate the data from the hyperviscous, forced shallow water equations:
@u(x, u(x, t). V (x - Vv Vu 2Ω u(x,t)
(31)
Ot
Oh - u(x, t) V (hu) vvth+ Ōt
(32)
where v is the hyperdiffusion coefficient, Ω is the Coriolis parameter, a is the velocity fields, H is the mean height, and denotes deviation from the mean height.
The forcing F is defined:
def
(t)) time of day = / day
time of year oearerr_-aaa max declination lon center = time of day*2+up. pi # Rescale sin to 0-1 then scale to up pi maammmax__eeeeee lat lat_center = np. sin (time_ot of yearr422ppppp pi) max declination
year .4 # Truncated from estimate of earth 's solar decline
lon anti = up. pi + lon center #2*np. #2=np.((mp sin sin(-time_od of day+2*mp pi)+1) / 2) s)ppi pi return center, lat lat_een lon antin lat center
def season day _or_forring// (phi, theta, t, f0): phi c, theta c, phi a, theta find center(1)
sigma up. pi/2
coefficients up. cos
cc *
-thtta = sigma ++2)
forcing = h 1o coefficients
return forcing
25

---

Published in Transactions on Machine Learning Research (11/2023)
This is not designed to mimic an exact physical process, but rather to force some level of daily/anmus pattern. It consists of two Gaussian blobs centered on opposite sides of the planet that circle on a daily basis. Each Gaussian increases/ picreases/decreasee decreases h with the intensity of the forcing increasing towards the center. The axis of rotation for these Gaussians varies over a model year which we define to be 1008 model hours.
Integration is performed forward in time using semi-implicit RK2 integrator. Step-sizes are computed using the CFLL-checker in Dedalus. The 3/2 rule is used for de-aliasing. Background orography taken from earth orography and passed through mean-pooling three times (until the simulations became stable empirically). Hyperdiffusinn is matched at L 96.
Initial conditions are randomly sampled from ERA5. u, U, are taken from the hpa 500 level with a used as h is the shallow water sel-up. We found that these conditions did not produce stable rollouts at the given hyperdiffusion level without pre-filtering the data, so prefiltering was performed by executing ten iterations of 50 steps followed by solving balance BVP. This was likely more than necessary, but given the quantity of data generated, we wanted to avoid in-depth manual inspection of the data.
As the diffusive system is decaying over time, we initialized each run with 500 burn-in hours then took the next 3 model years of data for total of 3024 samples per year. In total, the training set consisted of 25 trajectories of length 3024 for a total of 75,600 samples. Validation and test used an additional 2 and 3 trajectories respectively.
C.2.2 ReDFNO Architecture
The ReDFNO utilized four blocks structured as described in Figure 3. Each block is defined with embedding dimension 192. Our frequency-domain spectral normalization is used for spectral layers with conventional spectral normalization is used for spatial layers. Hybrid convolutions with DFS padding are used for the spectral block. We found a slight boost to stability by using a nonlinear spectral block as used in the AFNO, This structure is described in Figure 3. In total, this network contains approximately million paramcters.
We utilized time-dependent position embeddings. Position embeddings have previously been used for neural operators in both the AFNO (Guibas et al., 2021) and FCN, However, in cases where external forcing data is not directly available, static position embeddings are unable to fully represent seasonal behavior in earth systems. To address this, we introduce dynamic, time-dependent position embeddings. These embeddings are generated through the use of a shallow MLP that combines static embeddings with either sinusoidal features representing daily and yearly periods or more informative features directly computed for given problem like solar declination.
C.2.3 Comparison Configurations
As our model used 17 million parameters, we tuned and augmented comparison models until the sizes were roughly comparable when scaling laws allowed it, Otherwise for FNO based models, we used as similar settings as possible, though the vastly different parameter scaling led to higher parameter counts. We list model dimensions in Table The S2CNN model uses filters of [20, 16, 12, 8) for interpolates, and the DeepSphere uses kernel size 5 rather than the default 3, baseline uses ConvNext (Liu et al., 2022b) blocks with three blocks per stage instead of the classical blocks from Ronneberger et al. (2015) as we felt modernizing the architecture was important for fair comparison. We found this worked considerably better in tests.
found these to achieve better results. Our UNet
Note that persistence and climatology are two constant forecasts. Persistence uses the initial condition as the forecast while climatology implies using the mean estimate over the dataset.
C.3
ERA5
C.3.1
ReADFNO Architecture
This network uses FourCastNet as a base. The core processing unit consists of 8 AFNO blocks using our factorized nonlinear spectral convolutions rather than the standard complex MLP. As the AFNO architecture
26

---

Published in Transactions on Machine Learning Research (11/2023)
Table 2: Model dimensions for SWE experiments, along with the 1-step mean absolute error (MAE) and 100-step MAE. All MAE are reported in units of 10-2.
MODEL
STAGES
HIDDEN DIMENSIONS
#PARAMS
1-STEP MAE
AvG. 100-STEP MAE
UNET
4
[64, 128, 256, 512]
16M
.57
31.4
FNO
4
64
536M
0.80
19.2
SNO
4
64
268M
0.65
26.6
RESNO
4
192
9M
1.18
22.6
ReDENO
4
192
17M
0.49
4.1
S2CNN
[64, 128, 256,512]
18M
1.63
36.0
DEEPSPHERE
5 [64, 128, 256, 512, 1021]
38M
1.74
NAN
Table 3: Ablation path for SWE experiments evaluated validation set. In descending order, each adds new feature on top of previous inclusions until we reach the full stabilized architecture. Step is one training step (forward and backward) at batch size 16 on a NVIDIA A100 GPU. (DS=depth separable con- volution, DFS=Double Fourier Sphere transform, SN=Spectral Normalization, Rath=Reorder Shaping=Latitude-wiss spectral truncation)
MODEL
PARAMS
STEP (s)
MSE (T=1)
MSE (T=20)
MSE (T=40)
MSE (T=80)
FNO
536M
8.1u-5
2.6E-4
1.0p-2
INF
+DS (x3 WIDTH)
25M
.22
1.28-4
1.0E-2
INF
INF
+DFS
27M
.40
8.28-5
INF
INF
INF
SNN
27M
.46
8.7u-5
4.5k-
8.0g-4
1.45-3
PATTH
27M
.46
9.5g-5
3.6E-4
5.6E-4
1.0E-3
SHAPING (ReDFNO)
17M
.41
8.0e-5
2.3e-4
3.7e-4
6.3e-4
is more complication and non-isotropic, we made several further changes to eliminate aliasing-inducing operations:
Downsampling: Strided convolutions perform decimation which is addition in the spectral domain. To avoid this, downsample via. Fourier cetepolation/trunaatte in the DFS representation. Dur- ing downsampling, channels are expanded via 3x3 spatial convolution. This occurs at 1///////////// downsampling levels with the DeADFNO operations occuring at 8r downsampling as in FourCast- Net.
2. Upsampling We introduce a Unet-like structure here. At each intermediate resolution, the a downsampled version of the original data added back to the upsampling path then fed through a single layer spatial CNN. The spatial CNN used here as the spectral convolutions become increasingly memory-hungry as resolution increases so we can make best use of them at the lowest resolution.
LayerNorm - InstanceNorm LayerNorm divides cach pixel by a spatially varying standard devi- ation estimate. While we've observed the impact to be minor, this does produce new frequencies. We therefore replace it with InstanceNorm which computes the standard deviations over the spatial dimensions of channel rather than the channel dimensions of particular pixel.
The model was trained using the Nvidia Apex implementation of LAMB (You et al., 2020) following the schedule of FourCastNet. We found that the gradient rescaling in LAMB acted similarly to a trust region and avoided training instability we experienced with Adam in this case. However, we were able obtain similar partial training performance using version of Adam with added step size constraints, but chose stick with the known optimizer for this work.
27

---

Published in Transactions on Machine Learning Research (11/2023)
C.3.2 Comparison Configurations
NourCastNet comparisons use the pre-trained model weights provided by the paper authors. Similar SWE, the UNet comparison uses ConvNext, blocks rather than conventional ResNet structure. The UNet was trained using the Adam Kingma & Ba (2017) optimizer with learning rate selected by grid search over [le 5, be 5, Ic 5, 1, 3]. Epochs and schedule follow scttings in FoutCastNct.
The smaller number of comparisons here is due to the fact that the majority of the architectures used in the SWE experiments are incapable of scaling to the larger input sizes without seriously diminishing network capacity.
C.3.3 Additional Results
Full field by field RMSE can be found in Figure 9. Convergence to persistence was obtained for all fields except for Surface Pressure and sometimes Z50. In the FourCastNet training recipe, all fields (normalized) are weighted evenly. Surface Pressure has relatively low snapshot-loo-sappsholl variation but has significant outliers as it tends to be proportional to altitude. In the sertaaliedd-ut-not weiihtet scheme, this tends to lead to crrors in this field being undervalued. Z50, on the other hand, is slowly varying smooth field, but is quite distant (5KM) from the next closest pressure level recorded in this ERA5 subset. We suspect this underperforms for this reason.
28

---

Published in Transactions on Machine Learning Research (11/2023)
ReDFCN
FCN
UNet
Climatology
Persistence
u10
v10
t.2m
7.5
7.5
1500
RMSE
1000
2.5
2.5
500
0.0
1850
u1000
v1000
1500
7.5
10
10
RMSE
1000
500
2.5
0.0
0
7.1000
u850
v850
z850
1500
1500
RMSE
1000
10
1000
500
500
0
u500
v500
z500
t500
20
2000
7.5
RMSE
10
10
1000
2.5
z50
r500
r850
towv
15
4000
RMSE
2000
20
0
Step
Step
Step
Step
Figure 9: The restructured model matches the performance the two-step-ahead trained FCN, but demon- strates improved stability and asymptotically converges to the persistence forecast performance for all but the surface pressure field.
29


Reviews:
[{"review_id": "xgFOnhplVx", "reviewer": "Reviewer_Cgrj", "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\n1. The paper addresses the growing research field of ML for scientific applications, particularly weather prediction.\n2. Overall, The method looks sound and reasonable, and the improvements are substantial in some aspects of the applications tested. \n2. The experiments are thorough and adequate, the authors compare with a few baselines and performed ablation studies, clearly showing the contributions of each modification from FNO\n\n## Weaknesses\n1. The paper is very hard to follow. I really struggled to get a faint picture of the paper, and I am sure I will be missing technical details unless the authors submit a much-improved revision. This is probably due to my background being in machine learning rather than numerical methods or signal processing, so I might be missing some details, but for an ML journal and audience, I personally think the writing needs to be significantly improved with sufficient details presented more clearly. I feel like many important parts of the paper, such as Sec 4.2, are personal notes of the authors that are only accessible to experts. The paper has to be enriched with equations and clearly defined variables.\n2. If I understand correctly, the problem of aliasing arises during spectral convolution, and how this is related to Equations 3 & 4 of FNO is not clear at all. Is the matrix multiplication in Equation 4 the convolution? \n3. The claim that the model can achieve 800% improvement in stable prediction range feels like a overclaim. This is established on the results of a single channel out of 20 in the data. The prediction of other channels compared to the baselines need to highlighted, unless the field generally agrees to look at the best-predicted channel.", "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 paper proposes a new network architecture for autoregressive modelling of spatiotemporal physical systems described by partial differential equations. The authors first analysed aliasing issues in popular neural architectures, such as the Fourier Neural Operator, through the lens of pseudospectra methods. A key insight is that convolutions in frequency domains, a popular operation in these models, introduce\naliasing. The authors then proposed a few architectural changes and other application-specific designs. The experiments on simple synthetic and larges-scale real datasets show improvements of their methods compared to other methods. ", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Critical: please include these changes or address the questions. \n1. Please clearly establish the connection between the pseudospectral methods and the FNO in Eqns 3,4. Are these equations implementing 1x1 convolution in temporal and spectral spaces, mentioned later? \n2. The authors try to explain how FNO maps to Euler's method. Please write down how each part of (3) maps to the Euler, in equations. \n3. Write an equation for what is happening in **Depthwise Separable Spectral Convolution**. Which part of (3) maps to the separable spectral convolution, if they are related?\n4. Explain how spectral normalization can be implemented with polar form. I'm sure this is simple enough for the authors to explain with equations. \n5. Explain clearly what's done in the methods in Section 4.2. I have zero idea about what's happening here.\n6. Section 4.3 is better than 4.2, but I am still a little confused. Are the two bolded methods addressing the same or related issue caused by polar coordinates causing different spatial frequencies? \n\nOthers:\n1. The last sentence of the top paragraph on page 6 mentions distributional shift, but how is this dealt with? Did I miss this somewhere?\n2. It would be very useful to include arguments of functions in all differential equations (1,2,6,7)\n3. The term \"residual\" is mentioned in FNO context and also the residual network context. The former needs clarification. ", "broader_impact_concerns": "The ethical implications are discussed, and I think this is adequate. ", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "1OzQLil3S3", "reviewer": "Reviewer_gvhB", "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\n1. The key strength of this article is that it attempts to address the issue of unstable and unphysical behavior of geophysical surrogates during long-term rollouts. These limit the accuracy and computational gains of surrogates for practical applications and fall short of certain important challenges in the atmospheric forecasting community - subseasonal forecasting. I encourage the authors to continue their research in this direction. I also think the idea of adaptively learning filter transfer functions is very interesting.\n\n2. The connections drawn by the authors to well-known ideas in spectral methods and numerical analyses are timely and connect two rich fields. That being said, identifying aliasing errors and actually handling them appropriately are two different things and the techniques proposed by the authors are quite valuable for surrogate modeling of dynamical systems. Connections to computational complexity for different techniques are also appreciated.\n\n3. The experiments, appendix, and code provided by the authors are extensive and thorough and I commend them for encouraging the reproduction of their results.\n\n\n\nI do not have specific weaknesses but list out a few comments/questions below:\n\n1. Why does the UNet remain stable for the experiments in section 5.2 but diverge in Section 5.1? Have the authors looked into this? I'd be curious to learn why this is the case.\n\n2. Have the authors looked into how the filter transfer functions for their NL Spectral convolution evolve over time? Is there any connection to the typical frequencies expected in their physical experiments.\n\n3. What is the training and memory overhead (or benefit) of the authors proposed additions to the standard FNO? Is this a-priori quantifiable for a given problem?\n\n4. Can the authors comment on how a pretrained model with stability issues may be fixed in a post-hoc manner using some insights from this research?\n\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": "The authors investigate various techniques to improve the long-range forecast stability of Fourier Neural Operator techniques. They rely on insights from spectral and pseudospectral methods to identify architectural augmentations that lead to more accurate roll-outs of their techniques. The study is very well put together, has significant potential for impact, and I support its publication. ", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "See comments.", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "u1rd2h7J2F", "reviewer": "Reviewer_Hdii", "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\nThe paper studies an advanced topic with immense societal impact. Its proposed methodology sketched out in Figure 3 is simple, intuitive, and novel. The results reported on the ERA5 data set are intriguing. The 100 and 400 step results show a clear improvement over FCN. This may be a significant outcome for climate researchers. Figure 1 demonstrates the limitations of FCN very clearly, providing solid motivation for the proposed solution. Likewise, the illustration of the aliases in Figure 2 is also helpful to the reader. The paper provides a comprehensive survey of the related literature and positions its contribution properly.\n\n**Weaknesses:**\n\nI am not confident that the current shape of the paper speaks directly to the machine learning audience. Although it uses a good deal of machine learning tools, it builds on many concepts outside the scope of a typical machine learning researcher. The paper either assumes these terms to be known or gives too brief an explanation of them for an outsider to climate research. Example terms with this character: collocation point, advection equation, spectral differentiation matrix, pseudospectral methods, double Fourier sphere method, geometry-aware filter. It is also not clear in the storyline where exactly PDEs come to play and how they are related to the workflows given in Figure 3. The paper would also benefit from a discussion about the generalizeability of its findings to other modalities than climate data, and if possible also beyond spatiotemporal dynamics modeling.", "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": "The paper introduces a methodology for developing neural autoregressive models that can make long-term forecasts without accumulating integration error resulting in a total collapse. The paper takes modeling spatiotemporal dynamics, in particular Earth climate dynamics, as a use case. It reports strong results on the ERA5 data set, that is often treated as the most advanced climate model used by climate researchers.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "No", "audience_explanation": "", "requested_changes": "The following changes would improve the paper:\n\n - Provide a more detailed and more accessible background section where all terms not familiar to the machine learning research readership are properly and simply explained. See a nonexhaustive list of these terms under the “weaknesses” paragraph above.\n\n - Discuss the generalization aspects of the proposed methodology, if possible illustrated on a small additional data set.\n\n - Clarify the relationship of PDE modeling to the rest of the material more clearly.", "broader_impact_concerns": "The paper addresses the broad impact concerns in Section 2 under the “societal impact” paragraph. However, the provided discussion is from the perspective of the ethical risks associated with the studied problem setup and not about the particular solution the paper presents.\n", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The reviewers agreed that this is an interesting paper dealing with an increasingly important problem. While they initially had concerns about the accessibility of the paper to the machine learning community, the authors have addressed these in the new version.

Requested changes
1. Please incorporate the remaining changes promised to the reviewers, i.e. update the notation in the appendix.
2. Expand the caption in Figure 3 to describe each of the subfigures, and give them individual labels.
3. Update the last paragraph of Section 4.2 to refer to Figure 3 instead of Figure 8.
Decision: Accept (poster)