On the Benefits of Memory for Modeling Time-Dependent PDEs
Authors: Ricardo Buitrago
Tanya Marwah
Albert Gu
Andrej Risteski
Venue: ICLR 2025 Oral (2025)
Keywords: State Space Models
Partial Differential Equations
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)

Abstract: Data-driven techniques  have emerged as a promising alternative to traditional numerical methods for solving PDEs. For time-dependent PDEs, many approaches are Markovian---the evolution of the trained system only depends on the current state, and not the past states. In this work, we investigate the benefits of using memory for modeling time-dependent PDEs: that is, when past states are explicitly used to predict the future. Motivated by the Mori-Zwanzig theory of model reduction, we theoretically exhibit examples of simple (even linear) PDEs, in which a solution that uses memory is arbitrarily better than a Markovian solution. Additionally, we introduce Memory Neural Operator (MemNO), a neural operator architecture that combines recent state space models (specifically, S4) and Fourier Neural Operators (FNOs) to effectively model memory.  We empirically demonstrate that when the PDEs are supplied in low resolution or contain observation noise at train and test time, MemNO significantly outperforms the baselines without memory---with up to $6 \times$ reduction in test error. Furthermore, we show that this benefit is particularly pronounced when the PDE solutions have significant high-frequency Fourier modes (e.g., low-viscosity fluid dynamics) and we construct a challenging benchmark dataset consisting of such PDEs.

Published as a conference paper at ICLR 2025
ON THE BENEFITS OF MEMORY FOR MODELING TIME-DEPENDENT PDEs
Ricardo Buitrago Ruiz!? Tanya Marwahl, Cartesiaa Al
Albert Gu! Andrej Risteski'
Carnegie Mellon University
ABSTRACT
Data-driven techniques have emerged as promising alternative to traditional numerical methods for solving PDEs. For time-dependent PDEs, many ap- proaches are Markovian the evolution of the trained system only depends on the current state, and not the past states. In this work, we investigate the benefits of using memory for modeling time-dependent PDEs: that is, when past states are explicitly used to predict the future. Motivated by the Mori-Zwanzig theory model reduction, we theoretically exhibit examples PDEs, in which solution that memory is arbitrarily better than Markovian solution. Additionally, we introduce Memory Neural Operator (MemNO), a neural operator architecture that combines recent state space models (specifically, S4) and Fourier Neural Operators (FNOs) to effectively model memory. We em- pirically demonstrate that when the PDEs are supplied in low resolution contain observation noise at train and test time, MemNO significantly outperforms the baselines without memory- wwith up to 6x reduction in test error. Furthermore. we show that this benefit is particularly pronounced when the PDE solutions have significant high-free tighffreeueccy Fourier modes (e.g., low-viscosity fluid dynamics) and we construct challenging benchmark dataset consisting of such PDEs.
simple (even linear)
INTRODUCTION
Time-dependent partial differential equations (PDEs) are central to modeling various scientific and physical phenomena, necessitating the design of accurate and computationally efficient solvers. Re- cently, data-driven approaches based on neural networks (Li et al., 2024b; Lu al., 2019) have emerged attractive alternative to classical numerical solvers, such finite element and finite difference methods (Le Veque 2007). Classical approaches are computationally expensive in high dimension and struggle with PDEs which are sensitive to initial conditions. Learned approaches can often negotiate these difficulties better, at least for the PDE family they are trained on.
One example of data-driven approach is learning neural solution operator, which for time- dependent PDE learns to predict future states based on previous ones (Li et al., 2021; 2023a). Recent works (Tran al., 2023; Lippe al., 2023) suggest that optimal performance across various PDE families can be achieved by conditioning the models only on the immediate past state-i.c.. treating the system as Markovian. In contrast, other works propose architectures that explicitly use memory of past states (Li et al., 2021; 2023a; Hao et al., 2024). However, none of these works elucidate whether and when modeling memory helpful.
In this work, we demonstrate that when the solution of the PDE is only partially observed (e.g. observed at low resolution), explicitly modeling memory can be beneficial. Partial observability is natural in many practical settings. This could be due to limited resolution of the measurement devices collecting the data, inherent observational errors in the system, or prohibitive computa- tional difficulty in generating high-quality synthetic data. This can lead significant information loss, particularly in systems like turbulent flows (Pope, 2001) or shock formation in fluid dynam- ics (Christodoulou, 2007), where PDEs change abruptly in space and time. In such situations, classical results from dynamical systems (Mori-Zwanzig theory), suggest that the system becomes strongly non-Marken non-Markovian.
More precisely, Mori Mori- -waaziz theory (Mori, 1965; Zwanzig, 1961; Ma et al., 2018) is an ansatz to understand the evolution of subspace of system (e.g., the span of the & largest Fourier compo-

---

Published as a conference paper at ICLR 2025
nents). Under certain conditions, this evolution can be divided into Markovian term (the evolution of the chosen subspace under the PDE), a memory term (which is a weighted sum of the values of all previous iterates in the chosen subspace), and "unobservable" term, which depends on the values of the initial conditions orthogonal to the selected subspace.
The main focus of this paper is studying when explicitly modeling this memory term is useful. We give an example of a very simple (in fact, linear) PDE where we show theoretically that the solution which takes into account the memory term can be arbitrarily better than the Markovian solution. We also provide a way to operationalize the Mori-Zwanzig formalism by introducing Memory Neural Operator (MemNO), a neural operator architecture that combines a Markovian operator to model the spatial dynamics of the PDE (such as the Fourier Neural Operator (Li et al., 2021; Tran et al., 2023)), and sequence model to maintain a compressed representation of the past states (such as the S4 state space model (Gu et al., 2022; 2023)). We show that MemNO outperforms its Markovian (memoryless) counterpart in PDEs observed on low resolution grids or with observation noise achieving up to 6x less test error. Our contributions are follows:
We identify a setting in which explicitly modeling memory is helpful: namely, when there is a combination of lossy observations of the solution of the PDE (e.g., due to limited resolution or observation noise) and significant contributions from fregueryyy Fourier modes in the solu- tion.
Even in simple PDEs, we theoretically show the memory term can result in solution that is (arbitrarily) closer to the correct solution, compared to the Markovian approximation -in particular when the operator describing the PDE "mixes" the observed and unobserved subspace.
Across several families of one-dimensional and two-dimensional PDEs, we empirically demon- strate that when the input is supplied on low-resolution grid, or contains observation noise, neural operators with memory outperform Markovian operators by a significant margin. More precisely, to operationalize memory, we introduce MemNO, a neural operator architecture combining Fourier Neural Operators (FNOs) and S4, which achieves the best performance across several Markovian and memory baselines.
We observe that many current benchmarks for PDE solvers predominantly include PDEs in which there is little contribution from high-frequency Fourier modes. Consequently, we construct more challenging datasets where the solutions have significant high-frequency modes, which believe will be of broader significance to the community beyond testing the effects of memory- especially given recent meta-studies suggesting many current PDE benchmarks are too easy (McGreivy & Hakim, 2024).
2 RELATED WORK
Data-driven neural solution operators (Chen & Chen, 1995; Bhattacharya et al., 2021; Lu et al., 2019; Kovachki et al., 2023) have emerged the dominant approach for approximating PDEs, given their ability to model multiple families of PDEs at once, and their computational efficiency at inference time. Many architectures have been proposed improve their performance across different families of PDEs: Li et al. (2021) introduced the Fourier Neural Operator (FNO), a resolution invariant architecture that uses a convolution- (2023) later introduced the Factorized FNO (FFNO) architecture, which builds upon and improves the FNO architecture by adding separable spectral layers and residual connections; Cao (2021) proposed a Transformer method with linear attention over the spatial sequence; other recent works have used U-Net-based architectures (Gupta & Brandstetter, 2023; Rahman et al., 2023).
integral kernel evaluated in the Fourier space: Tran et al.
Focusing on memory, Tran et al. (2023) performed ablations that suggest the Markov assumption optimal and outperforms models that use the history of past timesteps as input. Lippe et al. (2023) performed a similar study for long rollouts of the PDE solution and concluded the optimal performance is indeed achieved under the Markovian assumption. We show that these findings can be replicated only when the resolution of the observation grid is high. On the other hand, we show that MemNO effectively models memory to achieve much superior performance than Markovian operators in low resolution, while not dropping performance in the high resolution case.
Our work is motivated by the Mori Swai-Zzanzig formalism (Zwanzig, 1961; Mori, 1965) which shows that a partial observation of the current state of the system can be compensated using memory of
2

---

Published as a conference paper at ICLR 2025
past states. Ma et al. (2018) draws parallels to the Mori-Zwanzig equations and LSTM (Hochreiter & Schmidhuber, 1997) to model the dynamics of the & largest Fourier components of single PDE. However, in our work, study the benefits of memory in neural operators that learn the dynamics of an entire family of PDE. Furthermore, we show conditions under which not maintaining memory result in arbitrarily large errors.
3 PRELIMINARIES
First, we introduce several definitions, as well as the Mori-ZWanzii formalism applied to our setting.
3.1 PARTIAL DIFFERENTIAL EQUATIONS (PDES)
Definition (Space of square integrable functions). For integers d, V and an open set Ω C Rd, define L (2;: R as the space of square integrable functions u : S2 - R such that || S oo,
where ||u]| (Sp [1u(z)112dd))
Notation 1 (Restriction). Given function Ω + RRV and subset A C Ω we denote 4nd as the restriction u to the domain A, i.c. UI : + RRV.
The general form of the PDEs we consider this paper will be the following:
Definition 2 (Time-Dependers PDE). For an open set Q C Rd and an interval T] C R, a Time- Dependent PDE is the following expression:
du
Øt (1, x)) = [(,](1,
1L E [0,T]. E 2
(1)
u(0, x) = un uz(x),
Vx E Ω Vt [0,
(2)
Blu
O,
where L: L (2; RV) + L (2; differential operator in æ which independent of time, un(x) E 1,2 (9; R and B is an operator defined on the boundary of an, commonly referred as the boundary condition.
In our theory and experiments, we will work with periodic boundary conditions (for precise defi- nition, see Definition 6). Finally, we will frequently talk about a grid of a given resolution:
Definition (Equispaced grid with resolution f). Let Ω = [0, Ljd. An equispaced grid with resolu- tion ff in the following set S C Rdd:
s (
os for 1 s
We will also denote by |S] the number of points in S.
3.2 MORI-ZWANZIG FORMALISM
The Mori-Zwanzig formalism (Zwanzig, 2001) considers the setting in which an equation is known for full system, but only a part of it is observed. It leverages the knowledge of past states of a system to compensate for the loss of information that arises from the partial observation of the current state. In work, partial observation only observing the Fourier modes up to a critical frequency. In the context time-dependent PDBs, the Mori-Zwanzig principle is formalized as the Nakajima-Zwanzz equation (Nakajima, 1958).
refer to observing the solution at discretized grid in space or
We will give an overview of the -awazi equation and set up the notation for the rest of
the paper. Assume have a PDE as in Definition Let P L (1; RY) L² (Ω: RR) be linear projection operator. We define @ = 1 P, where / is the identity operator. In our setting, for the PDE solution at timestep u, 2 (Ω; RV). P[ue] is the part of the solution that we observe and Q[u,] is the unobserved part. Thus, the initial information receive for the system is P (moo)-

---

Published as a conference paper at ICLR 2025
Applying P and Q to Equation 1 and using u Plu] + Q(ui, we get:
0
(lu](t, = PL]u)((, (
(4)
0t
0 [tu](t, @C(ul(t, x) ec
(5)
Solving for 5 yields 2)) [xxp)) ())[tt=) + ec (to
Plugging into 4, we obtain a Generalized Langevin Equation (Mori, 1965) for Plu]:
0 = + +PL /'" - ())ps))(((((()s))dss PLC PCC Q (uuo((t,)) Ot
We will refer to the first summand on the right hand side of Equation 6 as the Markovian term because it only depends on (t, .), the second summand as the memory term because it depends on Plu] (s. x) for I $ I t, and the third summand as the unobserved residual as it depends on Qluo] which is never observed.
Since Equation 6 is exact, is equivalent to solving the full system. The term that is typically most difficult compute is the memory term, and many methods approximate have been proposed.
In the physics literature, some techniques include a perturbation expansion of the exponential exp Qt s)} (Breuer & Petruccione, 2002), approximations using operators defined in P [ | (2; R (Shi & Geva, 2003; Zhang et al., 2006; Montoya-Castillo & Reichman, 2016; Kelly et al., 2016). In the classical numerical PDE solver literature, the memory term has been approx- imated by leveraging the structure of the orthogonal dynamics of the P semi-group (Gouasmi et al., 2017), and the Mori-Zwanzig formalism has been applied to a variety of fluid dynamics PDEs (Parish & Duraisamy, 2017). In the machine learning literature, some works approximate the mem- ory term with a neural network, which is then used as a part of a hybrid PDE solver (Ma et al., 2018; Beck et al., 2019; Pan & Duraisamy, 2018). Gupta Lermusiaux (2021) approximated both the Markovian and memory term with a neural network, yet the method required deriving and coding adjoint equations to perform backpropagation. In this work, we explain when modeling memory is expected to be helpful, and introduce a neural operator that leams to model the temporal (e.g. memory) and spatial dynamics of a PDE directly from data.
THEORETICAL MOTIVATION FOR MEMORY: A SIMPLE EXAMPLE
In this section, we provide simple, but natural example of (linear) PDE, along with (in the nomenclature of Section 3.2) a natural projection operator given by a Fourier truncation measure- ment operator, such that the memory term in the generalized Langevin equation (GLE) can have an arbitrarily large impact on the quality of the calculated solution. We will work with periodic functions over [0, 2n] which have a convenient basis:
Definition 4 (Basis for 2nT -pripoic functions). function : -> R is 2nT 2- if f (x). We can identify -prriooo functions with functions over the torus T == {e : E ER S C by the map )e) = f(r). Note that )e is basis for the set of 2m-perio functions.
We will define the following measurement operator:
Definition 5 (Fourier truncation measurement). The operator Pk (T;; (T; acts E (T; R), f(x) = -x eive as Pa > ene
For notational convenience, we will also define the functions {c, Jnez, where c, (x) := e ina + fine Now, consider the following operator define a linear time-dependent PDE:
Proposition 1. Let L (T; + L?( be defined as Lu(z) = -Au(x) + B (e + et ((()))) for B > 0. Then, we have:
A1 S n EN,, L(e.) n'e,, + B(en-1 ent1) & L(en) 2Be1
4

---

Published as a conference paper at ICLR 2025
The crucial property of this operator is that it acts by "mixing" the n-th Fourier basis with the (n 1)-th and (n + th: thus information is propagated to both the higher and lower-order part of the spectrum. Given the above proposition, can easily write down the evolution of a PDE with operator & in the basis Jez:
Proposition 2. Let C be defined as in Proposition 1. Consider the PDE
Ot att x) Zu(t.z)
u(0,0) Σ an(e))n
neNo
Let u(tu(tx)) = EnEro 1 Cn
en. Then, the coefficients (t) satisfy:
LA n
ōt a
=2aak a
B an- + an-1)
Ô (2) at 0
2Ba (1)
With this setup in mind, we will show that as B grows, the memory term in Equation can have an arbitrarily large effect on the calculated solution:
Theorem 1 (Effect of memory). Consider the operator L defined in Proposition I, the Fourier
truncation operator Pı, and let Q =1 1-PP P1. Let u(0, x) have the form in Proposition 2 for B
sufficiently large, and let on > 0, vn > 0. Consider the memoryless and memory-augmented PDEs:
0uj
PiLui
(9)
01
duz
Pi PiLuz + PIL S exp(([((-))) (000(((-))00000 ())2(u(())
(10)
with (0, = =u((0, x) PIu(O, x).. Then, 11 and u2 satisfy:
VI. > 0, Illl (4) u2(1)lla, 2 (1)la(()]ls 14 0, ||l (1) u22lllr, 2 Bt exp (22B))
(11)
(12)
Remark 1. Note that the conclusions the theorem mean that both the absolute difference, and the relative difference between the PDE including the memory term Equation 10 and not including the memory term Equation can be arbitrarily large B,t cc.
Remark 2. The choice of C is made for ease of calculation of the Markovian and memory term. Conceptually, we expect the solution to Equation 10 will differ a lot from the solution to Equation the action of the operator L tends "mix" components in the span of P and the span of Q
Remark 3. If we solve the equation will be on the order of exp(2Btt) This can be seen by writing the evolution of the coefficients of
Di (u(tx)) Lu(t, x) exactly,
&
calculate that [|u(t)]l La
au
u(() in the basis {en}, which looks like:
a = o
where O is roughly a tridiagonal
B
2
B
Toeplitz operator O
The largest eigenvalue of this oper-
0
B
(n + 1)²
B
ator can be shown to be on the order of at least 2B (Equation 4 in Noschese et (2013)). The Markovian term results in solution of order exp(v2Bt) Equation 19 and Equation 20), which is multiplicatively smaller by factor of exp - Bt). The result in this Theorem shows that the memory-based PDE Equation 10 results in a multiplicative "first order" correction which can be seen by Taylor expanding exp( (2B)) 2 /2Bt (V2B)R2 +

---

Published as a conference paper at ICLR 2025
5 EXPERIMENTAL SETUP
5.1 DATASET GENERATION
PDEs with high-frequeney Fourier modes: From the expression for the memory term Equa- tion 6 and the presence of high-frequency terms in the solution of the PDE of Theorem I, we should intuitively expect that memory will be most useful when the PDE solutions contain significant con- tributions from high-frequeng
Fourier modes'. Nevertheless, current benchmarks like PDEBench (Takamoto et al., 2023) rarely contain PDEs whose solutions have substantial high-frequency com- ponents,
quantitatively show in Appendix D. A solution which predominantly contains low- frequency Fourier modes can be accurately approximated by its Fourier truncation (Definition 5), so it be represented by solution @ u) in the notation Scction 3.2) should be small.
space, which implics that the unobserved part of the
Therefore, we construct new benchmark dataset which is specifically designed to contain PDEs in which the high-frequency Fourier modes have substantial contribution. Specifically, we gener- ate benchmark from solutions to the 6.1). In the case of Navier-Stokes (Section 6.2) and Burgers' equation (Section C), we directly take datasets from previous works. Details data generation procedure are provided in Appendix E.
equation with low viscosity (Section
Datasets with different resolutions: To construct our datasets, we first take discretized trajecto- ries of a PDE high resolution discretized spatial grid S C. Rd i.. u(t) R|SMR|, We then produce datasets that consist lower resolution versions of the above trajectories, i.c. on a grid SLR of lower resolution f. and show the performance of models that resolution. For 1-dimensional datasets, the discretized trajectory on SLR is obtained by cubic inter- polation of the trajectory the highest resolution grid. In 2D, the discretized trajectory is obtained by downsampling.
trained and tested such
5.2 TRAINING AND EVALUATION PROCEDURE
Task: Let u E C. (IO, T): I22 (2: be the solution of the PDE given by Definition Let be an equispaced grid in Ω with resolution f, and let T be another equispaced grid in [0, T' with N, points. Given uo(s)ss. our goal is to predict u(t, ()ts for E using a neural operator.
Training objective: As standard, we proceed by empirical risk minimization on dataset of trajectories. More specifically, given loss function L (RIS), RISi) R. dataset of training trajectories (()) im and parametrized maps G RIIS + RI for E T, we optimize:
N. Σ
θ = argming N
(u((,))00
up 1 (x) s)
N
Training and evaluation metric: Our training loss and evaluation metric is normalized Root Mean Squared Error (nRMSE):
ARMSE (u(t,z))s; a(()) |lu(t, r)is - a(t)
[u(t,,))ss1
where || ||2 is the Euclidean norm in SI
Further details on training hyperparameters are given in Appendix F.
5.3 ARCHITECTURE FRAMEWORK: MEMORY NEURAL OPERATOR
In this section we describe Memory Neural Operator (MemNO), deep learning framework to incor- porate memory into neural operators. A diagram is provided in Figure 9 and pseudocode in Figure 10.
Note, this is meant to be an intuitive rule-of of thumb rather than formal statement. In general, the "obser- vation" operator and the PDE will interact in complicated ways, but the combination of low-resolution grids and examining high-frequency components the Fourier basis seems very predictive in practice.

---

Published as a conference paper at ICLR 2025
Let NO, be neural operator with L layers, and denote NO," [uo] the prediction of the solution of the PDE at time t. We will assume that this Neural Operator follows the Markovian assumption, i.e. can write:
NO 2++1 [uo] = Yout o lz. o fz-1 o... o to o rin [N0][
(13)
where Tin R + R is the encoder and Tout : Risixh+t1 + R is the decoder; lj RiSix,, + RISIx are parametrized layers; and hy is the dimension of the j-th hidden layer. Essentially, the solution for each new timestep is obtained by applying the same L layers to the immediately previous predicted timestep.
Our goal is define a network Geeh that builds upon NO, and can incorporate memory. For this, we take inspiration from the Mori-Zwanzig formalism summarized in Section 3.2. Comparing Equa- tion 13 with Equation 6, we identify & f o ... o Ey with the Markov term which models the spatial dynamics. To introduce the memory term, we interleave an additional residual sequential layer M that acts on hidden representations of the solution at previous timesteps. Concretely, the MemNO architecture can be written as:
0 [uo] = Tou a L o o 4 Mol o... o Ey o rim (elloll9. G.. _[uo],... mn]
where- -1 < A: < L is a chosen hyperparameter." For notation, we will refer to (t') E Rliii as the hidden representation at the j-th layer for a timestep t' S t,, and (t) x) E Ih' as the value of such hidden representation at spatial point E S. Then, the spatial E layers are un- derstood to be applied timestep-wise wise,, i.e. lj [tt)) 2i (to)] == (t))...) (()))(((t)) and analogously for Tin and Tour- Thus, the E layers still follow the Markovian assumption. The memory is introduced through M, which is sequential model that uses the history of the previous timesteps to predict the next one. For computational efficiency, we consider a sequential model M x - A that is applied each element of the spatial dimension [S] independently, i.c. for each S, (4)), ((k) (to))) (x) == M[u(k) (), ..., ((8) (to, x)}.²
Note that our MemNO framework can be combined with any existing neural operator layer L, and with any (causal) sequential model M. Thus it provides a modular architecture design framework which we hope can serve as a useful tool for practitioners.
5.4 INSTANTIATING THE MEMORY NEURAL OPERATOR FRAMEWORK: S4FFNO
For our experiments, we introduce S4 Factorized Fourier Neural Operator (S4FFNO), which instan- tiates the MemNO framework by combining the Factorized Fourier Neural Operator (FFNO) (Tran et al., 2023) as the Markovian neural operator and S4 (Gu et al., 2022) as the sequential layer. We choose S4 models over recurrent architectures like LSTM (Hochreiter & Schmidhuber, 1997) due to superior performance in modeling long range dependencies (Gu et al., 2022; Tay et al., 2020), ease of training, and favorable memory and computational scaling with both state dimension and sequence length. An ablation comparing S4 to LSTM and Transformers is provided in Appendix G.
6 MEMORY HELPS IN LOW-RESOLLUTION AND INPUT NOISE: A CASE STUDY
In this section we present a case study for several PDEs of practical interest, showing that neural operators with memory confer accuracy benefits when the data is supplied in low resolution or with observation noise. We will use four Markovian baselines: Factformer (1D) (Li et al., 2023b), The Galerkin Transformer (GKT) (Cao, 2021), the U-Net Neural Operator (U-Net) (Gupta & Brandstetter, 2023), and the Factorized Fourier Neural Operator (FFNO) (Tran et al., 2023). For a memory-augmenter baseline, we consider the Multi Input Factorized Fourier Neural Operator (Multi input FFNO), which takes as input the last 4 timesteps of the solution of the PDE predict the next one, as proposed in the original FNO paper (Li et al., 2021), yet using the architectural design of FFNO. The architectural details for all the models are elaborated upon in Appendix B.
2 I. refers to inserting M after all the S layers, and k -1 refers to inserting M as the first layer. In Appendix G.2, we show our experiments are not very sensitive to the choice lt. We present analysis on some architecture modifications that model the local spatial structure more explicitly in Appendix I.
7

---

Published as a conference paper at ICLR 2025
aRMSE 4
Architecture
Uses memory
Resolution
KS
Burgers'
0.075
M 0.1
0.125
0.001
Factformer (ID)
X
0.436
0.391
0.149
0.190
GKT
X
0.588
0.601
0.314
0.356
U-Net
X
0,542
0,511
0.249
0.188
32
FENO
X
0,500
0.446
0.187
0.207
Multi Input FENO
0.364
0,308
0.092
0.099
S4FFNO (Ours)
0.139
0.108
0.031
0.053
Factformer (ID)
0.195
0.086
0.022
0.162
GKT
0.120
0.016
0.349
U-Net
0.747
0.062
0.022
0.171
64
FENO
0.107
0.033
0.004
0.146
Multi Input FENO
0.108
0.046
0.005
0.054
S4FFNO (Ours)
0.036
0.011
0.004
0.037
Factformer (ID)
X
0.058
0.030
0.017
0.117
GKT
X
0.028
0.013
0.007
0.307
U-Net
X
0.033
0.027
0.014
0.112
FENO
0.006
0.004
0.002
0.099
Multi Input FENO
0.057
0.052
0.023
0.028
S4FFNO (Ours)
0.008
0.005
0.003
0.030
Table 1: aRMSE values at different resolutions for Burgers' and KS with different viscosities. S4FFNO achieves up to 6x less error than its memoryless counterpart (FFNO) in KS at resolution 32. The final time of KS is 2.5 seconds and it contains 25 timesteps. The final times of Burgers' is 1.4 seconds and it contains 20 timesteps.
prediction at time t, S4FFNO has access timesteps. whereas Multi Input FFNO takes as imput the previous four timesteps. More details on training are given in Appendix F, and on the Burgers' experiment in Appendix C.
(compressed) memory all previous
6.1 KURAMOTO-SIVASHINKK EQUATION (ID): STUDY IN LOW-RESOLUTION
The Kuramoto Kuramoto-Sivashinse Sivashinsky equation (KS) is a nonlinear PDE that used as modeling tool in fluid dynamics, chemical reaction dynamics, and ion interactions. Due to its chaotic behavior it can model instabilities in various physical systems. For viscosity v R., is written as u, + War + Mpx + Mugarz O. We generated datasets for KS at different viscosities and resolutions, and show results Table 1. At resolutions 32 and 64, the memory models (S4FFNO and Multi Input FFNO) outperform the Markovian baselines. In particular, S4FFNO can achieve up to 6 less eтror than its Markovian counterpart (FFNO) and additionally 3x less error than Multi Input FFNO. Furthermore, in Appendix H.2 we show that S4FFNO still achieves 3x less error than FFNO models with 7x more parameters.
At resolution 128, FFNO has similar performance compared to S4FFNO, and it outperforms Multi Input FFNO. This is in agreement with other works which propose following the Markovian assump- tion in neural operators (Tran et al., 2023; Lippe al., 2023), where it is argued that incorporating previous timesteps as input is not necessary and can lead to difficulties in learning, as it seems to happen with Multi Input FFNO. By contrast, S4FFNO effectively models memory when it useful (resolutions 32 and 64) without compromising performance at higher resolutions.
In Figure we show the performance of all models across continuous range of resolutions. It can be seen that there is "cutoff" resolution at which memory models (i.e. S4FFNO) start outperforming Markovian ones (i.e. FFNO) by large margin. Very importantly, this cutoff resolution depends on the viscosity, being around 76 when = 0.075. 68 when и 0.1, and 52 when v 0.125. In the KS equation, a lower viscosity leads to the appearance of higher frequencies in the Fourier spectrum (see second row of Figure 1), which are not well captured at low resolutions. Thus, identify the resolution relative to the Fourier frequency spectrum of the solution as a key factor for the improved performance of MemNO over memoryless neural operators. We note that even if the initial condition does not contain high frequencies, in the KS equation high frequencies will appear the system evolves. We provide similar study ID Burgers equation in Appendix C.
Lastly, there are several architecture choices to model memory that improve performance in low resolution. In particular, in Appendix G.I show that using L.STM instead of S4 also brings similar
8

---

Published as a conference paper at ICLR 2025
performance improvements. Likewise, in Appendix G.3 we show that using S4 as the memory model with U-Net as the Markovian neural operator also outperforms the purely Markovian U-Net.
1.0
0.075 Resolation
jim 0.1.1 Seselution
KS, 0.125, flcsouution
Federw 1221
Pottrrn
08
08
PRNO S4P110 10s
1ourl
$4P150 30
0..-
02
(a) 0.075
(b) 0.1
(c) V/ 0.125
10'
10-
305-
Time Daceeaa
1o
50.
to
25
35"
1071
10"1
30
200
250
200
250
50
500
150
200
250
20
200
150
200
250
Progunncy
Proqurtyy
Prquerry
Figure 1: (First row) aRMSE for several models in the KS dataset different resolution, where each column different viscosity. The final time is T 2.5s and there are N₂ 25 timesteps. (Sccond row) visualization of the whole frequency spectrum at each of the 25 timesteps for single trajectory in the datasct. The spectrum is obtained with the ground truth solution at resolution 512.
6.2 NAVIER-STOKES EQUATION (2D): STUDY IN OBSERVATION NOISE
The Navier-Stokes equation describes the motion of a viscous fluid. Like in Li et al. (2021), we consider the incompressible form in the 2D unit torus, which is given by:
Ow(x,t) t) Øt
+ u(z,t) Vu(z, t) = v/ f(z),
E (0, t (0,T)
V ul, t) = 0,
x x (0, 1) tE [0, x (0. 1)²
w(x, 0) = wo(x))
Where w = V x u the vorticity, wa E L ((0 1)2; R)) is the initial vorticity, v E R, is the viscosity coefficient, and E ((0,1)2;133 is the forcing function. In general, the lower the viscosity, the more rapid the changes in the solution and the harder it is to solve it numerically or with neural operator. We investigate the effect of memory when adding i.i.d. Gaussian noise the inputs of our neural networks. The noise is sampled i.i.d. from a Gaussian distribution (N00 and then added to training and test inputs. During training, for each trajectory different noise (with the same σ) is sampled at each iteration of the optimization algorithm. The targets in training and testing represent our ground truth, and do not contain added noise. In Figure 2a, we show the results 2.048. S4FFNO-2D
for = 10 when adding noise levels from outperforms FFNO-2D across most noise levels, and the difference between the two is especially significant for noise levels beyond 0.128, where FFNO-2D is around 50% higher than S4FFNO-21 (note the logarithmic scale). For this viscosity, adding small levels of noise actually helps training, which was also observed in other settings in Tran al. (2023). Figure 2b shows the same experiment performed with = 10 10-5, Again, S4FFNO-2D outperforms FFNO-2D across most noise levels. FFNO-2D losses are similarly around 50% higher for noise levels above 0.032. In this viscosity, adding these levels of noise does not help performance.
0.0 (no noise)
6.3 RELATIONSHIP WITH FRACTION OF UNOBSERVED MODES
In this section, we provide a simple experiment to quantify the effect of the fraction of unobserved modes the performance of memory based models. Precisely, suppose E 1.2 T.2 RV) is the

---

Published as a conference paper at ICLR 2025
Neveer 10-1 nOUSE
Newer 10-3
('022
185
1200
200
(a) V = 10 T 16s, N. 32
(b) V = 10 = 3.28, N₁ 32
Figure 2: nRMSE of FFNO-2D and S4FFNO-2D trained on Navier-Stokes 2D with different noise standard deviations a added to training and test inputs. Two configurations viscosity v and final time are shown.
solution of a 1-dimensional PDE at a certain timestep, and 0n for n. E Z is its Fourier Transform. we observe it at resolution f, we can only estimate its top L modes.. Thus, we define wy as the ratio of unobserved modes resolution f:
Σ Linl=Lf) |anl2
wy Enezlam |nn
(14)
wy is an approximate indicator of the amount of information that is lost when the solu- tion of the PDE is observed at resolution f. In practice, wy be computed by ap- proximating the On with the discrete Fourier modes the solution the highest reso- lution available. We show that there is a positive correlation between wy and the dif- ference in nRMSE between FFNO and S4FFNO for the KS experiment in Figure 3, and also the for Burgers' experiments of Appendix C in Figure 5. This demonstrates the benefits of memory as a way to compensate for missing information in the observations.
[ol, AnR.E)) for KS u.1
7 CONCLUSION AND FUTURE WORK
40
70
We study the benefits of maintaining memory while mod- eling time dependent PDE systems. When we only ob- serve part of the initial conditions (for example, PDEs ob- served on low-resolution or with imput noise), the system longer Markovian, and the dynamics depend on memory term. Taking inspiration from the Mori-Zwanzig formalism, we introduce MemNO, an architecture that combines Fourier Neural Operators (FNO) to model the spatial dynamics of the PDE, and the $4 sequence model to incorporate memory of past states. Through our exper- iments on different ID and 2D PDEs, we show that the MemNO architecture outperforms the memoryless base- lines, particularly when the solution to the PDE has large components high frequecyy Fourier modes.
0.4
-00
61
00
0.0
0.1
0.2
0.3 0.4 0.5
Figure 3: Values of wy and the difference nRMSE between FFNO and S4FFNO for different resolutions in the KS experiment of Section 6.1 with v = 0.1. wy is aver- aged across all trajectories the dataset and across timesteps.
We present several avenues for future work. First, our experiments observation noise are limited the setting where the input noise is i..d. Further, extending the experiments and observing the effects memory in more real-world settings (for ex- ample, with non-i.i.d. noise or in the presence of aliasing) seems fertile ground for future work, and also necessary to ensure that the application of this method does not have unintended negative con- sequences when broadly applied in society. Lastly, while we primarily compare between Markovian and memory architectures, a study on the trade-offs between different memory architectures such as S4FFNO, LSTM- FFNO, S4U-Net and Multi Input FFNO is an interesting direction for future work.
*This consequence of the Syquist-Shannon sampling theorem.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGEMENTS
RBR supported by the "la Caixa" Foundation (ID 100010434). The fellowship code is LCF/BQ/BU22/1 1930090. TM is supported in part by CMU Software Engineering Institute via Department of Defense under contract FA8802-15-D-0002 AR is supported in part by NSF awards IIS-2211907, CCF-2238523, and Amazon Research.
REFERENCES
Shaojie Bai, Zico Kolter, and Vladlen Koltun. Deep equilibrium models. Advances in neural information processing systems, 32, 2019.
Andrea Beck, David Flad, and Claus-Dieter Munz. Deep neural networks for data-driven les clo- sure models. Journal of Computational Physics, 398:108910, 2019. ISSN 0021-9991. doi: http::/doi.orr//10 101 1001 sct.annee/atticeee//////00999999
URL https://wwet http://www.cieeceerreeeett. com/
Maximilian Beck, Korbinian Pöppel, Markus Spanring, Andreas Auer, Oleksandra Prudnikova, Michael Kopp, Günter Klambauer, Johannes Brandstetter, and Sepp Hochreiter. xlstm: Extended long short-term memory, 2024.
Kaushik Bhattacharya, Bamdad Hosseini, Nikola Kovachki, and Andrew M Stuart. Model reduc- tion and neural networks for parametric pdes. The SMAI journal of computational mathematics, 7:121-157, 2021.
Johannes Brandstetter, Max Welling, and Daniel E Worrall. Lie point symmetry data augmentation for neural pde solvers, 2022.
Heinz-Peter Breuer and Francesco Petruccione. The Theory of Open Quantum Systems. University Press, 2002.
Oxford
Shuhao Cao. Choose a transformer: Fourier or Galerkin. Advances in Neural Information Process- ing Systems (NeurIPS 2021), 34, 2021. URL https:
https:://penzews http://////p/peeeeeeeeeee net forum?id
ssohLemn4-r.
Tianping Chen and Hong Chen. Universal approximation to nonlinear operators by neural networks with arbitrary activation functions and its application to dynamical systems. IEEE transactions neural networks, 6(4):911-917 1995.
Demetrios Christodoulou. The formation of shocks in 3-dimensional fluids, volume 2. European Mathematical Society, 2007.
Cooley, Lewis, and P. Welch. The finite fourier transform. IEEE Transactions on Audio and Electroacoustics, 17(2):77-85, 1969. doi: 10.1
Tri Dao and Albert Gu. Transformers are SSMs: Generalized models and efficient algorithms through structured state space duality, In International Conference on Machine Learning (ICML), 2024.
Justin Gilmer, Samuel Schoenholz, Patrick Riley, Oriol Vinyals, and George Dahl. Neural message passing for quantum chemistry. In International conference on machine learning, pp- 1263-1272. PMLR, 2017.
A. Gouasmi, E.J. Parish, and K. Duraisamy. A priori estimation of memory effects in reduced- order models of nonlinear systems using the mori-zwanzig formalism. Proceedings of the Royal Society A: Mathematical, Physical and Engineering Sciences, 473(2205) 2017. doi: 10. 10.1098/rspa.20177
URL http://////// dxx doi wrg/o0 10. 1098 rspa. 2017. 0385.
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint ar21v:23/2.00752, 2023.
Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. In The International Conference Learning Representations (ICLR), 2022.

---

Published as a conference paper at ICLR 2025
Albert Gu, Isys Johnson, Aman Timalsina, Atri Rudra, and Christopher Re. How to train your HIPPO: State space models with generalized orthogonal basis projections. In International Con- ference Learning Representations, 2023.
A. Gupta and P. F. Lermusiaux. Neural closure models for dynamical systems. Proceedings of the Royal Society A: Mathematical, Physical and Engineering Sciences, 477(2252), Aug 2021.
Jayesh K Gupta and Johannes Brandstetter. Towards multi mualii-ppaiioempora-sscaaa generalized PDE modeling. Transactions on Machine Learning Research, 2023. ISSN 2835-8856.
E. Hairer and G. Wanner. Solving Ordinary Differential Equations II: Stiff and Differential-Algebred Problems. Berinner-Verlagee 1996. ISBN 3-540-60452-9.
Zhongkai Hao, Zhengyi Wang, Hang Su, Chengyang Ying, Yinpeng Dong, Songming Liu, Ze Cheng, Jian Song, and Jun Zhu. Gnot: A general neural operator transformer for operator learning. In International Conference on Machine Learning, pp. 2556-12569, PMLR, 2023.
Zhongkai Hao, Chang Su, Songming Liu, Julius Berner, Chengyang Ying, Hang Su, Anima Anand- kumar, Jian Song, and Jun Zhu. DPOT: Auto-regressive denoising operator transformer for large- scale PDE pre-training. March 2024. URL https: /tt:://gittt com/ tta-mm/DPOO..
Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (GELUs). June 2016. URL http:
//arxiv. owte:iiorrg/aaao abs www.aa/......eeeee//
Sepp Hochreiter and Jürgen Schmidhuber: Long Short-Term Memory. Neural Computation, doi: 1.211666/coo.1997..00 1735. URL https: ////oi.org//...1/ http://11.112//cncecee 1997. 8.375.
(8):1735-1780, 11 1997. ISSN 0899-7667.
Aaron Kelly, Andrés Montoya-Castillo, Lu Wang, and Thomas E. Markland. Generalized quantum master equations in and out of equilibrium: When can one win? The Journal of Chemical Physics, 144(18):184105, 05 2016. ISSN 0021-9606. doi: 10. 4948612. URL https:
//doi. org.o/gg////o/ 10. 1tt....013/...//...0. 4948612.
Thomas N Kipf and Max Welling. Semi-supervised classification with graph convolutional net- works. arXiv preprint arXiv: fttp:66000207 2016.
Nikola Kovachki, Zongyi Li, Burigede Liu, Kamyar Azizzadenesheli, Kaushik Bhattacharya, An- drew Stuart, and Anima Anandkumar. Neural operator: Learning maps between function spaces with applications to pdes. Journal of Machine Learning Research, 24(89):1-9) 1-77 2023.
Remi Lam, Alvaro Sanchez-Gonzaler Matthew Willson, Peter Wirnsberger, Meire Fortunato, Fer- ran Alet, Suman Ravuri, Timo Ewalds, Zach Eaton-Rosen, Weihua Hu, et al. Graphcast: Learning skillful medium-range global weather forecasting. arXiv preprint arXiv:2212.. 12794, 2022.
Randall LeVeque. Finite difference methods for ordinary and partial differential equations: steady- state and time-dependent problems. SIAM, 2007.
Zijie Li, Kazem Meidani, and Amir Barati Farimani. Transformer for partial differential equations'
operator learning, Transactions on Machine Learning Research, 2023a. ISSN 2835-8856. URL https: http///ppeeeeeeeeeeeeeeeeeeeeeee net/
Zijie Li, Dule Shu, and Amir Barati Farimani. Scalable transformer for pde surrogate modeling. 2023b.
Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Andrew Stuart, Kaushik Bhat- tacharya, and Anima Anandkumar. Multipole graph neural operator for parametric partial differ- ential equations. Advances in Neural Information Processing Systems, 33:6755-6766 2020.
Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, An- drew Stuart, and Anima Anandkumar. Fourier neural operator for parametric partial differential equations. In Proceedings of the International Conference on Learning Representations (ICLR). ICLR, 2021.
12

---

Published as a conference paper at ICLR 2025
Zongyi Li, Nikola Kovachki, Chris Choy, Boyi Li, Jean Kossaifi, Shourya Otta, Mohammad Amin Nabian, Maximilian Stadler, Christian Hundt, Kamyar Azizzadenesheli, et al. Geometry- informed neural operator for large-scale 3d pdes. Advances in Neural Information Processing Systems, 36, 2024a.
Zongyi Li, Hongkai Zheng, Nikola Kovachki, David Jin, Haoxuan Chen, Burigede Liu, Kamyar Azizzadenesheli, and Anima Anandkumar. Physics-inforee informed neural operator for learning partial differential equations. 1(3), 2024b. doi: 10.1145/3648506. URL https: //oi: org/ 10.
1145/364850
Phillip Lippe, Bastiaan S. Veeling, Paris Perdikaris, Richard E Turner, and Johannes Brandstetter. PDE-refiner: Achieving accurate long rollouts with neural PDE solvers. In Thirty-seventh Con- ference Neural Information Processing Systems, 2023.
Ilya Loshchilov and Frank Hutter. SGDR: Stochastic gradient descent with warm restarts. In Inter- national Conference on Learning Representations, 2017.
Lu Lu, Pengzhan Jin, and George Em Karniadakis. Deeponet: Learning nonlinear operators for iden- tifying differential equations based on the universal approximation theorem of operators. arXiv preprint 1XX0v:99000031900 2019.
Chao Ma, Jianchun Wang, and E Weinan. Model reduction with memory and the machine learning of dynamical systems. August 2018. URL
11008.8011000 .0..00222 04258.
Tanya Marwah, Ashwini Pokle, Zico Kolter, Zachary Lipton, Jianfeng Lu, and Andrej Risteski. Deep equilibrium based neural operators for steady-state pdes. Advances in Neural Information Processing Systems, 01115516-157737 2023.
Michael McCabe, Bruno Régaldo-Saint Blancard, Liam Holden Parker, Ruben Ohana, Miles Cran- mer, Alberto Bietti, Michael Eickenberg, Siavash Golkar, Geraud Krawezik, Francois Lanusse, Mariel Pettee, Tiberiu Tesileanu, Kyunghyun Cho, and Shirley Ho. Multiple physies pretrain- ing for physical surrogate models. October 2023. URL /ttt:://arxii org abs/2310. 02994.
Nick McGreivy and Ammar Hakim. Weak baselines and reporting biases lead to overoptimism in machine learning for fluid-related partial differential equations. Nature Machine Intelli- URL http::://ao..or/11.... s42256- 4 0222566-22-08877--0 Published on September 25, 2024.
gence, 2024. doi: 10. 000554442256-2--0887777
Andrés Montoya-Castillo and David Reichman. Approximate but accurate quantum dynamics from the Mori formalism: I. Nonequilibrium dynamics. The Journal of Chemical Physics, 144 (18):184104, 05 2016. ISSN 0021-9606. doi: 10.1063/1.4444408. URL https: / //doi. org/ 10.1063/11 11.100633..444 4.4944408
Hazime Mori. Transport, collective motion, and brownian motion. Progress of theoretical physics, 33(3):423 455, 1965.
Sadao Nakajima. On Quantum Theory of Transport Phenomena: Steady Diffusion. Progress of Theoretical Physics, (066))44-9955 12 1958. ISSN 0033-068X. doi: 10. 10.1144// URL https: http:///ooooog//////// org/10 11011114/PPP
20.948.
Silvia Noschese, Lionello Pusquini, and Lothar Reichel. Tridiagonal toeplitz matrices: properties and novel applications. Numerical linear algebra with applications, 20(2):302-326, 2013.
Antonio Orvieto, Samuel Smith, Albert Gu, Anushan Fernando, Caglar Gulcehre, Razvan Pas- canu, and Soham De. Resurrecting recurrent neural networks for long sequences. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scar- lett (eds.), Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 26670-26698 PMLR, 23-29 Jul 2023. URL ht https:/ hpps://proceedeeee /ttpp//rceeenngneee/e mlr Vitess/v202//rve/
Shaowu Pan and Karthik Duraisamy. Data-driven discovery of closure models. SIAM Journal on Applied Dynamical Systems, 17(4):2381-244 2018. doi: 10.1137/181117263 URL https:
/odo. /t/ wotooioorg/1... 1137/1 8M1 177263.
13

---

Published as a conference paper at ICLR 2025
Eric Parish and Karthik Duraisamy. Non-markovin closure models for large eddy simulations using the morizzwanzig formalism. Physical Review Fluids, 2:014604, 01 2017. doi: 10.1103/ PhyyseevFluids22014666
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Köpf. Ed- ward Yang. Zach Vito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. PyTorch: an imperative style, high-pprformance deep learning library. Curran Associates Inc., Red Hook, NY, USA, 2019.
Jaideep Pathak, Shashank Subramanian, Peter Harrington, Sanjeev Raja, Ashesh Chattopadhyay, Morteza Mardani, Thorsten Kurth, David Hall, Zongyi Li, Kamyar Azizzadenesheli, et al. Four- castnet: A global data-driven high-resolution weather model using adaptive fourier neural opera- tors. arXiv preprint arXiv:2202. 14, 2022.
Bo Peng, Eric Alcaide, Quentin Anthony, Alon Albalak, Samuel Arcadinho, Stella Biderman, Huanqi Cao, Xin Cheng. Michael Chung, Matteo Grella, Kranthi Kiran Gv, Xuzheng He, Haowen Hou, Jiaju Lin, Przemyslaw Kazienko, Jan Kocon, Jiaming Kong, Bartlomiej Koptyra, Hayden Lau, Krishna Sri Ipsit Mantri, Ferdinand Mom, Atsushi Saito, Guangyu Song, Xiangru Tang, Bolun Wang, Johan S Wind, Stanislaw Wozniak, Ruichong Zhang, Zhenyuan Zhang, Qihang Zhao, Peng Zhou, Qinghua Zhou, Jian Zhu, and Rui-Jie Zhu. RWKV: Reinventing RNNs for the transformer era. May 2023. URL http: www.i.xxivvvv.....a..aaaaa wttp://..rraabb///3/000 13048.
Stephen B Pope. Turbulent flows. Measurement Science and Technology, 12(1 22(11):220-202) 2001.
Md Ashiqur Rahman, Zachary Ross, and Kamyar Azizzadenesheli. U-NO: U-shaped neural operators. Transactions Machine Learning Research, 2023. ISSN 2835-8856.
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomed- ical image segmentation. In Medical image computing and computer-assisted intervention- MICCAI 2015: 18th international conference, Munich, Germany, October 5-9, 2015, proceed- ings, part III 18, pp. 234-241. Springer, 2015.
W. E. Schiesser. The Numerical Method of Lines. Academic Press, 1991. ISBN 1--2224330--.
Claude Shannon. Communication in the presence of noise. Proceedings of the Institute of Radio Engineers, 37(1):10-21, 1949. doi: 10.1 109/jrproc. 1949.2329909 Reprinted as a classic paper in Proc. IEEE, Vol. 86, No. 2, Feb 1998, Archived 2010-02-08 the Wayback Machine.
Junhong Shen, Tanya Marwah, and Ameet Talwalkar. Ups: Towards foundation models for pde solving via cross-modal adaptation, 2024.
Qiang Shi and Eitan Geva. A new approach to calculating the memory kernel of the generalized quantum master equation for an arbitrary system- bath coupling. The Journal of chemical physies, 119023):12063-1100 2003.
Makoto Takamoto, Timothy Praditia, Raphael Leiteritz, Dan MacKinlay, Francesco Alesiani, Dirk Pflüger, and Mathias Niepert. PDEBENCH: An extensive benchmark for scientific machine learnin. In ICLR 2023 Workshop on Physics for Machine Learning, 2023.
Yi Tay, Mostafa Dehghani, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang, Sebastian Ruder, and Donald Metzler. Long range arena: benchmark for efficient transformers. arXiv preprint arXiv:2011.0400 1.00066, 2020.
Alasdair Tran, Alexander Mathews, Lexing Xie, and Cheng Soon Ong. Factorized fourier neural operators. In The Eleventh International Conference on Learning Representations, 2023.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural informa- fion processing systems, 30, 2017.
14

---

Published as a conference paper at ICLR 2025
Pauli Virtanen, Ralf Gommers, Travis E. Oliphant, Matt Haberland, Tyler Reddy, David Courna- peau, Evgeni Burovski, Pearu Peterson, Warren Weckesser, Jonathan Bright, Stéfan J. van der Walt, Matthew Brett, Joshua Wilson, K. Jarrod Millman, Nikolay Mayorov, Andrew R. J. Nel- son, Eric Jones, Robert Kern, Hric Larson, C J Carey, Ilhan Polat, Yu Feng, Eric W. Moore, Jake VanderPlas, Denis Laxalde, Josef Perktold, Robert Cimrman, Ian Henriksen, E. A. Quintero, Charles R. Harris, Anne M. Archibald, António H. Ribeiro, Fabian Pedregosa, Paul van Mul- bregt, and SciPy 1.0 Contributors. SciPy 1.0: Fundamental Algorithms for Scientific Computing in Python. Nature Methods, 17:261-211 2020. doi: 10.1038/s4 599-019-0086--111
Ming-Liang Zhang, Being Ka, and Eitan Geva. Nonequilibrium quantum dynamics in the con- densed phase via the generalized quantum master equation. The Journal of chemical physics, 125 (4), 2006.
Robert Zwanzig. Memory effects in irreversible thermodynamics. Physical Review, 124(4):983, 1961.
Robert Zwanzig. Noneguilibrium Statistical Mechanics. Oxford University Press, New York, 2001. ISBN 9780195140187.
15

---

Published as a conference paper at ICLR 2025
A ADDITIONAL RELATED WORK
Neural Operators. The Fourier Neural Operator (FNO) neural operator that performs trans- formation in the frequency space of the input (Li et al., 2021). Other models have proposed differ- ent inductive biases for neural operators, including physics based losses and constraints (Li et al., 2024b), using Deep Equilibrium Model (DEQ) (Bai et al., 2019) design specialized architectures for steady-state (time-independen)
PDEs (Marwah et al., 2023), and using local message passing Graph Neural Networks (GNNs) (Gilmer et al., 2017; Kipf & Welling, 2016) encoders model ir- regular geometries (Li et al., 2020; 2024a). Other methodologies to solve PDEs include methods like (Gupta & Brandstetter, 2023; Rahman et al., 2023) that the U-Net (Ronneberger et al., 2015) architectures and works like Cao (2021); Hao et al. (2023) that introduce different Transformer- based (Vaswani al., 2017) neural solution operators for modeling both time-dependent and time- independent PDEs. While most of these methodologies are designed for dependetn PDEs, there is no clear consensus of how to use the past states to predict future states, and most of these methods predict the PDE states over time in an autoregressive way by conditioning the model on varying lengths the past states (Li et al., 2021; Tran et al., 2023; Hao et al., 2023).
Foundation models. There have been community efforts towards creating large-scale foundational models for modeling diverse PDE families (McCabe et al., 2023; Hao et al., 2024; Shen et al., 2024), and weather prediction (Pathak al., 2022; Lam et al., 2022).
B NETWORK ARCHITECTURES
For all our models, we use a simple spatial positional encoding E. In I-D, if the grid has f eq- uispaced points in 10, L), then E E R and the positional encoding is defined as E; = T for 0 i < f In 2-D, if we have f x points in a 2-D equispaced grid in [0, Lz] x [0, Lyl. the positional encoding is defined as Ei = (z,, Z,,) The input lifting operator (i.e. encoder) rin is linear layer that maps a concatenation of input and grid to the hidden dimension R² + RA, which is applied to each element of the spatial dimension independently. It is shared across all model architectures. Likewise, for the decoder Raus, we use another linear layer Rd - R.
Factorized Fourier Neural Operator (FFNO) (Tran al., 2023): This model is refinement over the original Fourier Neural Operator (Li et al., 2021). Given a hidden dimension h and a spatial grid S, its layers& &: R|
+ R|S) x are defined as:
{(v) == v + Linear
Linear, ,A K[[u]
(15)
where a is the GeLU activation function (Hendrycks & Gimpel, 2016) and h' is an expanded hidden dimension. K. is a kernel integral operator that performs a linear transformation in the frequency space. Denoting by FFT,, IFFT, the Discrete Fast Fourier Transform and the Discrete Inverse Fast Fourier Transform along dimension a (Cooley et al., 1969) respectively, the operator can be written as:
Klu| =
Σ
IFFT[R, FFT.[u]]
(16)
(I......)))
for learnable matrices of weights Ra E Ch2 Xhus, kmax is the maximum number of Fourier modes which are used in IC. We use all Fourier modes by setting kmax = =111-
In our experiments, The FFNO model consists of 4 FFNO layers. For experiments in ID, the hidden dimensions are all 128 (h2 128 for j 0, 1, 2, 3) and the expanded hidden dimension of FFNO's MLP A is 4. 128. For experiments in 2D, the hidden dimensions are all 64 and the expanded hidden dimension is 4. 64.
$4 - Factorized Fourier Neural Operator (S4FFNO): This model uses our MemNO framework. To isolate the effect of memory, all layers except the memory layer the same as FFNO. For the memory layer, we choose an S4 layer (Gu et al., 2022) with a state dimension of 64 and a diagonal S4 (S4D) kernel.5
'The S4 repository has two available kernels, the diagonal $4 (S4D) and the Normal Plus Low Rank S4 (S4NPIR). In our experiments, we didn't find significant difference between the two, and chose S4D for simplicity.
16

---

Published as a conference paper at ICLR 2025
Multi Input Factorized Fourier Neural Operator (Multi Input FFNO): This architecture uses the solution at the last K = 4 timesteps as input to predict the next timestep, as originally proposed by Li et al. (2021). Thus, this model uses the (uncompressed) memory of the four previous timesteps and it is not Markovian. We choose K = 4 because, in practice, the number of previous timesteps to which we have access is limited, if any. We also believe that Multi Input FFNO is advantaged by having access to four ground truth observations, whereas the rest of the models only have access to one. Thus, we consider K fairness in comparisons. On the implementation side, the only difference with the FFNO architecture resides in the input lifting operator Rin. which takes concatenation of Ut Vt Ut.-1 Ut as input to predict Vt.4.- In all our experiments, we choose the fourth timestep of the solution of the PDEs as initial condition for the rest of the models, whereas Multi Input FFNO is given access to the first, second, third, and fourth timesteps for its first prediction. The number of layers and hidden dimensions are the same as FFNO.
be reasonable choice when considering practical applicability and
Factformer 1D (Li et al., 2023b): This models uses four linear attention layers over the spatial sequence length and an MLP as output projection. We set the hidden dimension to 64, and each attention layer has 4 heads with hidden dimension of 1 128, thus expanding the dimension from 64 to 512. The implementation is taken from https:/ /ttp///tttttt
com/ com/Barat BBarat iLlb/ / /ubb/Factt FactF PoPPr ormer yet making slight modification for 1D instead of 2D inputs. Li et al. (2023b) deals with 2D inputs in the following manner: given a hidden state of a solution w with with spatial dimensions S, and S,, and hidden dimension H, two queries and keys are built from we by applying two different MLPs (MLP, x and MLP,) y and then taking the mean across Sy and Sz, respectively. Thus, we get and kz of shape (S,, H,, and qy and ky of shape (s,, .). The attention "values" of shape (S., Syy H) are obtained from w by a linear layer. Then two linear attention transformations are applied, first with da and h across the S5 dimension, and then qy and ky across the Sy dimension. For our ID case we do not have MLP,, y. nor qy, hig Concretely, we only have one MLP,. we do not take means to compute 4, and kn. and we only apply one linear attention per layer.
Galerkin Transformer (GKT) (Cao, 2021): This model uses four linear attention layers over the spatial sequence length. It includes positional information by concatenating the grid coordinates into the queries, keys and values After the attention layers, two FNO layers (using all Fourier modes) are used. The hidden dimension used in the experiments is 32 (both for the transformer encoders and the spectral regressor). For the experiments of Figure 1, GKT had unstable performance for some resolutions. Thus, for some resolutions we tried a different training setup: a dropout of 0.05 in the linear attention layer and 0.025 in the FFN layer and 50 training epochs instead of 200. We reported the nRMSE of the best configuration. Specifically, the dropout + reduced training epochs helped performance in resolutions [40-64] for / 0.075, [8-60] for
0.1 and [36-48] for
= 0.125 (all inclusive intervals). The implementation is based on the publicly available code https: /ttp:///ttttt thub. com/ sca.matt wwwphhhhwwalaannnaeeee. trann former.
U-Net Neural Operator (U-Net) (Gupta & Brandstetter, 2023): This model consists of four down- sample convolution blocks, middle convolution block, and four upsample convolution blocks. The upsample blocks have residual connections to the downsample blocks in the typical U-Net fashion. The downsample blocks have channel multipliers [1, 2, 2, 2] and no time embeddings are used. The first hidden dimension is 32. The implementation is based on the repository https: / //tptggttttt.e///// com/ httpp/paaaaaaaaarreeeee
S4- - U-Net Neural Operator (S4U-Net): This model also uses our MemNO framework. As before, all layers except the memory layer are the same as U-Net. The state dimension is 16 and the we used the S4D kernel. We apply the memory layer after the "middle" convolution block.
B. PARAMETER AND TRAINING TIMES FOR DIFFERENT ARCHITECTURES
The number of parameters of the different baselines and training times (forward backward) is shown in Table 2.
B.2
ALGORITTMIC COMPLEXITIES OF S4FFNO AND FENO
We present the theoretical complexities of the cores of the S4FFNO and FENO layers (i.c., the spectral convolution of FFNO and the convolution of S4FFNO). Let S be the spatial resolution, T
17

---

Published as a conference paper at ICLR 2025
Architecture
# Params (millions)
Training time (miliseconds)
Factformer (ID)
0.65
102
GKT
0.29
21
U-Net
2.68
FFNO
4.89
28
Multi Input FENO
4.89
28
S4FFNO
4.94
32
S4U-Net
2.82
25
Table Number of parameter and training times (forward and backward pass) of architectures for the experiments in Section 6.1 and Appendix G.3. The batch size is 32, the spatial resolution is 64 and the number of timesteps is 25. The GPU is an NVIDIA L40S.
the number of timesteps, H the hidden dimension and N the state dimension of S4. The core spectral convolution of FFNO (Eq. 16) has a Discrete Fourier Transform across the space dimension and a matrix multiplication the frequency space, which have complexities H5TH5) and O(TS respectively (tildes denote log factors). In contrast, the S4 layer has Discrete Fourier Transform across the time dimension and requires building the convolution kernel, which have complexities O(SHT) and O(SH (Ñ + +)) respectively (Gu et al., 2022). In our cases, ranges from 32 to 128, and T is either 20, 25 or 32. Thus, in most cases HTT < OS)..). As for the other term, we use H 128 and N = 64, so N +T < H and we also have OSS + +))) < O(TSH²). Thus, the S4 memory layer requires less computation than a spatial FFNO layer.
C BURGERS' EQUATION (1D): A STUDY ON LOW-RESOLUTION
The Burgers' equation with viscosity v E IR nonlinear PDE used as modeling tool in fluid mechanics, traffic flow, and shock waves analysis. It encapsulates both diffusion and advection processes, making it essential for studying wave propagation and other dynamic phenomena. is known for exhibiting rich variety of behaviors, including the formation of shock waves and the transition from laminar to turbulent flow. The viscous Burgers' equation is written as:
U1
WUr Vilga
We used the publicly available dataset of the Burgers' equation in the PDEBench repository (Takamoto al., 2023) with viscosity 0.001, which is available at resolution 1024.
We perform experiments at resolutions 64, 128, 256, 512 and 1024 and show results for the models Galerkin Transformer (GKT) (Cao, 2021), U-Net neural operator (U-Net) (Gupta & Brandstetter, 2023), Factorized Fourier Neural Operator (FFNO), Multi Input Factorized Fourier Neural Operator (Multi input FFNO) and our proposed model S4 Factorized Fourier Neural Operator (S4FFFNO). The results are shown in Figure 4a.
In low resolutions, memory- based architectures (Multi Input FFNO and S4FFNO) outperform the best Markovian baseline (FFNO). Specifically, S4FFNO achieves more than 4x less error than FFNO resolutions 32 and 64 (see Table 1). Additionally, S4FFNO has slightly better performance than Multi Input FFNO in high resolutions (512, 1024). Furthermore, we show the difference in nRMSE between FFNO and S4FFNO at each timestep in figure 4b. We observe that at the first timestep there is difference between the two models wwhchh is expected because S4FFNO has the exact same architecture as FFNO for the first timestep. Yet as the initial condition is rolled out, there is more history of the trajectory and the difference between FFNO and S4FFNO increases.
C.1 CORRELATION WITH FRACTION UNOBSERVED MODES
As mentioned in Section 6.3 we measure the correlation wy defined in Equation 14 with the difference in the nRMSE between FFNO and S4FFNO. The results can be seen in Figure 5.
18

---

Published as a conference paper at ICLR 2025
Resolution
Burgers' 0.001. HRMSE Reselution
32
128
FFNO
0.3
Mu Inpuc FFNG
S4FAND (Ours!
A ARRSSF
0.0
12
120
256
512
3024
(b) Difference between the aRMSE of FFNO and S4FFNO per timestep (higher difference means better performance of S4FFNO).
Resolution
(a) nRMSE of several models at different resolutions.
Figure 4: Results for the Burgers's PDEBench dataset with viscosity 0.001.
AnRMEE) Burgersi 0.001
1000
0.12
800
0.10
AnOWSE
0.08
0.08
0.04
Figure 5: Difference in nRMSE between FFNO and S4FFNO against wy (defined in Equation 14) for different resolutions of the Burgers' Equation. wy is averaged over all trajectories in the dataset and across all timesteps of the experiment. The value computed approximating the continuous Fourier modes with the Discrete Fourier modes of the solution in the highest resolution available (1024 for Burgers' Equation).
D ANALYSIS OF HIGH-FREQUENCY FOURIER MODES IN COMMON 1-D DATASETS
In Section 5.1, we explained that one of the main criteria for choosing the datasets of our experi- ments was the high contribution from high-frequency Fourier modes in the solutions of the PDEs. Intuitively, when the solution contains contributions from high-frequency Fourier modes, say higher than number k. then it cannot be approximated accurately from its first k Fourier components (see Definition 5). Therefore, when observed finite resolution f. only \ Fourier modes can be estimated which is not enough to approximate the solution when >>>>>>>>>> [ J. In this case, there an "unobserved" part of the solution (which corresponds to the fighffrequency components), and thus we can expect the memory term of the Mori- Mori-Zwanzig Equation to be non-negligible.
In order to quantitatively measure the importance of the high-frequency components of function, we propose using w/ from Equation 14. This quantity measures the fraction of Fourier modes (weighted by their amplitude) that are above the frequency | J. and thus 0 < W < 1. When wy is close then we expect the solution be very accurately approximated from the Fourier modes that are observed at resolution f. so the "unobserved" part of the function is very small and thus the
Thii consequence of Nyquist-Shannon Theorem (Shannon, 1949)
19

---

Published as a conference paper at ICLR 2025
memory term of the Mori-Zwanzig Equation is expected to be negligible. Conversely, if wy is large, we expect the memory term to be significant".
The results for wy are shown in Figure For most PDEBench datasets the values of wy are very small, even for very small resolutions like 16 (note that the original data is in resolution 1024). Therefore, based on our previous discussion we expect the memory term to be negligible. On several exploratory experiments
these datasets, indeed saw no benefit of using memory to model PDEs. The only exception is the Burgers' dataset with viscosity v = 0.001. where our experiments in Appendix C show superior performance of memory cemoory-augmeneed models over Markovian ones for resolutions 64, 128 and 256 (Figure 4a).
In the case of the Kuramoto Sivashiiiii (KS) dataset, again that the viscosities that are typical in other works, like V = 1.0 and y = 0.5 in PDE-Refiner (Lippe et al., 2023), do not have a large wf, unless the resolutions are low (16 or 32). For that reason, we generated our own datasets with lower viscosities, which yield higher values for w and thus a more challenging benchmark to compare Markovian and memory models. Besides the change of viscosity. PDE-Refiner generation method had warm-up of T = 72 seconds, while we did not consider warm-up. This warm-up explains the higher presence of high frequencies for viscosity v 0.5 compared to our viscosities at resolution 16. Details and code to generate our datasets are provided in Appendix E.
can be seen that wy depends on both the parameters of the PDE (i.e. viscosities in Burgers' and KS) and the observation resolution f. Thus, key to understanding the importance of the memory term is observing the rexolution relative to the Fourier frequency spectrum of the solution, as we noted in Section 6.1. Additionally, another important characteristic that affects wy is the frequency spectrum of the initial condition. While the initial condition for KS is generated as a superposition of sinusoidal waves, PDEBench also uses this superposition of waves but applies some transformations to it, like taking the absolute value (see Appendix D of Takamoto et al. (2023))- These transformations lead to the appearance of higher order frequencies in the initial conditional and thus also affect the frequency spectrum of later timesteps. We believe this is why Burgers' with 0.001 also exhibits high wy at resolutions 128 and 256 (Figure 6).
We hope wy can serve as a practical quantity to help practitioners and researchers explore whether to consider memory architectures or not.
E DATA GENERATION
E. KURAMOTO-SIVASHINSS EQUATION
The Kuramoto-Sivashinst (KS) equation is given by:
+ uur + Uga + Vulrrz = 0 (t, E [0, x L] u(0,z) = uo())) x E [0, L]
We use periodic boundary conditions. Our data generation method is very similar to the one used in PDERefiner (Lippe et al., 2023), except for the three following differences: (1) We do not have random At per trajectory (2) We set the initial condition to have eight Fourier modes in the spectrum, whereas Lippe et al. (2023) uses three (3) We do not discard the first generated timesteps of the solution of the PDE. We provide a forked repository with these changes in https: / //gitt tut..... com/r-buittne wwwbouttrrtrtieoo LPSDAA, which is based on the original repository of Brandstetter et al. (2022) https: hwt.:/ttht......b.oeooeeeeeeaeee www..///wwwonttttttooooooo www.dnddteeetetttteeeeeeeeeeeeee johannes //ss/LL /PSSDA The generation command for our datasets (change -vicoosity for the desired value):
pythen waweeraateggeeeatedddta.aa -experiment KS -ttrainssmmlees-20000
-vvalic.samples=256 -tesstssmples=00 --L=64 --nt=51 --nx=512
--nt.effective-5) --viscosity--.1 - 0 --1max-8
Now, we give an explanation of the generation proceue, We employ the method of lines (Schiesser, 1991), where the spatial dimension is discretized, and the PDE is transformed to system of Or-
'This is not a precise mathematical argument, but rather an intuition that has proven to be helpful in practice for the PDEs we have considered. In general, the interaction of the PDE and the frequency spectrum of the solution complex and wy by itself is enough determine the magnitude of the memory term.
20

---

Published as a conference paper at ICLR 2025
for different datasets and resolutions
Advection (PDEBench)
0.051
0.015
0.005
0.002
0.001
0.7
Reaction-Diffun Diftton-iiffusiiin (PDEBench)
Burgers' u 1.0 (PDEBench)
0.004
0.001
0.001
0.000
0.000
0.6
Burgers' = 0.1 (PDEBench)
0.009
0.001
0.000
0.000
0.000
0.3
0.01 (PDEBench)
0.148
0,056
0,013
0.002
0.000
Detaset
Burgers' (PDEBench)
0.4
KS v 1.0 (PDE-Refine))
0.553
0.076
0.001
0.000
0.000
KS V 0.5 (PDE-Refiner)
0.788
0.249
0.010
0.000
0.000
0.483
0.341
0.071
0.002
0.000
0.1 (Ours)
0.575
0.450
0.145
0.006
0.000
KS vm 0.075 (Ours)
0.558
0.454
0.203
0.015
0.000
0.0
64
128
256
Figure 6: wy for different resolutions f and datasets. wy measures the ratio of Fourier modes that are above frequency & (see Equation 14). The Advection, Diffusion-Reaction and Burgers' datasets come from PDEBench (Takamoto et al., 2023) (the Diffusion-SSrption dataset is not considered because does not have periodic boundary conditions). The KS datasets come from either PDE- Refiner (Lippe et al., 2023), they generated by ourselves following Section E. The PDE- Refiner datasets use a ua = 72s, that is, they discard all timesteps the numerical solvers up time 72s. In contrast, our generated KS datasets do have warm-up period. wy is averaged across all trajectories in the dataset, and also averaged across the first 20 timesteps. The values of wy are computed approximating the continuous Fourier modes with Discrete Fourier modes of the solution in the highest resolution available (512 for KS datasets and 1024 for all other PDEBench datasets).
21

---

Published as a conference paper at ICLR 2025
dinary Differential Equations (ODEs), one per point in the grid. In order compute the spatial derivative of the solution at each point in the grid, pseudospectral method is used, where deriva- tives are computed in frequency space and then converted to the original space through Fast Fourier Transform. This method is implemented in the di fiff method of the sci scipy. py. fftpack package (Virtanen al., 2020). Similarly, the system of ODEs is solved numerically with implicit Runge- Kutta method of the Radau IIA family of order (Hairer & Wanner, 1996), which is implemented in the solve holve_d deviivv ivp method of scipy. integrate. We refer the code provided in Brandstetter et al. (2022) to reproduce this data generation, however certain small modifications have to be made, like using a fixed At per trajectory and increasing the number of modes in the initial condition.
As for the PDE parameters, we use L = 64 and T = 2.5. For the initial condition, we use a superposition of sinusoidal waves:
20
2x x + Φ3
un (z)
Σ
A sin
L
where for each trajectory, the A; are sampled from continuous uniform in [-0.5, 0.5), the k; are sampled from a discrete uniform in 8). and the &i are sampled from a uniform uniform in [0, 2n]. We discretize [0, 2"] into 26 equispaced points separated by At 0.1. In the experiments in Section 6.1, for each of the four values of the viscosity (0.15, 0. 125, 0.1, 0.075), we generated a dataset with spatial resolution 512 with 2048 training samples and 256 test samples. For the experiment in the sequential model ablation in section G.I, we generated one dataset with viscosity 0.15 in resolution 256, 4096 training samples and 256 test samples.
F.2
BURGERS' ID EQUATION
The ID Burgers' equation can be written as:
Up + udy = Mugz E E [0, x L)
For the Burgers' equation, we take the publicly available Burgers' dataset of PDEBench (Takamoto et al., 2023) with viscosity 0.001. Out of the 10000 samples of the dataset, we use 10% for testing. For training, we found it sufficient to use 2048 samples. Additionally, for training and testing we only used the 20 first timesteps, since we observed that after the 20th timestep the diffusion term of the equation Uyz attenuates all high frequencies and the solution changes very slowly.
E.3 MAVIER-STOKER 2D EQUATION
The incompressible Navier-Stokes equation in the 2D unit torus is given by: E (0,1)).4 / (0,
dw(x,)) t) Øt. + u(x, 1) Vw(x, () = vAw(1, f(x),
V - u(1 t) = 0, w(x, 0) wo(x),
x E (0, 1) E [0, (0. 1)²
For the data generation, we follow the method of Li et al. (2021), yet with different tempo- ral and spatial grids. The initial conditions wi are sampled from Gaussian Random field
N (- -2.5 with periodic boundary conditions. The forcing term is f(11, 0.1 (sin (z) + x2) + cos 2 + x2)). At each timestep, the velocity is obtained from the vor- ticity by solving a Poisson equation. Then, spatial derivatives are obtained, and the non-linear term is computed in the physical space and then dealiased. A Crank-Nicholson scheme is used to move forward time, with a timestep of 10-4 4 We use 512x512 spatial grid which then downsampled to 64x64 for our experiments. For the viscosity
10-3, we use final time of 16 seconds and sample every 0.5 seconds. For the viscosity y = 10 5 we use a final time of 3.2 seconds and sample every 0.1 seconds. For more details on the data generation algorithm, refer to Li et al. (2021).
F TRAINING DETAILS
In this section, we will provide a detailed description of the training hyperparameters used in the KS experiments of Section 6.1, in the Burgers experimente of section C and the Navier Stokes
22

---

Published as a conference paper at ICLR 2025
experiments of section 6.2. We start with the training hyperparameters. All experiments used a learning rate of 0.001. For the number of epochs, in KS and Burgers, the training was done over 200 epochs with cosine annealing learning scheduling (Loshchilov & Hutter, 2017); whereas in Navier Stokes we trained for 300 epochs and halved the learning rate every 90. As for the number of samples, KS and Burgers
trained with 2048 samples and Navier Stokes with 1024 samples. Lastly, we observed that the batch size was a sensitive hyperparameter for both the memory and memoryless models (it seemed to affect both equally) so ran sweep at each experiment to select the best performing one. In the results shown in the paper, KS and Navier Stokes use a batch size of 32, and Burgers batch size of 64.
Another relevant detail is the memory length in training, that is, the number of past states that were fed to the memory layer in the MemÑO model. the KS and Burgers experiments, the maximum memory lengths are 20 and 25 (which are the same as the number of timesteps of the dataset). That means that for the last timestep, the previous or 24 states were fed into the memory layer. However, for GPU memory limitations in Navier Stokes the memory length was 16, half the number of timesteps of each trajectory in the dataset. In this case, the memory was reset after the 16th timestep, i.e. for the 16th timestep the 15 past states were fed to the memory model, yet for the 17th timestep only the 16th timestep was fed. Then, for the 18th timestep. the 17th and 16th were fed, and
As in (Tran et al., 2023), experiments were trained using teacher forcing. This means that for the prediction of the i-th timestep during training, the ground truth of the i previous steps was fed to the model (as opposed to the prediction of the model for such steps).
We ran our experiments on A6000/A6000-ACH GPUs. The Navier Stokes 2D experiments required around 34GB of GPU memory for the batch size of 32 and took around 5 hours to finish, whereas the rest of experiments in ID required lower GPU memory (less than 10GB) and each run took around 1 or 2 hours, depending on the resolution.
G ABLATIONS ON THE MEMNO ARCHITECTURE
In this section we present three ablations regarding the MemNO architecture
G.1 ABLATION: CHOICE OF SEQUENTIAL MODEL
In section 5.3 we introduced MemNO as an architecture framework which allowed the introduction of memory through any choice of a sequential layer, which we chose as S4 in the previous experi- ments. In this section, we explore two other candidates for the sequential layers: Transformer and an LSTM. We introduce Transformer-FFNT identical to S4FFNO except in the sequential layer, where a Transformer and an LSTM are used respectively. The Transformer layer includes causal masking and a positional encoding, which is defined for pos across the time dimension and i across the hidden dimension by:
(T-FFNO) and LSTM-FFNO as two models that are
pos 2t am, Eemaantaal
PE(pos, 2i) = sin
10000
pos
(Eopos, 2i+ 1) cos
22 10000 meddl
We show results for the KS dataset with viscosity v = 0.15 and different resolutions. This dataset was generated using a resolution of 256 and contains 4096 samples, twice as many compared the KS datasets of E, given that Transformers are known to perform better in high-data regimes. The results are shown Figure 7. TFFNO performs significantly worse than S4FFNO across almost all resolutions, and even performs worse than FFNO. In contrast, LSTM-FFNO outperforms FFNO, which shows that MemNO can work with other sequential models apart from S4. The memory term in Equation 6 is convolution in time, which is equivalent to the S4 layer and very similar to
Under this setup. the GPU memory requirements were around 34 GB. Using the full 32 timesteps for training would require memory beyond 48GB, which was heyond our GPU capacity (A6000/A6000-A4l)) GPUs).
23

---

Published as a conference paper at ICLR 2025
Recurrent Neural Network (RNN) style layer, showed in Gu al. (2022). We believe that this inductive bias in the memory layer is the reason why both S4FFNO and LSTM-FFNO outperform FFNO. However, S4 was designed with bias for continuous signals and has empirically proven better performance in these kind of tasks (Gu et al., 2022), which is in agreement with its increased performance over LSTMs in this experiment. Additionally, we observed that LSTMs were unstable to train in Navier Stokes 2D datasets.
Lastly, we make two remarks. Firstly, we believe that Transformers performed worse due over- fitting, given that the train losses were normally comparable or even smaller than the train losses of the rest of the models at each resolution. We hypothesize that the full access to the past of Trans- formers models might lead to exploiting spurious correlations during training. Modifications of the Transformer layer or to the training hyperparameters as in other works (Hao et al., 2024; Cao, 2021; Hao et al., 2023) might solve this issue. Secondly, recently there has been surge of new sequential models such as Mamba (Gu & Dao, 2023; Dao Gu, 2024), RWQK (Peng et al., 2023), xLSTM (Beck et al., 2024) or LRU (Orvieto et al., 2023). We chose S4 over Mamba-type architectures because in our experiments the PDE temporal dynamics do not change, and thus we do not expect the input-dependent selectivity mechanism to be necessary. However, study which of these sequential model has better overall performance, and hope that our study on the settings where the memory effect is relevant can help make accurate comparisons between them.
leave it future work to
KS, V= 0.15. nRMSE vs Resolution
1.4
FFNO
12
SAFFNO
T-FFNO
10-
LSTM-FFNO
nRMSE
0.8
0.6
0.4
0.2
0.0
10
18
22
26
34
38
50
54
58
66
70
74
Resolution
Figure 7: Performance of FFNO, S4FFNO and T-FFNO and LSTM-FFNO KS with viscosity
v= == 0.15.
G.2 ABLATION: MEMORY LAYER CONFIGURATION
In Section 5.3 we introduced the memory layer in MemNO as a single layer to be interleaved with neural operator layers. In our experiments, inserted it after the second layer of four layer neural operator. In this section, we explore the impact of having different layer configurations, including the possibility of having several memory layers. We will denote the configurations with sequence of S and T letters. S means a neural operator layer (some sort of Spatial convolution). and T a memory layer (some sort of Time convolution). For example, SS7SS denotes the architecture of our experiments, where we have 2 neural operators layers, followed by memory layer, followed by other 2 neural operator layers. Similarly. SSSST denotes 4 neural operators layers followed by a memory layer. In Table 3, we present the results for the KS dataset with v = 0.1 and final time of 4 seconds for several models. We include the S4FFNO model we used in previous experiments in the first row (with configuration SS7SS), and the FFNO model in the last row. In the middle rows, we show different configurations of memory and neural operator layers. It can be observed that all models with at least a memory layer outperform FFNO. There are slight differences between configurations, yet we focused mainly on the comparison to the memoryless model. For that reason, fixed SSTSS configuration in our previous experiment, which was the most efficient (only one memory layer) and symmetric. We leave as further work determining if there are settings where a given configuration pattern can be substantially better than the rest.
24

---

Published as a conference paper at ICLR 2025
nRMSE
Architecture
Resolution 32
Resolution 48
Resolution 64
S4FFNO (SS7SS)
0.123 + 0.011
0.086 + 0.004
0.015 0.001
S4FENO (SSSS7)
0.142 0.009
0.069 + 0.001
0.017 0.001
S4FFNO (STSSTS)
0,141 + 0.006
0.064 = 0.002
0.019 + 0.001
S4FFNO (STSTSTST)
0.113 + 0.006
0.070 + 0.004
0.017 + 0.001
S4FFNO (TSSSS)
0.129 0.007
0.080 = 0.003
0.017 0.001
FFNO
0.294 + 0.004
0.138 + 0.013
0.021 0.002
Table 3: KS, -= = 0.1. The final time is 4 seconds and the trajectories contain 20 timesteps. For each architecture, we tried learning rates (0.002, 0.001, 0.0005 and 0.00025, each with three different seeds. We present the results of the learning rate with the lowest nRMSE averaged across the three seeds. The standard deviation also with respect to the seeds.
G.3 ABLATION: S4U-NET
The experiments in Section 6 used S4FFNO our proposed memory model, which was the instan- tiation of the MemNO framework (Section 5.3) with FFNO as the Neural Operator and S4 as the memory layer. In the ablations of Section G.1 we showed that although S4 was the best performing memory model, LSTM also provided good performance, showing the versatility of the framework and the importance of adding memory (regardless of the specific architecture). In this section, we show that the MemNO can also use a different Neural Operators as the Markovian layer. In partic- ular, we instantiate the MemNO framework using U-Net as the Markovian Neural Operator, and S4 as the memory layer with state dimension of 16.
The results for the KS experiment (Section 6.1) and Burgers' experiment (section C) are shown in Table 4. As in the case of S4FFNO, S4U-Net also improves the performance of U-Net in the cases of low resolution. This shows that modeling memory is useful in low resolution for architectures other than FFNO. However, S4U-Net has worse performance than S4FFNO. We reiterate that the main contribution of our work is studying when modeling memory is helpful, as well as providing flexible ways to incorporate it into existing neural operators. We leave for future work the study of the trade-offs between different memory models under different setups.
aRMSE L
Architecture
Uses memory
Resolution
KS
Burgers'
0.075
0.1
0.125
0.001
U-Net
No
0.542
0.511
0.249
0.188
S4U-Net (Ours)
Yes
0.364
0.277
0.104
0.096
U-Net
No
0.147
0.062
0,022
0.171
64
S4U-Net (Ours)
Yes
0.114
0.052
0.026
0.070
U-Net
No
0.033
0.027
0.014
0.112
128
S4U-Net (Ours)
Yes
0.058
0.030
0.022
0.057
Table 4: nRMSE values for the S4U-Net architecture at different resolutions for Burgers' and KS with different viscosities. The values of U-Net are the same the ones in Table I and are provided here for context. More details on training are given in Appendix F, on the KS experiment on 6.1 and on the Burgers' experiment in Appendix C.
H ABLATIONS ON S4FFNO AND FFNO PERFORMANCE
H.1 ABLATION: MEMORY WINDOW LENGTH
In this section we present an ablation on the memory window length of the S4FFNO architecture. We recall that in our experiments of Section 6.1 we had discretized grid T = (to, 4, -.. (,), and
25

---

Published as a conference paper at ICLR 2025
nRMSE for S4FENO with varying memory window length
0.2
0.20
0.15
S4FFNO memory]
FFNO
010
(Mtemor window length)
Figure 8: nRMSE for S4FFNO with varying memory window length, for the KS experiment of Section 6.1 (u = 0.1 and resolution 32). A memory window of K means that the S4 model only has access the memory of the last K timesteps predict the next one. At training time, the sequence length is split into chunks of K timesteps and each chunk is trained independently. At inference time, the S4FFNO is given access to the last K predicted timesteps to make the next prediction.
in order to predict the solution at timestep 4 S4FFNO had access to the memory of all previous timesteps (i.e. the S4 model operated on the hidden dimensions v(ty) for 0 S i 1). In this section, we study what happens when S4 is only fed the last K timesteps, i.c. in order to predict t, S4 operates on the solution at timesteps 1- K, K+, 4. The results are shown Figure 8.
$4FFNO improves performance as the window length increases, illustrating that the reason for the increased performance of S4FFNO the capacity to model the memory of past timesteps.
H.2 ABLATION: FFNO MODEL SIZE
Now we consider what happens when we increase the model size of FFNO. Based on the results of Section 6.1, S4FFNO outperformed FFNO when they had similar compute budgets (see Table 2). However, S4FFNO still outperforms FFNO when FFNO has much higer compute and parameter budget, as it can be seen in Table 5. Thus, we conclude that S4FFNO has superior performance due to the possibility of modeling memory from past states.
nRMSE.. +
# Params (millions)
Architecture
Hidden Dimension
Layers
Resolution
Resolution 48
Resolution
Resolution
S4FFNO
128
0.108
0.045
2.8
3.9
FFNO
128
0.440
0.238
2.8
3.8
FFNO
128
0.361
0.181
5.5
7.6
FFNO
256
0.435
0.252
11.1
15.3
FFNO
256
0.346
0.194
22.2
30.6
Table 5: Performance of S4FFNO and different model sizes FFNO on the KS experiment with viscosity 0.1 and resolutions 32 and 48. The experimental details are the same 6.1.
1 ABLATION ON THE MODELING OF LOCAL SPATIAL INFORMATION
In Section 5.3 we presented the MemNO framework as a way to build neural operators that model memory, which is inspired by the Mori-Zwanzig formalism. In Equation 6, it can be seen that the memory term depends on the differential operator L, which can depend on the spatial derivatives of the input. The memory layer of MemNO is applied independently hidden representation the input, and thus it does not model spatial derivatives explicitly, although the hidden dimension can in principle encode such local information implicitly.
each spatial dimension of the
26

---

Published as a conference paper at ICLR 2025
In this section, present several modifications to the S4FFNO architecture, which attempt to pro- vide an inductive bias for modeling spatial information more explicitly. In particular, we try the following modifications:
S4FFNO + Input Gradients: The numerical gradients of the input are fed to the model encoder. Specifically, Wy is represented as a vector of spatial shape S, and the spatial gradients are approximated using second-order accurate central differences method (using the method torch. gradient of the pytorch library (Paszke et al., 2019)). This would directly facilitate storing information about the first order spatial gradients in the hidden dimension.
S4FFNO + Convolutional Encoder: Instead of having linear layer as encoder, we sub- stitute it with ID Convolution of kernel size 3 across the spatial dimension.
S4FFNO + Convolution before memory layer: Before feeding the sequence of hidden presentations (Ug, vi,.... 0t] the memory layer S4, we apply a 1D Convolution across the spatial dimension with kernel size of 3.
The performance of such modifications on the KS experiment of Section 6.1 with viscosity V = 0.1 shown in Table 6. It can be that none of the modifications improve performance compared to the baseline significantly. We hypothesize that training becomes more difficult with these architec- ture modifications, and thus they do not provide an improvement in performance, We leave future work to consider other architectures to model spatial information more explicitly.
#RMSE
Architecture
Resolution
Resolution 32
Resolution 48
Resolution 64
Resolution 80
FFFNO
1.106
0.4461
0,2325
0.328
0.0040
S4FFNO (Base)
0.3318
0,1081
0.0455
0.0111
0.0048
S4FFNO Input Gradients
0.4433
0.1120
0.0482
0.0120
0.0054
S4FENO
Convolutional encoder
01,3440
0.1053
0.0571
S4FFNO Convolution before memory layer
0.4283
0.1224
0.0491
0.0142
0.0057
Table 6: Performance comparison between FFNO, S4FFNO and several architecture modifications S4FFNO aimed at introducing an inductive bias to model spatial information, see Section The experiment is performed on the the KS PDE with viscosity 0.1 under the same setup as Section 6.1. The architecture details of S4FFNO and FFNO are provided in Appendix B.
J MEMNO DIAGRAM
In Figure 9 we provide a diagram for the MemNO framework (Section 5.3) in ID.
K QUANTIFYING THE EFFECT OF MEMORY
We include the proof for Theorem
Proof. We proceed to the Equation first. Note that un (t), Vt 2 0 can be written as uz
40 c + 0j e1. Moreover, by Proposition 1, we have
dag (t)
àt = 2Ba
(17)
0 = a + Bag
(18)
Øt
In matrix form, these equations form linear matrix ODE:
0
2B
(8)
=
B
a (t)
27

---

Published as a conference paper at ICLR 2025
1/3
W
Encoder
Noura
Decader
Poation godd
Vienery Sager
WWM
Neural Operanor
Naura Opanstor
Decoder
Paeltengge
MMM
HAPH
Neraa Opeartr
heurs Openttor
Figure 9: Diagram of the MemNO framework in ID (Section 5.3). denotes spatial dimension, H denotes hidden dimension, and L number of layers. The memory layer inserted the middle, although the framework works with other configurations.
= exp (
{
2B
Gg
(0)
The solution of this ODE is given by
By the first statement
1
aj (0)
(0)
(0)
of Lemma and the non-negativity of
dg
ai
get:
0 10e
(t)
VBr
0g + 0
(0) (0)
(19)
(t) a1 S 10e
2Bt
(0)
dg
+
(20)
We proceed to Equation 10. Note that for any & 2 0, we can write u2(s) = ab ep + a ei with
(0) (0) and a dj (0) By Proposition I, we have
dg dg
QCug(x) = Bag (2(2)
Moreover, given function v(x), the action of the operator exp [L(1)) on vr is given by the solution w(f,x)) to the PDE
w(1, x) = QLw(ltxx)
01
w(0, = u(x))
If w((,, = EnEN, (1) e, and Vn No, bn 2 0, we are interested in solving the previous PDE with initial conditions b/ Ba(a)
(5) and bn Vn +
We claim that the coefficients a!
> and Vn E {0. 1). For this is by definition, and we will prove it for all / by way of contradiction. Suppose the claim is not true, then there exists and some n² E {0, 1) such that a(4h) and a((²)
vn {0, 1} and Vs < t". But from 0. However, it can be easy to see that if ((4) 0 Vs < t', then PiLu,(t) > 0 and PiC fo exp QC(t 3) (s)ds > 0. Therefore, from Equation 10, à (t') 5r dn 0, which is contradiction.
continuity this implies that there exists 0 t' < t* such that a (('))
o7 a
This claim implies that 6 2 Vn E N, and in turn it implies that (1) > 0 Vn E N,t > 0. Applying QL results in the following inequalities for the coefficients / (1)
b ()) (t),
28

---

Published as a conference paper at ICLR 2025
0
1/ 2 11 + Bbb 2 Bby
(21)
0t
1 > (6()
4b (1) Bby z Bbb + Bbq
(22)
0 (1)
0/
B62
(
96 2 Bb2 (t)
(1)
(23)
Thus, we can write linear matrix ODE for the vector
(t)
b2
(1) (())):
0
0 B
0
B
B
(24)
dt
0
B
0
Therefore, using Lemma 2, for sufficiently large B we have (1-5) Be (2 (e)
01
Hence, if we write f exp( - })} QCuz(s)dd in the basis (entneNo the coefficient for e will be lower bounded by
L' 1 Be B(t 00 ds
10
Applying the second statement of Lemma and using the non-negativity of d and (0) we have a 2 The
(0)
V
I
+ ). Hence, the coefficient for e2 is lower bounded by
[ Be (2
10 L Go + ( ds 2
BI vBBI
(0) (o + (0)
10
100
We finally need to consider what happens after applying the outermost operator PiL. Because of Proposition again, applying L makes the coefficient in front of e at least Finally, applying P1 preserves the coefficient in front of e1. Hence, equation Equation 10 results in the following evolution inequalities:
100 e v2Bi
B
(1
(0)
(0)
(00
a1
0ā!
2Ba (t)
(25)
0t
0a (t)
2 d Báy (0) +
Bat
V2BB (0) + a(T))
(26)
100
Using the second statement of Lemma again we have that @n 2 1 e (u)0) a Thus,
dropping the (positive) term a in equation 26, we have:
da (t)
1
Bt
Be Be
/BB
0 + +a(
(0) 10)
(27)
Öt
10
100
Integrating this equation yields:
a
a
(0)
/2B
(2BBt+ 10v2-11 (a. +
(28)
200
Thus, have (0)
2 Bie LBt
d
(0)
+ a
(0)
Together with equation 19, the claim of the Theorem
follows.
2Bt
Lemma 1. There exists B > 0 sufficiently large such that for all
0 the matrix
Bt
satisfies:
Vi.j E (1,2), exp
(0:
Bt
2Bt
< 10exp
(v2Bt)
(29)
i.j
Vi,j E (1,2),
2Bt
2
exp
v2BBt)
(30)
Bt
10
ij
29

---

Published as a conference paper at ICLR 2025
Proof. By direct calculation, have:
(0r
2Bt Bt t )
V8B² -(B()) 2Bh(B,t)
4Bh(B,t)
exp
2/8B² +1
V8B² + 1g(B, + k(B,t)
where:
(BBtt = et (VSB+++++) +e -((BB+++I-1))
b(B,t) t) = ed (V8B++1++)) _ (V88*+1-1)
Thus, the statement follows.
0
B
0
Lemma 2. For all B > O, the matrix
B
0
B
satisfies:
0
B
0
)
B
0
exp sop(VBB)
vi,j E (1,2,3),
B
0
B
(31)
B
0
10
4
Proof. By direct calculation:
0 B
exp
B
=
B
i
2e + e22B +1
V2e2V2B - -22
-2eV2B + 22VBB +1
1
e -22B
V2e2VBB -2
2(e2 +1)
V2e2/BB - V2
4
V2B
2e
+ e22BB
+1
22eB V2
2e V2B + 222B +1
Thus, the statement follows.
L DEFINITION OF PERIODIC BOUNDARY CONDITIONS
For completeness, we give a precise definition of periodic boundary conditions for the PDE defindd in Definition 2:
Definition 6 (Periodic Boundary Conditions). For PDE given by Definition with Ω |0, L] we define the periodic boundary conditions as the condition:
u( ... -k-1,, 0, Tk+1. -*a) = u(z1, .. -x-1, L, 2K+1, ..))
for all (Zi, Tt Tk41,"'1 (F) E [o, L]d-1 and all i: 1,... d.
30

---

Published as a conference paper at ICLR 2025
M MEMNO PSEUDOCODE
In Figure 10 we provide pseudocode for the MemNO framework (Section 5.3) in PyTorch (Paszke al., 2019).
from import List einops import rearsange
import borch
import torch.nn ac
Motation:
batch
Ti terporai
St spatial dimension (10)
52
hiddon dimension
iiiiiiitttt (salf, narkoriar_tayeze, List Mod.Noaaaa,, namozy_ na.Jedelo, monoey_position: int)
Arga:
markovian layers: List nadlldulo
inguts
outouts
nonory layert anoHoouie maps inputs 57 outpets (2, Z, 81.
This must de a causal and sequent rai nodel.
monory_porition Position of monery layez in the model.
aupez 1
self solf.markorisn_ipper self self.menory layee nenocy_
Moduichist (narkovizn
self.nenory_pssiti moncry_position
def (usva at Lorch. Pemnor) -2 Lorch.Tera01: Tersos:
encoded representation the soducion PDS 8, (1.0. fi_0, ... (L(I-1]])
Oukpot:
Prediction the solution of the
Linestep
(1.0. [ihatju]_ ... (aat
B, 2, x.shape "b
markav__aayee self.mackovian layers metoyyy Lon
x layar (x]
1b a5 My bel)
self. sel..ownury_
(sasrangu (b: a1 (b 6) br,, beB)
for narkov. layer in sel F.narkovi an layera Feerolllliieeooyg.oooitiii
mackov_layer
hi, be21
Figure 10: Pseudocode for the MemNO framework (see Section 5.3) 1-D, Encoder and Decoder layers are omitted for clarity (see details in Appendix B).
31


Reviews:
[{"review_id": "ddNnrUuBpA", "reviewer": "Reviewer_8n1K", "summary": "The paper proposes a new form of neural operators (NOs) that are inspired by the Mori–Zwanzig (MZ) formalism. This operator is to model spatio-temporal dynamics that arise in time-dependent (nonlinear) partial differential equations (PDEs). The paper points out that the previous roll-out-of-NOs which is caused by memorylessness and, thus, proposes to implement the memory term similarly in the MZ formalism. The specific example architecture has been investigated: Fourier NO + S4, which are empirically shown to be more effective in capturing the dynamics. Moreover, the paper points that the previous benchmark PDEs are easy (including relatively only frequency) and, thus, proposes a new set of benchmark datasets that gives challenge to previous NOs.", "questions": "Please see the questions in the weaknesses above.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper is well-motivated and proposes solution methods as well as a new set of datasets.\n\nThe paper is well-written, providing necessary details without going into too much detail and explaining why including the memory term can be beneficial in modeling complex dynamics. \n\nThe paper introduces a new challenging dataset (although not extensive), the dataset could be used as a next important benchmark for future works.\n\nThe empirical results seem to show that the proposed method is more effective in modeling the dynamics. Interestingly, the results are somewhat aligned with the effect of the MZ formalism in heterogeneous systems; that is, on a coarse mesh (on a coarse-grained or homogenized system in MZ), the Markovian surrogate model can be aided by the memory term to capture the influence the state defined on a finer mesh (the eliminated variables in MZ).", "weaknesses": "It is not easy to see exactly how the architecture is designed. Including diagrams showing the forward computational path will be very helpful. \n\nAlthough taking a cue from the MZ formalism is good, some discussions on the related topic such as closure modeling for surrogate models seems to be missing.\n\nFNOs are known to select low-frequency modes and are expected to struggle with high-frequency solutions. Are transformer-based neural operators expected to behave similarly? Could the authors provide why GKT performs poorly? \n\nAlso, could the authors provide experimental results with some advanced transformers-based neural operators? such as OFormer (Li, et al, TMLR, 2023) and their variants (for example, FactFormer, Li et al, NeurIPS 2023, 2023)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "yLgxvb8TPK", "reviewer": "Reviewer_vkVf", "summary": "The authors propose a variant of a factorized Fourier neural operator with memory and show how it can improve over the standard Markovian neural operator for modeling PDEs with small grid sizes or noise added to the observations. Additionally, some theory motivated by the Mori-Zwanzig theory of model reduction is given to motivate and explain the success of the addition of memory (Eq. 6). This success in certain regimes is observed in experiments.", "questions": "See Weakness 2.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1) The problem of whether to use memory to model time-dependent PDEs is well-defined and well-motivated.\n2) The theoretical motivation (Eq. 6) and example (Section 4) are clear, concise, and show the possible situations where adding memory in Neural Operators could improve performance.\n3) The two experiments varying resolution and additive noise nicely demonstrate regimes where memory can be helpful. Figure 6 also provides good context in this regard.", "weaknesses": "1) There should be a brief discussion about the algorithmic cost of adding memory to the FNO. The authors use the diagonal matrix parameterization of S4, so this should be a reasonable increase in cost. Empirical timing would also be helpful to make this point.\n2) In the added noise experiments (Section 6.2) it is expected that memory improves performance because the memory unit can approximately filter the input data. The discussion surrounding this experiment is less fleshed out than that of the resolution experiments in Section 6.1 and does not have an obvious connection to the theory and example presented in sections 3 and 4. What practical situation is this intended to model?\n3) lines 1119-1121: this is a strange training setup, but maybe necessary due to the GPU memory limitations.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "aMclqF6xPg", "reviewer": "Reviewer_oCpa", "summary": "This paper analyzes the advantages of incorporating memory into AI models for solving time-dependent partial differential equations (PDEs). Traditional AI-driven approaches often rely on a Markovian assumption, where future states are predicted solely from the current state. However, the authors argue that including memory (i.e., information from past states) improves model accuracy, particularly when solutions are low-resolution or contain noise.\n\nThe authors further present the Memory Neural Operator (MemNO), a novel architecture that combines state space models (namely S4) and Fourier Neural Operators (FNOs) to capture both spatial and temporal dynamics effectively. Theoretical foundations based on the Mori-Zwanzig theory suggest that memory-inclusive models can outperform Markovian models in specific scenarios, particularly with high-frequency Fourier components. Empirical results show MemNO achieves up to 6x lower test error compared to memoryless models.", "questions": "See weaknesses", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "* This paper seeks to formalize and theoretically analyze the role of memory in learning time-dependent PDEs, distinguishing itself from prior studies that have approached memory in this context primarily from an empirical standpoint.\n* The experiments and ablations are thorough and well-executed. Additionally, the empirical results are presented using relative errors, which is exemplary and should be the standard in this field. Unfortunately, this point deserves emphasis, as many other studies overlook this.\n* the paper is well written and clearly structured\n* the paper introduces recent state-of-the-art state-space models to the field of learning PDE solutions", "weaknesses": "* The paper claims to improve learning the solution of time-dependent PDEs theoretically grounded in the \n    the Mori-Zwanzig formalism. However, equation 6 of the paper clearly states that the 'memory term' (i.e., second summand on the right hand side of the equation) involves the differential operator of the PDE, i.e., a local spatial operator. But the proposed method does not take into account spatial interactions and applies the memory cell (S4 layer) to each element of the spatial dimension independently. This is clearly not covered by the formalism and should be stated very clearly. That being said, it would be interesting to incorporate spatial interactions via Graph Neural Networks acting on the nearest neighbors for instance.\n* The experiments are interesting and extensive. However, there are some issues that need to be addressed. It's not clear up to which extend the performance increase is due to the memory or simply due to the very expressive S4 model. One good way to analyze this empirically would be to plot the performance of the proposed method on the y-axis, while on the x-axis the size of the look-back window is varied, i.e., how many past instances are fed into the memory cell. Currently, all available past instances are fed to the memory cell. If the performance does not change significantly by varying that, this would indicate that memory is not as important and the performance difference can be explained by the different architectures. Figure 7 in the appendix already suggests that the architecture choice influences the performance.\n* A general weakness: Why do we need an advanced AI model with 5M parameters only to solve a simple 1D Burgers equation that can be solved in milliseconds with traditional methods? I acknowledge that the aim of the paper is to improve AI models applied to PDEs. However, it would still be important to draw some connections to traditional methods, which in particular for the experiments considered in this paper are vastly superior. A good starting point to look into this is the referenced paper by McGreivy and Hakim (2024).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "WgYGZSUzZA", "reviewer": "Reviewer_akXT", "summary": "This work applies a state space model to model the memory term of a PDE that (under assumption) decomposed by the Mori-Zwanzig theory. The method is mostly well-motivated and demonstrates improvement in solving PDEs in low resolution. Overall the paper addresses an important issue in the field of neural operators.", "questions": "1. The main question I have is regarding the flexibility of the proposed method. The author showcased that the memory term modeled by S4 achieves improved performance, but in Appendix G, the memory term modeled by LSTM can be almost just as good. I hope the authors can also highlight this in the main text as I think this shows that adding the memory term is important, no matter the architecture. \n2. Related to the above, can similar memory terms be added to the other benchmarks like U-Net neural operator? If so this can also be highlighted as an ablation study such that the importance of the memory term can be further recognized. \n3. When the high-frequency mode matters, would it simply be more effective to increase the complexity of FFNO rather than applying the memory term? Can the author briefly discuss how this trade-off affects the choice of using a memory term (or not) in this case?\n4. How well can the model perform for zero-shot super-resolution once the model is trained?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper is theoretically well motivated. \n2. The proposed method for memory operators can be implemented easily.\n3. The method demonstrates improvement in performance for solving PDEs.", "weaknesses": "1. The motivation for using S4 as the model of choice is not convincing.\n2. Lack of numerical comparison against benchmarks in 2D cases.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Oral)