On the Benefits of Memory for Modeling Time-Dependent PDEs

## Abstract

Abstract Data-driven techniques
have emerged as a promising alternative to traditional numerical methods
for solving partial differential equations (PDEs).
These techniques frequently offer a better trade-off between computational cost and accuracy for many PDE families of interest. For time-dependent PDEs, existing methodologies typically treat PDEs as Markovian systems, i.e.,
the evolution of the system only depends on the “current state”, and not the past states. However, distortion of the input signals — e.g., due to discretization or low-pass filtering — can render the evolution of the
distorted
signals non-Markovian.
In this work, motivated by the Mori-Zwanzig theory of model reduction,
we investigate the impact of architectures with memory
for modeling PDEs: that is, when past states are explicitly used to predict the future.
We introduce Memory Neural Operator (MemNO),
a network based on the recent SSM architectures and Fourier Neural Operator (FNO).
We empirically demonstrate on a variety of PDE families of interest that when the input is given on a low-resolution grid, MemNO significantly outperforms the baselines without memory,
achieving more than 6 × fragments 6 6\times less error on unseen PDEs. Via a combination of theory and experiments, we show that the effect of memory is particularly significant when the solution of the PDE has high frequency Fourier components (e.g., low-viscosity fluid dynamics), and it also increases robustness to observation noise.

## 1 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.
Recently, data-driven neural network based approaches (Li et al., ; Lu et al., ) have emerged as an attractive alternative to classic numerical solvers, such as finite element and finite difference methods (LeVeque, ).
Classical approaches are computationally expensive in high dimension, and struggle with PDEs which are very sensitive to initial conditions. Learned approaches can often negotiate these difficulties better, at least for the family of PDEs they are trained on.

One example of a data-driven approach is learning a *neural solution operator*, which for a time-dependent PDE learns a time
evolution map that predicts the solution of the PDE for future time steps.
The operators are frequently autoregressively parametrized, such that the network predicts future states based on previous ones (Li et al., , ), and the number of past states the model is conditioned on serves as “memory” and is treated as
a tunable hyperparameter.
Recent works (Tran et al., ; Lippe et al., )
suggest that optimal performance across various PDE families
can be achieved by conditioning the models only on the immediate past
state—i.e., treating the system as Markovian—though this is in settings in which the training data is very high-resolution.

In many practical settings, we expect to only observe a part of the system. This could be due to limited resolution
of the measurement devices collecting the data, inherent observational errors in the system, or prohibitive computational difficulty generating high-quality synthetic data.
This can lead to significant information loss, particularly in systems like turbulent flows (Pope, ) or shock formation in fluid dynamics (Christodoulou, ),
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-Markovian.

More precisely, Mori-Zwanzig theory (Mori, ; Zwanzig, ; Ma et al., ) is an ansatz to understand the evolution of a subspace of system (e.g., the top $k$ Fourier components). Under certain conditions, this evolution can be divided into a *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 an *“unobservable” term*, which depends on the values of the initial conditions orthogonal to the chosen subspace.

In this paper, we study the effects of explicit
memory when deploying neural operators for time-dependent PDEs. By memory
we loosely mean a representation of the previous states
of a PDE.
We focus on PDE families that take the form $\partial_{t}u(x,t)={\mathcal{L}}u(x,t)$
where $u(x,t):\Omega\times[0,T]\to\mathbb{R}$ is a time-dependent function
defined over the domain $\Omega$, and ${\mathcal{L}}$ is a (possibly non-linear) operator.
This is a generic form of a time-dependent PDE system and contains many PDE families of interest in different application domains (e.g., heat diffusion, Navier-Stokes, Kuramoto-Sivashinsky, Black-Scholes, Schrödinger equation to name a few).

We introduce Memory Neural Operator (MemNO),
an architecture which combines Fourier neural operator (FNO) (Li et al., ; Tran et al., )
and
the S4 architecture (Gu et al., , ).
The MemNO architecture can be seen as an adaptation of the FNO (Li et al., )
architecture where the FNO layers model the spatial dynamics of the PDE
while the S4 layers (Gu et al., , )
maintain a compressed memory of the past states.
We choose S4 models over
recurrent architectures like LSTM (Hochreiter and Schmidhuber, )
due to superior performance in modeling long range dependencies (Gu et al., ; Tay et al., ), ease of training,
and favorable memory and computational scaling with both state dimension
and sequence length.
Through our experiments we
show that for PDEs observed on low resolution grids and/or with observation noise,
MemNO outperforms their Markovian (memoryless) baselines—achieving $6\times$ less
loss on unseen PDEs. Our contributions are as follows:

- •
We introduce MemNO, a Memory Neural Operator architecture which uses a *combination of FNO layers and S4 layers* to model the spatial and temporal dynamics of a PDE. The S4 layers explicitly model memory, by introducing a weighted average over the past (compressed) states.
- •
Even in relatively simple linear PDEs, we theoretically show the *memory term can result in a 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 show
that when the input is supplied on a
*low-resolution grid*, or has additional *observation noise*, memory based architectures outperform the best performing FNO based baselines by a significant margin.
- •
Finally, we empirically show that this effect is more pronounced for PDEs which result in *solutions with high order
frequency modes*, and introduce a metric which indicates
when memory based models will have most impact.

## 2 Related Work

Data-driven neural solution operators (Chen and Chen, ; Bhattacharya et al., ; Lu et al., ; Kovachki et al., )
have emerged as the dominant approach for
approximating
PDEs,
given their ability to model multiple families of PDEs at once,
and relatively fast computation at inference time.

Recently, many architectures have been proposed to improve the performance of neural
operators across multiple families of PDEs,
Li et al. () designed
the Fourier Neural Operator (FNO), a resolution invariant
architecture
that uses a convolution based integral kernels
evaluated in the Fourier space.  Tran et al. () later
introduced Factorized FNO (FFNO) architecture, that builds upon and improves
the FNO architecture by adding separable spectral layers and residual connections.
Additionally, they perform extensive ablations that point out
that strategies like the Markov
assumption—i.e., predicting a state only from its immediate prior state—is optimal
and outperforms models that use the history of past timesteps as input.  Lippe et al. () performed a similar study for long rollouts of the PDE solution and concluded the the optimal performance is indeed achieved under the Markovian assumption.
However, we show that
when there is a loss of information in the observation a PDE,
a model that the uses the history
of past states
outperforms its
Markovian counterpart, often achieving $6\times$ less error on unseen PDEs from the same family.

Our work is motivated by the Mori-Zwanzig formalism (Zwanzig, ; Mori, )
which shows that
a partial observation of the current state of the system
can be compensated using memory of past states.
Our work is also inspired by a previous study on the effects of memory
in modeling PDE dynamics by  (Ma et al., ). Here the authors
draw parallels to the Mori-Zwanzig equations and LSTM (Hochreiter and Schmidhuber, )
to model the dynamics of the top Fourier components of a time-dependent 1D Kuramoto-Sivashinsky
and 2D shear flow equations, one single PDE at a time.
However, in our work, we study the benefits of memory in neural operator settings, i.e,
we have a single model that learns the dynamics of an entire family of PDE at once.
Furthermore, we use
the S4 state space model architecture (Gu et al., ) to model
the temporal dependencies, which in our experiments has better performance and is more stable than LSTMs.

## 3 Preliminaries

In this section, we introduce several definitions, as well as background on the Mori-Zwanzig formalism as applied to our setting.

### 3.1 Partial Differential Equations (PDEs)

###### Definition 1 (Space of square integrable functions) .

For integers $d$, $V$ and an open set $\Omega\subset\mathbb{R}^{d}$, we define $L^{2}\left(\Omega;\mathbb{R}^{V}\right)$ as the space of square integrable functions $u:\Omega\rightarrow\mathbb{R}^{V}$ such that $\|u\|_{L^{2}}\leq\infty$, where $\|u\|_{L^{2}}=\left(\int_{\Omega}\|u(x)\|^{2}_{2}dx\right)^{\frac{1}{2}}$.

###### Notation 1 (Restriction) .

Given a function $u:\Omega\rightarrow\mathbb{R}^{V}$ and a subset $A\subset\Omega$, we denote $u\hskip 1.0pt\rule[-2.15277pt]{0.43057pt}{5.38193pt}\hskip 1.0pt_{A}$ as the restriction of $u$ to the domain $A$, i.e. $u\hskip 1.0pt\rule[-2.15277pt]{0.43057pt}{5.38193pt}\hskip 1.0pt_{A}:A\rightarrow\mathbb{R}^{V}$.

The general form the PDEs we consider in this paper will be as follows:

###### Definition 2 (Time-Dependent PDE) .

For an open set $\Omega\subset\mathbb{R}^{d}$ and an interval $[0,T]\subset\mathbb{R}$, a Time-Dependent PDE is the following expression:

$$ $\displaystyle\frac{\partial u}{\partial t}(t,x)$ $\displaystyle=\mathcal{L}[u](t,x),\qquad\;\;\;\forall t\in[0,T],x\in\Omega,$ (1) $\displaystyle u(0,x)$ $\displaystyle=u_{0}(x),\qquad\qquad\;\;\forall x\in\Omega,$ (2) $\displaystyle\mathcal{B}[u\hskip 1.0pt\rule[-2.15277pt]{0.43057pt}{5.38193pt}\hskip 1.0pt_{\partial\Omega}](t)$ $\displaystyle=0,\qquad\qquad\qquad\;\;\forall t\in[0,T]$ (3) $$

where $\mathcal{L}:L^{2}\left(\Omega;\mathbb{R}^{V}\right)\rightarrow L^{2}\left(\Omega;\mathbb{R}^{V}\right)$ is a differential operator in $x$ which is independent of time, $u_{0}(x)\in L^{2}\left(\Omega;\mathbb{R}^{V}\right)$ and $\mathcal{B}$ is an operator defined on the boundary of $\partial\Omega$, commonly referred as the boundary condition.

Unless otherwise stated, both in the experiments and in the theory we will largely work with *periodic* boundary conditions:

###### Definition 3 (Periodic Boundary Conditions) .

For $\Omega=[0,L]^{d}$, we define the periodic boundary conditions as the condition:

$$ $u(x_{1},\cdots,x_{k-1},0,x_{k+1},\cdots x_{d})=u(x_{1},\cdots,x_{k-1},L,x_{k+1},\cdots x_{d})$ $$

for all $(x_{1},\cdots,x_{k-1},x_{k+1},\cdots,x_{L})\in[0,L]^{d-1}$ and all $k=1,\cdots,d$.

Finally, we will frequently talk about a grid of a given resolution:

###### Definition 4 (Equispaced grid with resolution $f$ ) .

Let $\Omega=[0,L]^{d}$. An equispaced grid with resolution $f$ in $\Omega$ is the following set $\mathcal{S}\subset\mathbb{R}^{d}$:

$$ $\mathcal{S}=\left\{\left.\left(i_{1}\frac{L}{f},\cdots,i_{k}\frac{L}{f}\right)\right|0\leq i_{k}\leq f-1\ \text{for}\ 1\leq k\leq d\right\}.$ $$

We will also denote by $|\mathcal{S}|$ the number of points in $\mathcal{S}$.

###### Definition 1 (Space of square integrable functions) .

###### Notation 1 (Restriction) .

###### Definition 2 (Time-Dependent PDE) .

###### Definition 3 (Periodic Boundary Conditions) .

###### Definition 4 (Equispaced grid with resolution f 𝑓 f ) .

### 3.2 Mori-Zwanzig

The Mori-Zwanzig formalism (Zwanzig, ) deals with cases where an equation is known for a full system, yet only a part of it is observed. It leverages the knowledge of past observed states of a system to compensates for the loss of information that arises from the partial observation. In our paper, partial observation can refer to observing the solution at a discretized grid in space or only observing the Fourier modes up to a critical frequency. In particular, the Mori-Zwanzig formalism in the context of time-dependent PDEs is well-known in the Physics literature as the Nakajima–Zwanzig equation ((Nakajima, ))

Now, we will apply the Nakajima–Zwanzig equation to our setting. Assume we have a PDE as in Definition [2](#Thmdefi2). Let $\mathcal{P}:L^{2}\left(\Omega;\mathbb{R}^{V}\right)\rightarrow L^{2}\left(\Omega;\mathbb{R}^{V}\right)$ be a linear projection operator. We define $\mathcal{Q}=I-\mathcal{P}$, where $I$ is the identity operator. In our setting, for the PDE solution at timestep $t$ $u_{t}\in L^{2}\left(\Omega;\mathbb{R}^{V}\right)$, $\mathcal{P}[u_{t}]$ is the part of the solution that we observe and $\mathcal{Q}[u_{t}]$ is the unobserved part. Thus, the initial information we receive for the system is $\mathcal{P}[u_{0}]$. Applying $\mathcal{P}$ and $\mathcal{Q}$ to Equation [1](#S3.E1) and using $u=\mathcal{P}[u]+\mathcal{Q}[u]$, we get:

$$ $\displaystyle\frac{\partial}{\partial t}\mathcal{P}[u](t,x)$ $\displaystyle=\mathcal{P}\mathcal{L}[u](t,x)=\mathcal{P}\mathcal{L}\mathcal{P}[u](t,x)+\mathcal{P}\mathcal{L}\mathcal{Q}[u](t,x)$ (4) $\displaystyle\frac{\partial}{\partial t}\mathcal{Q}[u](t,x)$ $\displaystyle=\mathcal{Q}\mathcal{L}[u](t,x)=\mathcal{Q}\mathcal{L}\mathcal{P}[u](t,x)+\mathcal{Q}\mathcal{L}\mathcal{Q}[u](t,x)$ (5) $$

Solving for [5](#S3.E5) yields $\mathcal{Q}[u](t,x)=\int_{0}^{t}\exp{\mathcal{Q}\mathcal{L}(t-s)}\mathcal{Q}\mathcal{L}\mathcal{P}[u](s,x)ds+e^{\mathcal{Q}\mathcal{L}t}\mathcal{Q}[u_{0}](t,x).$

Plugging into [4](#S3.E4), we obtain a *Generalized Langevin Equation* (Mori, ) for $\mathcal{P}[u]$:

$$ $\displaystyle\frac{\partial}{\partial t}\mathcal{P}[u](t,x)$ $\displaystyle=\mathcal{P}\mathcal{L}\mathcal{P}[u](t,x)+\mathcal{P}\mathcal{L}\int_{0}^{t}\exp{\mathcal{Q}\mathcal{L}(t-s)}\mathcal{Q}\mathcal{L}\mathcal{P}[u](s,x)ds+\mathcal{P}\mathcal{L}e^{\mathcal{Q}\mathcal{L}t}\mathcal{Q}[u_{0}](t,x)$ (6) $$

We will refer to the first summand on the right hand side of [6](#S3.E6) as the Markovian term because it only depends on $\mathcal{P}[u](t,x)$, the second summand as the memory term because it depends on $\mathcal{P}[u](s,x)$ for $0\leq s\leq t$, and the third summand as the unobserved residual as it depends on $\mathcal{Q}[u_{0}]$ which is never observed.

We note that Equation [6](#S3.E6) is exact, not an approximation, so it is equivalent to solving the full system. Typically, the term that is most difficult to compute is the exponential of the memory term, and thus several methods to approximate it have been proposed.
In the physics literature, the memory term has been approximated through a perturbation expansion of the exponential (Breuer and Petruccione, ), or by approximating the operator $\exp{\mathcal{Q}\mathcal{L}(t-s)}:L^{2}\left(\Omega;\mathbb{R}^{V}\right)\rightarrow L^{2}\left(\Omega;\mathbb{R}^{V}\right)$ through operators defined in $\mathcal{P}\left[L^{2}\left(\Omega;\mathbb{R}^{V}\right)\right]$ (Shi and Geva, ; Zhang et al., ; Montoya-Castillo and Reichman, ; Kelly et al., ). In the machine learning literature, Ma et al. () develop the equations for the case when the operator $\mathcal{P}$ kept only the top-k modes, and designed a hybrid approach where the memory term was approximated with an LSTM (Hochreiter and Schmidhuber, ), and then used as an additional input of a numerical solver. In this work, we treat the whole memory term as an operator $\mathcal{M}:\mathcal{C}\left([0,T],\mathcal{P}\left[L^{2}\left(\Omega;\mathbb{R}^{V}\right)\right]\right)\rightarrow\mathcal{P}\left[L^{2}\left(\Omega;\mathbb{R}^{V}\right)\right]$(^1^11 Here $\mathcal{C}\left(A,B\right)$ denotes the space of continuous functions $u:A\rightarrow B$.) to be learnt by a parametrized sequential layer of a Neural Operator.

## 4 Our approach

### 4.1 Training procedure

First, we describe the high-level training scaffolding for our method, namely the way the data is generated, and the loss we use.

Training data: Let $u\in\mathcal{C}\left([0,T];L^{2}\left(\Omega;\mathbb{R}^{V}\right)\right)$ be the solution of the PDE given by Definition [2](#Thmdefi2). Let $\mathcal{S}$ be an equispaced grid in $\Omega$ with resolution $f$, and let $\mathcal{T}$ be another equispaced grid in $[0,T]$ with $N_{t}+1$ points. Given $u_{0}(x)\hskip 1.0pt\rule[-2.15277pt]{0.43057pt}{5.38193pt}\hskip 1.0pt_{\mathcal{S}}$, our goal is to predict $u(t,x)\hskip 1.0pt\rule[-2.15277pt]{0.43057pt}{5.38193pt}\hskip 1.0pt_{\mathcal{S}}$ for $t\in\mathcal{T}$ using a Neural Operator.

Training loss: As it is standard, we proceed through empirical risk minimization on a dataset of trajectories. More specifically, given a loss function $\ell:(\mathbb{R}^{|\mathcal{S}|},\mathbb{R}^{|\mathcal{S}|})\rightarrow\mathbb{R}$, a dataset of training trajectories $\left(u(t,x)^{(i)}\right)_{i=0}^{N}$, and parametrized maps $\mathcal{G}^{\Theta}_{t}:\mathbb{R}^{|\mathcal{S}|}\rightarrow\mathbb{R}^{|\mathcal{S}|}$ for $t\in\mathcal{T}$, we define:

$$ $\displaystyle\Theta^{*}=\min_{\Theta}\frac{1}{N}\sum_{i=0}^{N-1}\frac{1}{N_{t}}\sum_{t=1}^{N_{t}}\ell\left(u(t,x)\hskip 1.0pt\rule[-2.15277pt]{0.43057pt}{5.38193pt}\hskip 1.0pt_{\mathcal{S}},\mathcal{G}^{\Theta}_{t}\left[u_{0}(x)\hskip 1.0pt\rule[-2.15277pt]{0.43057pt}{5.38193pt}\hskip 1.0pt_{\mathcal{S}}\right]\right)$ $$

We then aim to find an adequate architecture choice such that $\mathcal{G}^{\Theta^{*}}$ has low test error on unseen trajectories of the same PDE.

### 4.2 The Architecture: Memory Neural Operator

In this section we describe Memory Neural Operator (MemNO), a Deep Learning framework to incorporate memory into Neural Operators.
Let $\text{NO}^{\Theta}_{t}$ be a Neural Operator with $L$ layers,
and denote $\text{NO}^{\Theta}_{t}[u_{0}]$ the prediction of the solution of the PDE at time $t$.
We will assume that this Neural Operator follows the Markovian assumption, i.e. we can write:

$$ $\text{NO}^{\Theta}_{t_{i+1}}[u_{0}]=r_{\text{out}}\circ\ell_{L}\circ\ell_{L-1}\circ...\circ\ell_{0}\circ r_{\text{in}}[\text{NO}^{\Theta}_{t_{i}}[u_{0}]]$ (7) $$

Where $r_{\text{in}}:\mathbb{R}^{|\mathcal{S}|}\rightarrow\mathbb{R}^{|\mathcal{S}|\times h_{0}}$ and $r_{\text{out}}:\mathbb{R}^{|\mathcal{S}|\times h_{L+1}}\rightarrow\mathbb{R}^{|\mathcal{S}|}$ are projector operators; $\ell_{j}:\mathbb{R}^{|\mathcal{S}|\times h_{j}}\rightarrow\mathbb{R}^{|\mathcal{S}|\times h_{j+1}}$ are parametrized layers; and $h_{j}$ is the dimension of the $j$-th hidden layer.

Our goal is to define a network $\mathcal{G}^{\Theta}_{t}$
that builds upon $\text{NO}^{\Theta}_{t}$ and can incorporate memory. For this
we
take inspiration from the Mori-Zwanzig theory exposed in Section [3.2](#S3.SS2).
Comparing Equation [7](#S4.E7) with Equation [6](#S3.E6),
we identify $\ell_{L}\circ\ell_{L-1}\circ...\circ\ell_{0}$ with the Markov term
which models the spatial dynamics.
To introduce the memory term, we interleave an additional residual sequential layer $\mathcal{M}$ that acts on hidden representations of the solution at previous timesteps. Concretely, the MemNO architecture can be written as:

$$ $\mathcal{G}_{t_{i+1}}^{\Theta}[u_{0}]=\mathcal{R}_{\text{out}}\circ\mathcal{L}_{L}\circ...\circ\mathcal{L}_{k+1}\circ\mathcal{M}\circ\mathcal{L}_{k}\circ...\circ\mathcal{L}_{0}\circ\mathcal{R}_{\text{in}}\left[\mathcal{G}_{t_{i}}^{\Theta}[u_{0}],\mathcal{G}_{t_{i-1}}^{\Theta}[u_{0}],...,u_{0}\right]$ $$

Where $-1\leq k\leq L$ is a chosen hyperparameter.(^2^22 $k=L$ referst to inserting $M$ after all the $S$ layers, and $k=-1$ refers to inserting $M$ as the first layer. As we show in Appendix [F.2](#A6.SS2), our experiments are not very sensitive to the choice of $k$.)
Now, the spatial $\mathcal{L}_{j}$ layers are understood to be applied timestep-wise.
That is, if $v^{(j)}(t^{\prime})$ is the hidden representation at the $j$ layer for a timestep $t^{\prime}\leq t_{i}$, then $\mathcal{L}_{j+1}\left[v^{(j)}(t_{i}),...,v^{(j)}(t_{0})\right]=\left[\ell_{j}[v_{i}^{(j)}(t_{i})],...,\ell_{j}[v_{0}^{(j)}(t_{0})]\right]$,
and analogously for $\mathcal{R}_{\text{in}}$ and $\mathcal{R}_{\text{out}}$.
Thus, the $\mathcal{L}_{j}$ layers still follow the Markovian assumption. The memory is introduced through $\mathcal{M}$, which consists of a sequential layer $m$ that is applied to the time history of the hidden representation of the $k$-th layer, that is $\mathcal{M}:\mathbb{R}^{i\times|\mathcal{S}|\times h_{k}}\longrightarrow\mathbb{R}^{|\mathcal{S}|\times h_{k}}$ with $\left(\mathcal{M}[v^{(k)}(t_{i}),...,v^{(k)}(t_{0})]\right)_{sh}=m\left[v_{sh}^{(k)}(t_{i}),...,v_{sh}^{(k)}(t_{0})\right]$. Note that $m$ is the same for each of the $|\mathcal{S}|\times h_{k}$ elements of the hidden layer.
The main motivation of our MemNO framework is that it can be utilized with *any*
existing neural operator layer $\ell$, and with any
(causal) sequential model $\mathcal{M}$.
Thus it provides a modular architecture design
which we hope can serve as a useful tool for practitioners.

## 5 Theoretical motivation for memory: a simple example

In this section, we provide a simple, but natural example of a (linear) PDE, along with (in the nomenclature of Section [3.2](#S3.SS2)) a natural projection operator given by a *Fourier truncation measurement 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,2\pi]$ which have a convenient basis:

###### Definition 5 (Basis for 2 $\pi$ -periodic functions) .

A function $f:\mathbb{R}\to\mathbb{R}$ is $2\pi$-periodic if $f(x+2\pi)=f(x)$. We can identify $2\pi$-periodic functions with functions over the torus $T:=\{e^{i\theta}:\theta\in\mathbb{R}\}\subseteq\mathbb{C}$ by the map $\tilde{f}(e^{ix})=f(x)$. Note that $\{e^{ixn}\}_{n\in\mathbb{Z}}$ is a basis for the set of $2\pi$-periodic functions.

We will define the following measurement operator:

###### Definition 6 (Fourier truncation measurement) .

The operator $\mathcal{P}_{k}:L^{2}(T;\mathbb{R})\to L^{2}(T;\mathbb{R})$ acts on $f\in L^{2}(T;\mathbb{R})$, $f(x)=\sum_{n=-\infty}^{\infty}a_{n}e^{inx}$ as $\mathcal{P}_{k}(f)=\sum_{n=-k}^{k}a_{n}e^{inx}.$

We will also define for notational convenience the functions $\{\mathbf{e}_{n}\}_{n\in\mathbb{Z}}$, where $\mathbf{e}_{n}(x):=e^{-inx}+e^{inx}.$
Now we consider the following operator to define a linear time-dependent PDE:

###### Proposition 1 .

Let $\mathcal{L}:L^{2}(T;\mathbb{R})\to L^{2}(T;\mathbb{R})$ be defined as
$\mathcal{L}u(x)=-\Delta u(x)+B\cdot(e^{-ix}+e^{ix})u(x)$
for $B>0$. Then, we have:

$$ $\displaystyle\forall 1\leq n\in\mathbb{N},\hskip 5.0pt\mathcal{L}(\mathbf{e}_{n})$ $\displaystyle=n^{2}\mathbf{e}_{n}+B(\mathbf{e}_{n-1}+\mathbf{e}_{n+1})\quad\&\quad\mathcal{L}(\mathbf{e}_{0})=2B\mathbf{e}_{1}$ $$

The crucial property of this operator is that it acts by “mixing” the $n$-th Fourier basis with the $(n-1)$-th and $(n+1)$-th: thus information is propagated to both the higher and lower-order part of the spectrum.
Given the above proposition, we can easily write down the evolution of a PDE with operator $\mathcal{L}$ in the basis $\{\mathbf{e}_{n}\}_{n\in\mathbb{Z}}$:

###### Proposition 2 .

Let $\mathcal{L}$ be defined as in Proposition [1](#A7.EGx5). Consider the PDE

$$ $\displaystyle\frac{\partial}{\partial t}u(t,x)$ $\displaystyle=\mathcal{L}u(t,x)$ $\displaystyle u(0,x)$ $\displaystyle=\sum_{n\in\mathbb{N}_{0}}a_{n}(0)\mathbf{e}_{n}$ $$

Let $u(t,x)=\sum_{n\in\mathbb{N}_{0}}a_{n}^{(t)}\mathbf{e}_{n}$.
Then, the coefficients $a_{n}^{(t)}$ satisfy:

$$ $\displaystyle\forall 1\leq n\in\mathbb{N},\hskip 5.0pt\frac{\partial}{\partial t}a_{n}^{(t)}$ $\displaystyle=n^{2}a_{n}^{(t)}+B\left(a_{n-1}^{(t)}+a_{n+1}^{(t)}\right)$ (8) $\displaystyle\frac{\partial}{\partial t}a_{0}^{(t)}$ $\displaystyle=2Ba_{1}^{(t)}$ (9) $$

With this setup in mind, we will show that as $B$ grows, the memory term in Equation [6](#S3.E6) can have an arbitrarily large effect on the calculated solution:

###### Theorem 1 (Effect of memory) .

Consider the Fourier truncation operator $\mathcal{P}_{1}$ and let $\mathcal{Q}=I-\mathcal{P}_{1}$. Let $u(0,x)$ have the form in Proposition [2](#A7.EGx7) for $B>0$ sufficiently large, and let $a^{(0)}_{n}>0,\forall n>0.$ Consider the memoryless and memory-augmented PDEs:

$$ $\displaystyle\frac{\partial u_{1}}{\partial t}$ $\displaystyle=\mathcal{P}_{1}\mathcal{L}u_{1}$ (10) $\displaystyle\frac{\partial u_{2}}{\partial t}$ $\displaystyle=\mathcal{P}_{1}\mathcal{L}u_{2}+\mathcal{P}_{1}\mathcal{L}\int_{0}^{t}\exp{\mathcal{Q}\mathcal{L}(t-s)}\mathcal{Q}\mathcal{L}u_{2}(s)ds$ (11) $$

with $u_{1}(0,x)=u_{2}(0,x)=\mathcal{P}_{1}u(0,x)$. Then, $u_{1}$ and $u_{2}$ satisfy:

$$ $\displaystyle\forall t>0,\|u_{1}(t)-u_{2}(t)\|_{L_{2}}$ $\displaystyle\gtrsim Bt\|u_{1}(t)\|_{L_{2}}$ (12) $\displaystyle\forall t>0,\|u_{1}(t)-u_{2}(t)\|_{L_{2}}$ $\displaystyle\gtrsim Bt\exp(\sqrt{2}Bt)$ (13) $$

###### Remark 1 .

Note that the two conclusions of the theorem mean that both the absolute difference, and the relative difference between the PDE including the memory term Equation [11](#S5.E11) and not including the memory term Equation [10](#S5.E10) can be arbitrarily large as $B,t\to\infty$.

###### Remark 2 .

The choice of $\mathcal{L}$ is made for ease of calculation of the Markov and memory term. Conceptually, we expect the solution to Equation [11](#S5.E11) will differ a lot from the solution to Equation [10](#S5.E10) if the action of the operator $\mathcal{L}$ tends to “mix” components in the span of $\mathcal{P}$ and the span of $\mathcal{Q}$.

###### Remark 3 .

If we solve the equation $\frac{\partial}{\partial t}u(t,x)=\mathcal{L}u(t,x)$ *exactly*, we can calculate that $\|u(t)\|_{L_{2}}$ will be on the order of $\exp(2Bt)$. This can be seen by writing the evolution of the coefficients of $u(t)$ in the basis $\{\mathbf{e}_{n}\}$, which looks like:
$\frac{\partial}{\partial t}\begin{pmatrix}a_{0}\\
a_{1}\\
\dots\end{pmatrix}=\mathcal{O}\begin{pmatrix}a_{0}\\
a_{1}\\
\dots\end{pmatrix}$
where $\mathcal{O}$ is roughly a tridiagonal Toeplitz operator $\mathcal{O}=\begin{pmatrix}&\vdots&\vdots&\vdots&\vdots&\\
\dots&B&n^{2}&B&0&\dots\\
\dots&0&B&(n+1)^{2}&B&\dots\\
&\vdots&\vdots&\vdots&\vdots&\end{pmatrix}.$
The largest eigenvalue of this operator can be shown to be on the order of at least $2B$ (equation (4) in (Noschese et al., )). The Markov term results in a solution of order $\exp(\sqrt{2}Bt)$ ( Equation [18](#A7.E18),Equation [19](#A7.E19)), which is multiplicatively smaller by a factor of $\exp(\left(2-\sqrt{2}\right)Bt)$. The result in this Theorem shows the memory-based PDE Equation [11](#S5.E11) results in a multiplicative “first order” correction which can be seen by Taylor expanding $\exp(\sqrt{2}Bt)\approx 1+\sqrt{2}Bt+\frac{1}{2}(\sqrt{2}B)^{2}t^{2}+\dots$.

###### Definition 5 (Basis for 2 π 𝜋 \pi -periodic functions) .

###### Definition 6 (Fourier truncation measurement) .

###### Proposition 1 .

###### Proposition 2 .

###### Theorem 1 (Effect of memory) .

###### Remark 1 .

###### Remark 2 .

###### Remark 3 .

## 6 Memory helps with low-resolution data and input noise: a case study

In this section we present a case study for several common PDEs of practical interest, showing that MemNO brings accuracy benefits when the data is supplied in low resolution.
Through our experiments we show
the difference in the performance between
a baseline “memoryless” architecture, which we choose to be Factorized Fourier Neural Operator (FFNO) (Tran et al., )
and a memory-augmented architecture using S4 (Gu et al., ), which we denote as the S4-Factorized Fourier Neural Operator (s4FFNO).
The architectural details for both the
architectures are elaborated upon in
Appendix [B](#A2).

### 6.1 Setup: Training and evaluation procedure

To construct our datasets, we first produce discretized trajectories of a PDE for $N_{t}$ timesteps, i.e. $(u(t))_{t=0}^{N_{t}}$ in a *high resolution* discretized spatial grid $\mathcal{S}^{HR}\subset\mathbb{R}^{d}$, i.e. $u(t)\in\mathbb{R}^{|\mathcal{S}^{HR}|}$.
We then produce datasets that consist of *lower resolution* versions of the above trajectories, i.e. on a grid $\mathcal{S}^{LR}$ of lower resolution $f$. For 1-dimensional datasets, the discretized trajectory on $\mathcal{S}^{LR}$ is obtained by cubic interpolation of the trajectory in the high resolution grid. In 2D, the discretized trajectory is obtained by downsampling.
We will show results in different resolutions, in which case both train and test trajectories are at such resolution, and the loss function is also computed at the chosen resolution.
Our training loss and evaluation metric is *normalized Root Mean Squared Error (nRMSE)*:

$$ $\text{nRMSE}\left(u(t,x)\hskip 1.0pt\rule[-2.15277pt]{0.43057pt}{5.38193pt}\hskip 1.0pt_{\mathcal{S}^{LR}},\hat{u}(t)\right)=\frac{\|u(t,x)\hskip 1.0pt\rule[-2.15277pt]{0.43057pt}{5.38193pt}\hskip 1.0pt_{\mathcal{S}^{LR}}-\hat{u}(t)\|_{2}}{\|u(t,x)\hskip 1.0pt\rule[-2.15277pt]{0.43057pt}{5.38193pt}\hskip 1.0pt_{\mathcal{S}^{LR}}\|_{2}},$ $$

where $\parallel\cdot\parallel_{2}$ is the euclidean norm in $\mathbb{R}^{|\mathcal{S}^{HR}|}$. More details on training are given in appendix [E](#A5).

### 6.2 Kuramoto–Sivashinsky equation (1D): a study in low-resolution

The Kuramoto-Sivashinsky equation (KS) is a nonlinear PDE that is used as a 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 a viscosity $\nu$, it is written as $u_{t}+uu_{x}+u_{xx}+\nu u_{xxxx}=0.$
We generated datasets for KS at different viscosities and resolutions. The results are shown in Table [1](#S6.T1). We can see s4FFNO outperforms FFNO across these viscosities and resolutions, having an nRMSE that can be more than six times smaller.
We also note that, since the memory layer is applied element-wise in the time dimension, it has very few parameters compared to the spatial layers, as seen in the difference of parameters between s4FFNO and FFNO in column 3.

Figure: (a) $\nu=0.075$
Refer to caption: /html/2409.02313/assets/Images/KSres_0.075.png

**Table 1: nRMSE values at different resolutions for KS with different viscosities. The final time is 5 seconds and the trajectories contain 25 timesteps.**
| Architecture | Resolution | Parameters | nRMSE $\downarrow$ |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
| $\nu=0.075$ | $\nu=0.1$ | $\nu=0.125$ | $\nu=0.15$ |  |  |  |
| FFNO | 32 | 2.79M | 0.650 | 0.376 | 0.254 | 0.187 |
| FFNO | 48 | 3.84M | 0.330 | 0.196 | 0.131 | 0.068 |
| FFNO | 64 | 4.89M | 0.179 | 0.024 | 0.006 | 0.003 |
| s4FFNO | 32 | 2.85M | 0.233 | 0.079 | 0.042 | 0.028 |
| s4FFNO | 48 | 3.90M | 0.071 | 0.035 | 0.020 | 0.010 |
| s4FFNO | 64 | 4.94M | 0.049 | 0.011 | 0.005 | 0.003 |

The key factor for the improved performance of MemNO over memoryless Neural Operators in not the absolute resolution, but rather the *resolution relative to the frequency spectrum of the solution*. The lower the viscosity, the higher the frequencies that appear in the spectrum. This can be clearly seen in Figure [1](#S6.F1): in the top row, as viscosities increase, the resolution at which there is a significant difference between s4FFNO and FFNO decreases. In the second row of the figure, we show a visualization of the frequency spectrum of a solution at those viscosities.

We note that even if the initial condition does not contain high frequencies, in the KS equation high frequencies will appear as the system evolves—indeed, this dataset was generated with initial conditions whose maximum Fourier mode was 8. This is in qualitative agreement with the theoretical motivation in Section [5](#S5)—albeit the KS equation is substantially more complicated since it is nonlinear, so it is hard to fully theoretically analyze.
We provide a similar study on 1D Burgers equation
in the Appendix [C](#A3).

### 6.3 Navier Stokes equation (2D): study in observation noise

The Navier Stokes equation describes the motion of a viscous fluid.
Like in Li et al. (), we consider the incompressible form in the 2D unit torus, which is given by:

$$ $\displaystyle\frac{\partial w(x,t)}{\partial t}+u(x,t)\cdot\nabla w(x,t)$ $\displaystyle=\nu\Delta w(x,t)+f(x),$ $\displaystyle x\in(0,1)^{2},t\in(0,T]$ $\displaystyle\nabla\cdot u(x,t)$ $\displaystyle=0,$ $\displaystyle x\in(0,1)^{2},t\in[0,T]$ $\displaystyle w(x,0)$ $\displaystyle=w_{0}(x),$ $\displaystyle x\in(0,1)^{2}$ $$

Where $w=\nabla\times u$ is the vorticity, $w_{0}\in L^{2}((0,1)^{2};\mathbb{R})$ is the initial vorticity, $\nu\in\mathbb{R}_{+}$ is the viscosity coefficient, and $f\in L^{2}((0,1)^{2};\mathbb{R})$ 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 and with a Neural Operator.
We investigate the effect of memory when adding IID Gaussian noise to the inputs of our neural networks. Noise would represent the observation noise arising from the intrisic error of the measurement device. The noise $\epsilon\in\mathbb{R}^{|\mathcal{T}|\times|\mathcal{S}^{LR}|}$ is sampled IID from a Gaussian distribution $\epsilon_{ts}\sim\mathcal{N}(0,\sigma)$, and then added to training and test inputs. During training, for each trajectory a different noise with the same $\sigma$ is sampled at each iteration of the optimization algorithm. The targets in training and testing represent our ground truth and are not added noise.

Figure: (d) $\nu=10^{-3}$, $T=16s$, $N_{t}=32$
Refer to caption: /html/2409.02313/assets/Images/NS3-noise.png

In Figure [2](#S6.F2), we show the results for $\nu=10^{-3}$ when adding noise levels from $\sigma=0.0$ (no noise) to $\sigma=2.048$. s4FFNO-2D 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-2D (note the logarithmic scale). For this viscosity, adding small levels of noise actually helps training, which was also observed in other settings in Tran et al. (). Figure [2](#S6.F2) shows the same experiment performed with $\nu=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.

### 6.4 Relationship with fraction of unobserved information

In this section, we provide a simple experiment to quantify the
effect of the fraction of unobserved information on the performance of memory based models.
Given a grid of resolution $f$, we define the Fourier truncation measurement $\mathcal{P}_{\lfloor\frac{f}{2}\rfloor}$ as in Section [5](#S5),
which simply keeps the top $\lfloor\frac{f}{2}\rfloor+1$ modes and discards the other high frequency modes.
Assume $u(t)\in L^{2}(\Omega;\mathbb{R}^{V})$ is the solution of a 1-dimensional PDE at time $t$, and $a_{n}^{(t)}$ for $n\in\mathbb{Z}$ is its Fourier Transform. We define the quantity:

$$ $\displaystyle\omega_{f}=\frac{1}{N_{t}}\sum_{i=1}^{N_{t}}\frac{\sum_{|n|\geq\frac{f}{2}}|a_{n}^{(t_{i})}|^{2}}{\sum_{n\in\mathbb{Z}}|a_{n}^{(t_{i})}|^{2}}$ (14) $$

$\omega_{f}$ is approximate indicator of the
amount of information that is lost when the solution of the PDE is observed at resolution $f$ across time.
We show that there is a positive correlation between $\omega_{f}$ and the difference in nRMSE between FFNO and s4FFNO for the KS experiment in Figure [3](#S6.F3), and also the for Burgers’ experiments of Appendix[C](#A3) in Figure [5](#A3.F5).
This demonstrates the benefits
of memory as a way to compensate missing information in the observation.

Figure: Figure 3: Values of $\omega_{f}$ and the difference in nRMSE between FFNO and s4FFNO for different resolutions in the KS experiment of Section [6.2](#S6.SS2) with $\nu=0.1$. $\omega_{f}$ is averaged over all trajectories in the dataset. The value is computed approximating the continuous Fourier modes with Discrete Fourier modes of the solution in the highest resolution available (512 for KS).
Refer to caption: /html/2409.02313/assets/Images/KS-omega.png

## 7 Conclusion and Future Work

We study the benefits of maintaining memeory while modeling
time dependent PDE systems. Taking inspiration from the Mori-Zwanzig formulation, we show that when
we only observe part of the PDE initial condition (for example, PDEs observed on low-resolution or with
input noise), the system is no longer Markovian, and the dynamics depend on a
*memory term*.
To this end, we introduce
MemNO, an architecture that combines Fourier Neural Operator (FNO) and
the S4 architecture.
Through our experiments on different 1D and 2D PDEs, we show that the MemNO
architecture outperforms the memoryless baselines.

We present several avenues for future work. First, our experiments on
observation noise are limited to the setting where the input noise is IID.
Further, extending the experiments and observing the effects of memory in
more real-world settings (for example, with non-IID noise or in the presence of aliasing) is a fertile ground for future work, and also necessary to ensure that the application of this method does not have unintended negative consequences when broadly applied in society.
Lastly, while we limit our
study of the effects of memory to FNO based architectures, performing
similar studies for different architectures like Transformer based neural operators (Hao et al., )
and diffusion based operators (Lippe et al., )
is an interesting direction for future work.

## Acknowledgements

RBR is supported by the “la Caixa” Foundation (ID 100010434). The fellowship code is LCF/BQ/EU22/11930090.
TM is supported in part by CMU Software Engineering Institute via Department of Defense under contract FA8702-15-D-0002.
AR is supported in part by
NSF awards IIS-2211907, CCF-2238523, and Amazon Research.
The authors also thank Cartesia AI for their generous provison of computational resources.

## Appendix A Extended Related Work

Neural Operators. The Fourier Neural Operator (FNO) is a Neural Operator that performs a transformation in the frequency space of the input [Li et al., ]. Other models have proposed different inductive biases for Neural Operators, including physics based losses and constraints [Li et al., ],
using Deep Equilibrium Model (DEQ) Bai et al. []
to design specialized
architectures for steady-state (time-independent)
PDEs Marwah et al. [], and using local message passing
Graph Neural Networks (GNNs) [Gilmer et al., , Kipf and Welling, ]
based encoders to model irregular geometries [Li et al., , ].
Other methodologies to solve for
PDEs include methods like  Gupta and Brandstetter [], Rahman et al. []
that use the U-Net [Ronneberger et al., ]
type architectures
and works like [Cao, , Hao et al., ]
that introduce different Transformer Vaswani et al. [] based
neural solution operators for modeling both time-dependent and time-independent PDEs.
While most of these methodology
are designed for time-dependent PDEs,
there is no clear consensus of how to model
the past-states to predict future states, and most
of these methods predict the PDE states over time
in an auto-regressive way by conditioning
the model on varying lengths of the past states Li et al. [], Tran et al. [], Hao et al. [].

Foundation models. Lately, there have been community efforts
towards
creating large scale
foundational models for modeling multiple PDE families [McCabe et al., , Hao et al., , Shen et al., ], and weather prediction Pathak et al. [], Lam et al. [].
We hope that our study is useful in informing the architectural
design of future models.

## Appendix B Network Architectures

Factorized Fourier Neural Operator (FFNO) (Tran et al. []): This model is a refinement over the original Fourier Neural Operator (Li et al. []). Given a hidden dimension $h$ and a spatial grid $\mathcal{S}$, its layers $\ell:\mathbb{R}^{|\mathcal{S}|\times h}\rightarrow\mathbb{R}^{|\mathcal{S}|\times h}$ are defined as:

$$ $\ell(v)=v+\text{Linear}_{h,h^{\prime}}\circ\sigma\circ\text{Linear}_{h^{\prime},h}\circ\mathcal{K}[v]$ (15) $$

where $\sigma$ is the GeLU activation function [Hendrycks and Gimpel, ] and $h^{\prime}$ is an expanded hidden dimension. $\mathcal{K}$ is a kernel integral operator that performs a linear transformation in the frequency space. Denoting by
$\text{FFT}_{\alpha}$, $\text{IFFT}_{\alpha}$ are to the Discrete Fast Fourier Transform and the Discrete Inverse Fast Fourier Transform along dimension $\alpha$ [Cooley et al., ], it can be written as:

$$ $\mathcal{K}[v]=\sum_{\alpha\in\{1,...,d\}}\text{IFFT}[R_{\alpha}\cdot\text{FFT}_{\alpha}[v]]$ $$

for learnable matrices of weights $R_{\alpha}\in\mathbb{C}^{h^{2}\times k_{\text{max}}}$. $k_{\text{max}}$ is the maximum number of Fourier modes which are used in $\mathcal{K}$. We use all Fourier modes by setting $k_{\text{max}}=\lfloor\frac{f}{2}\rfloor$.

In our experiments, The FFNO model consists of 4 FFNO layers. For experiments in 1D, the hidden dimensions are all 128 ($h_{j}=128$ for $j=0,1,2,3$) and the expanded hidden dimension of FFNO’s MLP $h^{\prime}$ is $4\cdot 128$. For experiments in 2D, the hidden dimensions are all 64 and the expanded hidden dimension is $4\cdot 64$.

S4 - Factorized Fourier Neural Operator (s4FFNO): This model uses our MemNO framework. To discern the effect of memory, all layers except the memory layer will be the same as FFNO. For the memory layer, we choose an S4 layer Gu et al. [] with a state dimension of 64 and a diagonal S4 (S4D) kernel.(^3^33The S4 repository has two available kernels, the diagonal S4 (S4D) and the Normal Plus Low Rank S4 (S4NPLR). In our experiments, we didn’t find a significant difference between the two, and chose S4D for simplicity.)

For all our models, we use a simple spatial positional encoding $E$. In 1D, if the grid has $f$ equispaced points in $[0,L]$, then $E\in\mathbb{R}^{f}$ and the positional encoding is defined as $E_{i}=\frac{i}{L}$ for $0\leq i\leq f-1$. In 2D, if we have a $f\times f$ points 2D equispaced grid in $[0,L_{x}]\times[0,L_{y}]$, the positional encoding is defined as $E_{ij}=(\frac{i}{L_{x}},\frac{j}{L_{y}})$. In the input lifting operator $\mathcal{R}_{\text{in}}$, the input and the grid are stacked and a Linear layer from $\mathbb{R}^{2}\rightarrow\mathbb{R}^{h_{0}}$ is applied element-wise. For the decoder $\mathcal{R}_{\text{out}}$, we use another Linear layer (without positional encoding).

## Appendix C Burgers’ Equation (1D): A study on low-resolution

The Burgers’ equation with viscosity $\nu$ is a nonlinear PDE used as a 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. It is known for exhibiting a 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:

$$ $u_{t}+uu_{x}=\nu u_{xx}$ $$

We used the publicly available dataset of the Burgers’ equation in the PDEBench repository (Takamoto et al. []) with viscosity $e^{-3}$, which is available at resolution 1024. We compare our models at resolutions 64, 128, 256, 512 and 1024 and show the results at figure [4(a)](#A3.F4.sf1). As in the case of KS, s4FFNO outperforms FFNO, especially at low resolutions. Furthermore, we show the differnece in nRMSE at each timestep in figure [4(b)](#A3.F4.sf2). We observe that at the first timestep there is no difference between the two models - this makes sense 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.

Figure: (a) nRMSE of FFNO and s4FFNO at different resolutions.
Refer to caption: /html/2409.02313/assets/Images/Burgers-resolution.png

### C.1 Relationship with fraction of unobserved information

As mentioned in Section [6.4](#S6.SS4)
we measure the correlation of
$\omega_{f}$ defined in Equation [14](#S6.E14) with the difference in the nRMSE
between FFNO and s4FFNO. The results can be seen in Figure [5](#A3.F5)

Figure: Figure 5: Values of $\omega_{f}$ (defined in Equation [14](#S6.E14)) and the difference in nRMSE between FFNO and s4FFNO for different resolutions in two of our experiments for Burgers’ Equation. $\omega_{f}$ is averaged over all trajectories in the dataset. The value is computed approximating the continuous Fourier modes with the Discrete Fourier modes of the solution in the highest resolution available (1024 for Burgers’ Equation).
Refer to caption: /html/2409.02313/assets/Images/Burgers-omega.png

## Appendix D Data Generation

### D.1 Kuramoto–Sivashinsky equation

For Kuramoto-Sivashinsky (KS) equation is given by:

$$ $\displaystyle u_{t}+uu_{x}+u_{xx}+\nu u_{xxxx}$ $\displaystyle=0\quad(t,x)\in[0,T]\times[0,L]$ $\displaystyle u(0,x)$ $\displaystyle=u_{0}(x)\quad x\in[0,L]$ $$

We use periodic boundary conditions. Our data generation method is very similar to the one of Brandstetter et al. []. We employ the method of lines [Schiesser, ], where the spatial dimension is discretized, and the PDE is transformed to a systems of Ordinary Differential Equations (ODEs), one per point in the grid. In order to compute the spatial derivative of the solution at each point in the grid, a pseudospectral method is used, where derivatives are computed in frequency space and then converted to the original space through a Fast Fourier Transform. This method is implemented in the diff method of the scipy.fftpack package [Virtanen et al., ]. Similarly, the system of ODEs is solved numerically with a implicit Runge-Kutta method of the Radau IIA family of order 5 [Hairer and Wanner, ], which is implemented in the solve_ivp method of scipy.integrate. We refer to the code provided in Brandstetter et al. [] to reproduce this data generation, however certain small modifications have to be made, like using a fixed $\Delta t$ per trajectory and increasing the number of modes in the initial condition.

As for the PDE parameters, we use $L=64$ and $T=5$. For the initial condition, we use a superposition of sinusodial waves:

$$ $u_{0}(x)=\sum_{i=0}^{20}A_{i}\sin\left(\frac{2\pi k_{i}}{L}x+\phi_{i}\right)$ $$

where for each trajectory, the $A_{i}$ are sampled from a continuous uniform in $[-0.5,0.5]$, the $k_{i}$ are sampled from a discrete uniform in $\{1,2,...,8\}$, and the $\phi_{i}$ are sampled from a uniform uniform in $[0,2\pi]$. We discretize $[0,T]$ into 26 equispaced points separated by $\Delta t=0.2$(^4^44 We reiterate that, as opposed to [Brandstetter et al., ], we don’t have a random $\Delta t$ per trajectory), keep the 25 temporal points generated by the numerical solver and discard our initial condition.
In the experiments in section [6.2](#S6.SS2), 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 [F.1](#A6.SS1), we generated one dataset with viscosity 0.15 in resolution 256, 4096 training samples and 256 test samples.

### D.2 Burgers’ 1D equation

The 1D Burgers’ equation can be written as:

$$ $u_{t}+uu_{x}=\nu u_{xx}\quad(t,x)\in[0,T]\times[0,L]$ $$

For the Burgers’ equation, we take the publicly available Burgers’ dataset of PDEBench [Takamoto et al., ] with viscosity $0.001$. Out of the 10000 samples of the dataset, we use a $10\%$ for testing. For training, we found it sufficient to use 2048 samples. Additionaly, 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 $u_{xx}$ attenuates all high frequencies and the solution changes very slowly.

### D.3 Navier Stokes 2D equation

The incompressible Navier Stokes equation in the 2D unit torus is given by:

$$ $\displaystyle\frac{\partial w(x,t)}{\partial t}+u(x,t)\cdot\nabla w(x,t)$ $\displaystyle=\nu\Delta w(x,t)+f(x),$ $\displaystyle x\in(0,1)^{2},t\in(0,T]$ $\displaystyle\nabla\cdot u(x,t)$ $\displaystyle=0,$ $\displaystyle x\in(0,1)^{2},t\in[0,T]$ $\displaystyle w(x,0)$ $\displaystyle=w_{0}(x),$ $\displaystyle x\in(0,1)^{2}$ $$

For the data generation, we follow the method of Li et al. [], yet with different temporal and spatial grids. The initial conditions $w_{0}$ are sampled from a Gaussian Random field $\mathcal{N}\left(0,7^{\frac{3}{2}}(-\Delta+49I)^{-2.5}\right)$ with periodic boundary conditions. The forcing term is $f(x_{1},x_{2})=0.1\left(\sin{2\pi(x_{1}+x_{2})}+\cos{2\pi(x_{1}+x_{2})}\right)$. At each timestep, the velocity is obtained from the vorticity 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 in time, with a timestep of $10^{-4}$. We use a 512x512 spatial grid which is then downsampled to 64x64 for our experiments. For the viscosity $\nu=10^{-3}$, we use a final time of 16 seconds and sample every 0.5 seconds. For the viscosity $\nu=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, we refer to Li et al. [].

## Appendix E Training details

In this section, we will provide a detailed description of the training hyperparameters used in the KS experiments of Section [6.2](#S6.SS2), in the Burgers experimente of section [C](#A3) and the Navier Stokes experiments of section [6.3](#S6.SS3). We start with the training hyperparameters. All our 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 and Hutter, ]; 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 were 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 we run a 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 a 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 MemNO model. In the KS and Burgers experiments, the maximum memory length was 25 (which is the same as the number of timesteps of the dataset). That means that for the last timestep, the previous 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 so on.

As in [Tran et al., ], 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-1$ previous steps was fed to the model (as opposed to the prediction of the model for such steps).

We run our experiments on A6000/A6000-Ada 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 1D required a lower GPU memory (less than 10GB) and each run normally took less than an hour.

## Appendix F Ablations on the Memory layer

In this section we present two ablations regarding the memory layer of MemNO.

### F.1 Ablation: Choice of sequential model

In section [4.2](#S4.SS2) 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 experiments. In this section, we explore two other candidates for the sequential layers: a transformer and an LSTM. We introduce Transformer-FFNO (T-FFNO) and LSTM-FFNO as two models that are identical to s4FFNO except in the sequential layer, where a transformer and an LSTM are used respectively. The LSTM model only has one layer and the transformer layer includes causal masking and a positional encoding. The positional encoding for $pos$ across the time dimension and $i$ across the hidden dimension is given by:

$$ $\displaystyle PE(pos,2i)$ $\displaystyle=\sin\left(\frac{pos}{10000^{\frac{2i}{\text{dim_model}}}}\right)$ $\displaystyle PE(pos,2i+1)$ $\displaystyle=\cos\left(\frac{pos}{10000^{\frac{2i}{\text{dim_model}}}}\right)$ $$

We show results for the KS dataset with viscosity $\nu=0.15$ and different resolutions. This dataset was generated using a resolution of 256 and contains 4096 samples, twice as many compared to the KS datasets of [D](#A4), given that transformers are known to perform better in high-data regimes. The results are shown in Figure [6](#A6.F6). TFFNO performs significantly worse than s4FFNO across almost all resolutions, and even performs worse than FFNO. In constrast, LSTM-FFNO outperforms FFNO, which shows that MemNO can work with other sequentials model apart from S4. The memory term in Equation [6](#S3.E6) is a convolution in time, which is equivalent to the S4 layer and very similar to a Recurrent Neural Network (RNN) style layer, as showed in Gu et al. []. 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 a bias for continuous signals and has empirically has proven better performance in these kind of tasks [Gu et al., ], which is in agreement with its increased performance over LSTMs in this experiment. Additionally, we observed that LSTMs were instable to train in Navier Stokes 2D datasets.

Lastly, we make two remarks. Firstly, we believe that transformers performed worse due to overfitting, given that the train losses were normally comparable or even smaller than the train losses of the rest of the models at each resolution. Modifications of the transformer model or to the training hyperparameters as in other works [Hao et al., , Cao, , Hao et al., ] might solve this issue. Secondly, recently there has been a surge of new sequential models such as Mamba [Gu and Dao, ], RWQK [Peng et al., ], xLSTM [Beck et al., ] or LRU [Orvieto et al., ]. We leave it as future work to 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.

Figure: Figure 6: Performance of FFNO, s4FFNO and T-FFNO and LSTM-FFNO in KS with viscosity $\nu=0.15$.
Refer to caption: /html/2409.02313/assets/Images/KS-resolution.png

### F.2 Ablation: memory layer configuration

In Section [4.2](#S4.SS2) we introduced the memory layer in MemNO as a single layer to be interleaved with Neural Operator layers. In our experiments, we inserted it after the second layer of a 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 a 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, SSTSS denotes the architecture of our experiments, were we have 2 Neural Operators layers, followed by a memory layer, followed by other 2 Neural Operator layers. Similarly, SSSST denotes 4 Neural Operators layers followed by a memory layer. In Table [2](#A6.T2), we present the results for the KS dataset with $\nu=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 SSTSS), 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, we 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.

**Table 2: KS, $\nu=0.1$. The final time is 4 seconds and the trajectories contain 20 timesteps. For each architecture, we tried 4 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 is also with respect to the seeds.**
| Architecture | nRMSE $\downarrow$ |  |  |
| --- | --- | --- | --- |
| Resolution 32 | Resolution 48 | Resolution 64 |  |
| s4FFNO (SSTSS) | 0.123 $\pm$ 0.011 | 0.086 $\pm$ 0.004 | 0.015 $\pm$ 0.001 |
| s4FFNO (SSSST) | 0.142 $\pm$ 0.009 | 0.069 $\pm$ 0.001 | 0.017 $\pm$ 0.001 |
| s4FFNO (STSSTS) | 0.141 $\pm$ 0.006 | 0.064 $\pm$ 0.002 | 0.019 $\pm$ 0.001 |
| s4FFNO (STSTSTST) | 0.113 $\pm$ 0.006 | 0.070 $\pm$ 0.004 | 0.017 $\pm$ 0.001 |
| s4FFNO (TSSSS) | 0.129 $\pm$ 0.007 | 0.080 $\pm$ 0.003 | 0.017 $\pm$ 0.001 |
| FFNO | 0.294 $\pm$ 0.004 | 0.138 $\pm$ 0.013 | 0.021 $\pm$ 0.002 |

## Appendix G Appendix: Quantifying the effect of memory

###### Proof.

We proceed to the Equation [10](#S5.E10) first. Note that $u_{1}(t),\forall t\geq 0$ can be written as $u_{1}(t)=a_{0}^{(t)}\mathbf{e}_{0}+a^{(t)}_{1}\mathbf{e}_{1}$. Moreover, by Proposition [1](#A7.EGx5), we have

$$ $\displaystyle\frac{\partial a^{(t)}_{0}}{\partial t}$ $\displaystyle=2Ba^{(t)}_{1}$ (16) $\displaystyle\frac{\partial a^{(t)}_{1}}{\partial t}$ $\displaystyle=a_{1}^{(t)}+Ba_{0}^{(t)}$ (17) $$

In matrix form, these equations form a linear matrix ODE:

$$ $\frac{\partial}{\partial t}\begin{pmatrix}a_{0}^{(t)}\\ a_{1}^{(t)}\end{pmatrix}=\begin{pmatrix}0&2B\\ B&1\end{pmatrix}\begin{pmatrix}a_{0}^{(t)}\\ a_{1}^{(t)}\end{pmatrix}$ $$

The solution of this ODE is given by $\begin{pmatrix}a_{0}^{(t)}\\
a_{1}^{(t)}\end{pmatrix}=\exp\left(t\begin{pmatrix}0&2B\\
B&1\end{pmatrix}\right)\begin{pmatrix}a_{0}^{(0)}\\
a_{1}^{(0)}\end{pmatrix}$. By the first statement of Lemma [1](#A7.EGx23) and the non-negativity of $a_{0}^{(0)},a_{1}^{(0)}$, we get:

$$ $\displaystyle a_{0}^{(t)}$ $\displaystyle\leq 10e^{\sqrt{2}Bt}\left(a_{0}^{(0)}+a_{1}^{(0)}\right),$ (18) $\displaystyle a_{1}^{(t)}$ $\displaystyle\leq 10e^{\sqrt{2}Bt}\left(a_{0}^{(0)}+a_{1}^{(0)}\right)$ (19) $$

We proceed to Equation [11](#S5.E11). Note that for any $s\geq 0$, we can write $u_{2}(s)=\hat{a}_{0}^{(s)}\mathbf{e}_{0}+\hat{a}_{1}^{(s)}\mathbf{e}_{1}$ with $\hat{a}_{0}^{(0)}=a_{0}^{(0)}$ and $\hat{a}_{1}^{(0)}=a_{1}^{(0)}$. By Proposition [1](#A7.EGx5), we have

$$ $\mathcal{Q}\mathcal{L}u_{2}(x)=B\hat{a}_{1}^{(s)}\mathbf{e}_{2}(x)$ $$

Moreover, given a function $v(x)$,
the action of the operator $\exp{\mathcal{Q}\mathcal{L}(\tilde{t})}$ on $v$ is given by the solution $w(\tilde{t},x)$ to the PDE

$$ $\displaystyle\frac{\partial}{\partial t}w(t,x)$ $\displaystyle=\mathcal{Q}\mathcal{L}w(t,x)$ $\displaystyle w(0,x)$ $\displaystyle=v(x)$ $$

If $w(t,x)=\sum_{n\in\mathbb{N}_{0}}b_{n}^{(t)}\mathbf{e}_{n}$ and $\forall n\in\mathbb{N}_{0},b_{n}^{(0)}\geq 0$, we are interested in solving the previous PDE with initial conditions $b^{(0)}_{2}=B\hat{a}_{1}^{(s)}$ and $b^{(0)}_{n}=0\ \forall n\neq 2$.

We claim that the coefficients $\hat{a}_{n}^{(t)}\geq 0$ $\forall t>0$ and $\forall n\in\{0,1\}$. For $t=0$ this is by definition, and we will prove it for all $t$ by way of contradiction. Suppose the claim is not true, then there exists a $t^{*}>0$, and some $n^{*}\in\{0,1\}$ such that $\hat{a}_{n^{*}}^{(t^{*})}=0$, and $\hat{a}_{n}^{(s)}>0$ $\forall n\in\{0,1\}$ and $\forall s<t^{*}$.
But from continuity this implies that there exists $0<t^{\prime}<t^{*}$ such that $\frac{\partial}{\partial t}\hat{a}_{n^{*}}^{(t^{\prime})}<0$. However, it can be easy to see that if $\hat{a}_{n}^{(s)}>0\ \forall s\leq t^{\prime}$, then $\mathcal{P}_{1}\mathcal{L}u_{2}(t^{\prime})>0$ and $\mathcal{P}_{1}\mathcal{L}\int_{0}^{t^{\prime}}\exp{\mathcal{Q}\mathcal{L}(t-s)}u_{2}(s)ds>0$. Therefore, from Equation [11](#S5.E11), $\frac{\partial}{\partial t}\hat{a}_{n^{*}}^{(t^{\prime})}>0$, which is a contradiction.

This claim implies that $b_{n}^{(0)}\geq 0\ \forall n\in\mathbb{N}$, and in turn it implies that $b_{n}^{(t)}\geq 0\ \forall n\in\mathbb{N},t>0$.
Applying $\mathcal{Q}\mathcal{L}$ results in the following inequalities for the coefficients $b_{1}^{(t)}$, $b_{2}^{(t)},b_{3}^{(t)}$:

$$ $\displaystyle\frac{\partial}{\partial t}b^{(t)}_{1}$ $\displaystyle\geq b^{(t)}_{1}+Bb^{(t)}_{2}\geq Bb_{2}^{(t)}$ (20) $\displaystyle\frac{\partial}{\partial t}b^{(t)}_{2}$ $\displaystyle\geq Bb^{(t)}_{1}+4b^{(t)}_{2}+Bb^{(t)}_{3}\geq Bb^{(t)}_{1}+Bb^{(t)}_{3}$ (21) $\displaystyle\frac{\partial}{\partial t}b^{(t)}_{3}$ $\displaystyle\geq Bb^{(t)}_{2}+9b^{(t)}_{3}\geq Bb^{(t)}_{2}$ (22) $$

Thus, we can write a linear matrix ODE for the vector $(b_{1}^{(t)},b_{2}^{(t)},b_{3}^{(t)})$:

$$ $\displaystyle\frac{\partial}{\partial t}\begin{pmatrix}b_{1}^{(t)}\\ b_{2}^{(t)}\\ b_{3}^{(t)}\end{pmatrix}\geq\begin{pmatrix}0&B&0\\ B&0&B\\ 0&B&0\end{pmatrix}\begin{pmatrix}b_{1}^{(t)}\\ b_{2}^{(t)}\\ b_{3}^{(t)}\end{pmatrix}$ (23) $$

Therefore, using Lemma LABEL:l2:growthorder, for sufficiently large $B$ we have $b_{2}^{(t-s)}\geq\frac{Be^{\sqrt{2}B(t-s)}}{10}\hat{a}_{1}^{(s)}.$

Hence, if we write $\int_{0}^{t}\exp{\mathcal{Q}\mathcal{L}(t-s)}\mathcal{Q}\mathcal{L}u_{2}(s)ds$ in the basis $\{\mathbf{e}_{n}\}_{n\in\mathbb{N}_{0}}$, the coefficient for $\mathbf{e}_{2}$ will be lower bounded by

$$ $\int_{0}^{t}\frac{1}{10}Be^{B(t-s)}a^{(s)}_{1}ds$ $$

Applying the second statement of Lemma [1](#A7.EGx23) and using the non-negativity of $a^{(0)}_{0}$ and $a^{(0)}_{1}$, we have $\hat{a}^{(s)}_{1}\geq\frac{1}{10}e^{\sqrt{2}Bs}\left(a_{0}^{(0)}+a_{1}^{(0)}\right)$. Hence, the coefficient for $\mathbf{e}_{2}$ is lower bounded by

$$ $\int_{0}^{t}\frac{1}{10}Be^{\sqrt{2}B(t-s)}\frac{1}{10}e^{\sqrt{2}Bs}\left(a_{0}^{(0)}+a_{1}^{(0)}\right)ds\geq\frac{Bt}{100}e^{\sqrt{2}Bt}\ \left(a_{0}^{(0)}+a_{1}^{(0)}\right)$ $$

We finally need to consider what happens after applying the outermost operator $\mathcal{P}_{1}\mathcal{L}$. Because of Proposition [1](#A7.EGx5) again, applying $\mathcal{L}$ makes the coefficient in front of $\mathbf{e}_{1}$ at least $\frac{B^{2}t}{100}e^{\sqrt{2}Bt}\left(a_{0}^{(0)}+a_{1}^{(0)}\right)$. Finally, applying $\mathcal{P}_{1}$ preserves the coefficient in front of $\mathbf{e}_{1}$.

Hence, equation Equation [11](#S5.E11) results in the following evolution inequalities:

$$ $\displaystyle\frac{\partial\hat{a}^{(t)}_{0}}{\partial t}$ $\displaystyle\geq 2B\hat{a}^{(t)}_{1}$ (24) $\displaystyle\frac{\partial\hat{a}^{(t)}_{1}}{\partial t}$ $\displaystyle\geq\hat{a}_{1}^{(t)}+B\hat{a}_{0}^{(t)}+\frac{B^{2}t}{100}e^{\sqrt{2}Bt}\left(a_{0}^{(0)}+a_{1}^{(0)}\right)$ (25) $$

Using the second statement of Lemma [1](#A7.EGx23) again we have that $\hat{a}_{0}(t)\geq\frac{1}{10}e^{\sqrt{2}Bs}\left(a_{0}^{(0)}+a_{1}^{(0)}\right)$. Thus, dropping the (positive) term $\hat{a}_{1}^{(t)}$ in equation [25](#A7.E25), we have:

$$ $\displaystyle\frac{\partial\hat{a}^{(t)}_{1}}{\partial t}$ $\displaystyle\geq\left(\frac{1}{10}+\frac{Bt}{100}\right)Be^{\sqrt{2}Bt}\left(a_{0}^{(0)}+a_{1}^{(0)}\right)$ (26) $$

Integrating this equations yields:

$$ $\displaystyle\hat{a}^{(t)}_{1}\geq a^{(0)}_{1}+\frac{1}{200}e^{\sqrt{2}Bt}\left(\sqrt{2}Bt+10\sqrt{2}-1\right)\left(a_{0}^{(0)}+a_{1}^{(0)}\right)$ (27) $$

Thus, we have $a_{1}^{(t)}\gtrsim Bte^{\sqrt{2}Bt}\left(a_{0}^{(0)}+a_{1}^{(0)}\right)$. Together with equation [18](#A7.E18), the claim of the Theorem follows.
∎

###### Lemma 1 .

There exists $B>0$ sufficiently large such that for all $t>0$ the matrix $\begin{pmatrix}0&2Bt\\
Bt&t\end{pmatrix}$ satisfies:

$$ (28) (29) $$

###### Proof.

###### Proof.

###### Lemma 1 .

###### Proof.

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: Please see the questions in the weaknesses above.
soundness: 3
strengths: The paper is well-motivated and proposes solution methods as well as a new set of datasets.

The 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. 

The paper introduces a new challenging dataset (although not extensive), the dataset could be used as a next important benchmark for future works.

The 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).
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.
weaknesses: It is not easy to see exactly how the architecture is designed. Including diagrams showing the forward computational path will be very helpful. 

Although 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.

FNOs 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? 

Also, 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)

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: See Weakness 2.
soundness: 3
strengths: 1) The problem of whether to use memory to model time-dependent PDEs is well-defined and well-motivated.
2) 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.
3) 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.
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.
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.
2) 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?
3) lines 1119-1121: this is a strange training setup, but maybe necessary due to the GPU memory limitations.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: See weaknesses
soundness: 2
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.
* 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.
* the paper is well written and clearly structured
* the paper introduces recent state-of-the-art state-space models to the field of learning PDE solutions
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.

The 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.
weaknesses: * The paper claims to improve learning the solution of time-dependent PDEs theoretically grounded in the 
    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.
* 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.
* 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).

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
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. 
2. 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. 
3. 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?
4. How well can the model perform for zero-shot super-resolution once the model is trained?
soundness: 3
strengths: 1. The paper is theoretically well motivated. 
2. The proposed method for memory operators can be implemented easily.
3. The method demonstrates improvement in performance for solving PDEs.
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.
weaknesses: 1. The motivation for using S4 as the model of choice is not convincing.
2. Lack of numerical comparison against benchmarks in 2D cases.

Meta-review: This paper introduces a novel neural operator with memory for modeling time-dependent PDEs. By combining state-space models and Fourier neural operators, the authors establish a sound framework that extends beyond standard Markovian neural operators. The approach is not only novel, but also demonstrates its significance through strong experimental results. In addition, Mori-Zwanzig theory provides a theoretical motivation for the proposed model. Overall, the key strengths of this paper lie in its innovative architecture, theoretical underpinnings, and strong empirical findings. 

All reviewers agree that this is a very strong submission. In particular, I believe that this paper has the potential to inspire interesting future research. Thus, I strongly recommend accepting this paper.

Decision: Accept (Oral)