Learning semilinear neural operators: A unified recursive framework for prediction and data assimilation.
Authors: Ashutosh Singh
Ricardo Augusto Borsoi
Deniz Erdogmus
Tales Imbiriba
Venue: ICLR 2024 poster (2024)
Keywords: Neural operator
PDEs
semi-linear evolution
sequential learning
filtering
data assimilation
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
TL;DR: We propose a flexible recursive neural operator approach for prediction and data assimilation with semilinear PDEs.

Abstract: Recent advances in the theory of Neural Operators (NOs) have enabled fast and accurate computation of the solutions to complex systems described by partial differential equations (PDEs). Despite their great success, current NO-based solutions face important challenges when dealing with spatio-temporal PDEs over long time scales. Specifically, the current theory of NOs does not present a systematic framework to perform data assimilation and efficiently correct the evolution of PDE solutions over time based on sparsely sampled noisy measurements. In this paper, we propose a learning-based state-space approach to compute the solution operators to infinite-dimensional semilinear PDEs. Exploiting the structure of semilinear PDEs and the theory of nonlinear observers in function spaces, we develop a flexible recursive method that allows for both prediction and data assimilation by combining prediction and correction operations. The proposed framework is capable of producing fast and accurate predictions over long time horizons, dealing with irregularly sampled noisy measurements to correct the solution, and benefits from the decoupling between the spatial and temporal dynamics of this class of PDEs. We show through experiments on the Kuramoto-Sivashinsky, Navier-Stokes and Korteweg-de Vries equations that the proposed model is robust to noise and can leverage arbitrary amounts of measurements to correct its prediction over a long time horizon with little computational overhead.

Published as a conference paper at ICLR 2024
LEARNING SEMILINEAR NEURAL OPERATORS: A UNI- FIED RECURSIVE FRAMEWORK FOR PREDICTION AND DATA ASSIMILATION
Ashutosh Singh !.* Ricardo Borsoi
Deniz Erdogmus'
Tales Imbiriba 2.3
ABSTRACT
Recent advances in the theory of Neural Operators (NOs) have enabled fast and accurate computation of the solutions complex systems described by partial differential equations (PDEs). Despite their great success, current NO-based so- lutions face important challenges when dealing with spatio-temporal PDEs over long time scales. Specifically, the current theory of NOs does not present a sys- tematic framework perform data assimilation and efficiently correct the evolu- tion of PDE solutions over time based on sparsely sampled noisy measurements. this paper, we propose a learning-basse solution operators anfinit--dimessinne semilinear PDEs. Exploiting the struc- ture of semilinear PDEs and the theory of nonlinear observers in function spaces, we develop flexible recursive method that allows for both prediction and data assimilation by combining prediction and correction operations. The proposed framework is capable of producing fast and accurate predictions over long time horizons, dealing with irregularly sampled noisy measurements to correct the so- lution, and benefits from the decoupling between the spatial and temporal dy- namics of this class of PDEs. We show through experiments on the Kuramoto- Sivashinsky, Navier-Stokes and Korteweg-de Vries equations that the proposed model is robust to noise and can leverage arbitrary amounts of measurements correct its prediction over long time horizon with little computational overhead.
state-space approach to compute the
INTRODUCTION
The evolution of many dynamical systems in science and engineering can be described by partial differential equations (PDEs), where modeled quantities are often a function of both space and time. Evolving PDEs over time can be computationally very intensive, especially for fine spatiotemporal grids and large scale systems. In this context, great effort has been recently devoted to provide neural-network (NN) approximations integral operators as solution to such PDEs (Lu et al., 2019 Kovachki et al., 2021: Morrill et al. [2021: Kidger et al. (020), namely, Neural Operators (NOs), allowing for efficiently approximating the solution to PDE systems.
Despite the efficiency and approximation power of current NOs, the literature lacks systematic framework for data assimilation (Asch et al. a (016)) capable of improving the estimation quality and correcting the evolving trajectories based on small amount of noisy measurements from the system. This is very important due to advances of sensing technologies and the proliferation of data from dynamical systems many fields such Earth surface temperature (Tiag et al. 2023) remote sens- ing imaging (Weikmann et al. [2021] 1 Borsoi et al. 2021). traffic concentration Thodi et al.) 2023).
fMRI dynamics (Singh et al. [202]] | Buxton 2013) and video dynamics (Guen & Thome 2020). to
name but a few. Afshar et al. (2023) discussed the design of observers for the so-called semilinar
PDEs (a class of common PDEs with specific form of spatiotemporal dynamics) and its close con- nection to extended Kalman filters (Smith et al. 1962. Särkkä & Svensson| 2023). The designed observer an important tool that can estimate the state of a system governed by some underlying
denotes equal contribution.
1 Department of Electrical and Computer Engineering, Northeastern University. Boston MA, 02115, USA.
CNRS. CRAN, Université de Lorraine, Vandoeuvre-lés-Nane F-54000 Nancy, France.
3 Institute for Experiential AI, Northeastern University, Boston MA, 02115, USA. edu, raborsoi 3gma com Emails: {s ngg..asse ashu, d. erdoggus,, ta es woormeettttttttnnncmmmmmm ia

---

Published as a conference paper at ICLR 2024
dimenssonni PDEs based measurements, showing natural connection to data assim- ilation. However, standard approaches based on Kalman filters and observers require tremendous computational efforts, limiting their application to large-scale systems. Recently, a few works aimed at approximating the Kalman filter updates using NNs for video prediction (Guen & Thome (020), while the Kalman updates were approximated in (Revach al. 2022). Nevertheless, general the- ory and framework for data assimilation exploiting the structure of semilinear PDEs is still missing,
In this paper, we extend the theory of NOs by exploiting the observer design of semilinear PDEs and provide a systematic approach for data assimilation. The resulting method, namely NO with Data Assimilation (NODA), recursive NO approach that can be used for both estimation, when noisy measurement data is available, and prediction. To overcome the computational challenges inherent to the design and implementation of infniiteddimensional Kalman observers, we propose a data-driven NN approximation leveraging NOs (Li et al.) (020a) and saaringg-assed Kalman esti- mators (Guen & Thome, 2020). Thus, NODA recursively estimates the system's 's states over time through prediction and update steps typical to Bayesian filtering approaches but with small com- putational cost. We demonstrate through extensive simulations that NODA leads to better prediction performance when compared with other closely related NO approaches, while also being able to assimilate data with arbitrary sampling rates. This can have significant impact in practical appli- cations including, e.g., weather and Earth surface temperature forecast (Pathak et al. 2022 Jiang] eet al.] (023), remote sensing imaging /Weikmann et al.) 2021 Borsoi et al. (021) and fMRI dy-
namics (Singh et al., 2021 Buxton 2013). The contributions of this work can be summarized as:
we extend the NO theory by leveraging the observer design of semilinear PDEs;
we break the observer solution into prediction and update steps, allowing for systematic way of performing both prediction and data assimilation, and devise data-driven solution;
the resulting framework can estimate solutions using arbitrary amounts of measurements.
Related work: Traditionally, many methods have been proposed to approximate the solution of PDEs, as can be found in (Larson & Bengzon 2013) and (Evans 2022). A major drawback of traditional solvers is that they are computationally taxing. A class of these methods resorts to data- driven approaches that aim approximating the solution operator of the underlying PDE from the snapshots of the states (Kutz et al. 20166] Krstic & Smyshlyaev] 2008). Recently, NO learning has become a popular strategy for learning solution operators to broad family of parametric PDEs
(Li et al. [220a b) Kovachki et al.] 2021 Brandstetter et al.] 2022a). These methods have been
further used in many applications such as for weather forecast ( (Pathak et al. al.] (222)) and to track CO2
migration (Wen et al. 2022) and coastal floods (Jiang et al. 2021). ). among others.
NOs are also being adopted in different machine-learning paradigms. Guibas et al. (021) present work on using a Fourier NO for token mixing in transformers. Goswami et al. (2022) motivates a physics-informan
learning of NOs. |Brandstetter et al. (202bb) use the symmetries of Lie groups for data augmentation. Gupta et al. (2021) proposed mullwwavelet-based operator leaning framework. Rotman et al. (2023) and Kaltenbach et al. (2023) present a semi-supervised learning approach for NO learning. Hao et al. (2023) propose the use of transformers for NO learning. CCen et al. (2023) adapts NOs for different geometries. In (Magnani et al.] 2022) the authors provide a Bayesian treatment of NO for uncertainty estimation. One of the key drawbacks of these methods is that for
long time scales, the solutions tend to deviate from the true trajectory. Li et al. (2022) present a Markov NO for tracking long-term trajectories of chaotic dissipative systems. However, such works have not considered the presence of noisy measurements, which is common in real-world datasets. This motivates the need for new frameworks that can address data assimilation.
Similar to the proposed work. Salvi et al. (2022) attempt to exploit the structure of the solutions of semi-lineat stochastic PDEs and propose learning frameworks based on solving a fixed-poin prob-
lem using ODE numerical solvers. However, the method on (Salvi et al., 2022) differs from the proposed framework as it does not address the need for correcting predictions using noisy measure- ment available during test time and is not recursive in nature. The capability of learning solution operator that is ficxible, recursive and robust to noisy data is essential for performing fast and accu- rate predictions and data assimilation over long time horizons.
Data assimilation has been well studied in the context of dynamical systems (Cheng et al. 2023 Farchi et al. 2021: Levine & Stuart 2022 Levine, 2023). Most methods discussed in these works exploit auto-regressive formulations, which can provide good approximations of the dynamical be- havior typically seen in such systems. Recently, |Fion et al. (2023) extended deep dynamic mode
2

---

Published as a conference paper at ICLR 2024
decomposition learn Koopman operators through data assimilation. In the case of infinite dimen- sional systems, Afshar et al. (2023) discusses the observer design of semilinear PDEs and its close relationship with difnite-dimensioon extended Kalman filter, which is closely connected to data assimilation methods. We leverage these principles to propose NODA as it will become clear in the remainder of this paper.
BACKGROUND
Learning Solution Operators to PDEs: Many applications ranging from weather forecasting to modelling molecular dynamies involve finding the solution to PDEs both spatial and temporal variables, herein denoted by Rs and by t E R, respectively. The generic PDE considered in recent works is of the form
(La=) 1) = fxx,
ED, (z. E AD.
(I)
=(x, () 0,
Here = E Zzt represents the solution of the PDE that belongs to Banach space Z C C(D, ("), its domain given by a bounded, open set D C I29 & with a Dirichlet boundary condition imposed over 0 E Z is a known function, L Zzt + Z is a differential operator with a set of parameters a, and 2 is the dual space of Zzt- Recent methods in neural operator learning focus on computing parametric mapping Go + which approximates the solution to the PDE in equation based on set of pointwise evaluations {as, 2j )=1- Such approximation performed by minimizing the empirical risk (Kovachki et al. 2021)
min
N Di=
1| Go(az)l2
where || l|z denotes an appropriate norm on Z. A closely related problem consists in learning an operator that computes snapshot solution of the PDE at a time t >> 0, 2(-, t), based on some initial condition z(-, 0) and on set of parameters a.
Different NOs were recently proposed to define the parametric mappings Go. motivated by the kernel integral formulation of the solution operator to linear PDEs (Li et al.) [220b] Kovachki al. 202I) These are given in the form of the composition of different layer operators, such as:
g(a) = (Q W2 ..... Wi o (a),
with W/ being the operator for layer f, P is a lifting operator, mapping a to the first hidden represen- tation, and Q a projection operator, mapping the last hidden representation to the output. Different forms for the operator layers have been proposed, such the graph and multipole graph NOs (Ko- vachki et al. (221)), the U-shaped NO (Rahman et al. 2023), and the DeepONets (Lu et al. 2019).
A notable example is the Fourier neural operator (FNO) layer (Li et al. 2020a). The FNO was motivated by constraining the kernel in the integral operator as a shift-invariant function (Li et al. 20206), thereby making the integral operator equivalent convolution that be implemented in the frequency domain. An FNO layer, applied to an input signal and evaluated at (z, f) E D, can be represented generically
(x = t))
(Z,)))
(4)
where a is a componentwise nonlinear activation function, W is linear operator, F is the Fourier transform, and Rs is the Fourier transform of the kernel, parametrized by 6. The FNO can be easily discretized for form NO layer, and be made invariant to the discretization level (Li al. 2020a)..
Many applications involve PDEs which exhibit a very specific spatiotemporal structure, where the time variable has linear dynamics, while the spatial variables can be governed by more complex and possibly highly nonlinear equations (Curtain & Zwart 2020). This specific spatiotemporal decoupling, in the so-called semi-linear form, is represented as (Rankin 1993):
Semilinear PDEs:
Qz(t) Az(t) G(z(t),t). Øt
20) 20
(5)
where z(t) Z, denotes the solution time index t [0,ty], which is itself function of the
spatial variables x E R D(A) + Zn be a linear bounded operator with domain D(A) ((A) C Z, that generates continuous semi-group operator T(I)= Zz 2, (Curtain & Zwart 2020), and
3

---

Published as a conference paper at ICLR 2024
G: G:Z,x x (0.tyl + Z2 is operator which is strongly continuous [0, and possibly non-linear in Zz, satisfying G(O,t) t) = O. The operator (()) is closely linked to the temporal evolution of the PDE solution 2(t), as will become clearer in the following section. Note that 2(t) is still a function
of the spatial variables æ. Examples of PDEs with this structure include the Navier-Stokes and the -iraatto-Svvsshinky equations, which can be seen in Section 4
3 PREDICTION AND DATA ASSIMILATION WITH RECURSIVE NOS
In this section, will present the proposed framework for a recursive method to handle both pre- diction of the PDE solution as well as data assimilation. Although semilinear PDEs can describe the time evolution of many infinite dimensional systems, in many practical applications such as in weather forecast (Pathak et al. 2022) or tracking of coastal floods or CO2 migration (Wen et al.) 2022 Jiang et al. 2021p, one does not have direct access to measures discrete snapshots of the true solution (t). Instead, we only measure noisy or degraded version thereof in Euclidean space, which we denote by (t) E I.. Thus, we consider the following model:
dz(()
Ōt A=(() + ((=()))))++))
(6)
=00) 20:
(7)
v(() Cz(() + n()-
(8)
Note that, besides a dynamical evolution based on a semilinear PDE in equations (() and 77. the second part of the model in equation describes how the measurements are generated from the true solution (or states) z(t), with r7()) E TRP representing the measurement noise and C E ((2,, IR?)
being mapping from the solution space to the measurement space. w(t) Zz is an unknown perturbation signal which describes uncertain knowledge of the system dynamics.
As mentioned previously, in many applications we only observe v(() sparse, intermittent) time samples T C [0, ly] in an interval of length L,. Given an initial condi- tion/initialization 20 and the set of discrete measurements {t((t) E T,, our goal is to learn an online model capable of recovering/estimatin the solution xt) recursively over time, using only present past measurements {u(T) a flexible method, capable of performing prediction in order to recover 2(t) at time instants when no measurements are available (using only the past estimates of =(T) for 7 t, thus dealing with missing data), and to correct the solution with the measured data y(t) on the small number of time instants when it observed, performing data assimilation.
discrete set of (possibly
f E ' at every time t. This entails the need for
In the following, we will first investigate the particular form of the solutions to the PDE in 6 (8) and of observer designs that can provide consistent estimate of its solution 2(t). This will be paramount to support the data-driven solution proposed thereafter.
3.1 OBSERVER DESIGN FOR SYSTEM WITH SEMILINEAR PDES
The specific structure of the semi-linear PDE allows us write the solution in general form under mild conditions. Thus, assuming that the nonlinear operator G' admits Fréchet derivative that is globally bounded as well as locally Lipschitz, uniformly in time, there exists a solution 2(t) E Zz
to equation 5 for a time interval 0 < 177 < l,. that can be written as (Weissler 1979 Rankin 1993)
=(1) = T + ['' T(t - G(z(s))( s)ds.
When considering the availability of measurements y(t) as in equation these can be used to correct the state evolution of the system by designing an observer. This is performed by designing a PDE that contains a copy of the original system's dynamics with the addition of a correction term which is based on the error between the predicted and actual measurements, hereby denoted by C.i(t) y(t),
by means of observer operator gain (Afshar et 2023). For equations 6 (8. the general form of the observer is given by
a²(t)
Ōt A&(t) + G(8(t),t) + (()[v(t) C2(t))
2(0) 20
(10)
where K(t)= RP + Z, is the observer gain, mapping the error in the measurement space to the PDE dynamics, and a() is the observer solution.
4

---

Published as a conference paper at ICLR 2024
By solving nonlinear infiite-dimmnsionne Riccati equation, one can design an operator K ((t) based on the PDE in equation 10 that can compute the solution 2(t) over time, such that under the absence of noise and disturbances, and under mild additional conditions, the solution satis- + 0 as - oo for sufficiently small initial errors (see Theorem 5.1 in (Afshar al., 2023 ). Furthermore, the authors show that under bounded disturbances w(t) and ntt), the estimation error can also be bounded for all t (see Corollary 5.2 in (Afshar et al. 2023 ). Similarly the previous case, when (t) is strongly continuous the analytical form of solution 2(t) to (10 is given by:
fies (t= 8(1) z.
2(1) = T + J [[Tu- T) ()((()))).)) + {())(((s) -C(()
(11)
The observer solution, however, challenging to implement. First, the numerical solution to the infnite-dimensional Riccati equation and computation of equation [T is difficult and computa- tionally expensive. Moreover, it requires measurements ytt to be available at all time instants
t E [0, tyl. and cannot handle missing data or perform prediction. These challenges will be ad- dressed in the next section by the learning-basee solution.
3.2 LEARNING-BASEE RECURSIVE PREDICTION- COEEIIITOO-OOORIIIIN NO
In this section, we provide a fully data-driven framework to compute a recursive solution to ( which can perform prediction and handle missing data. First, we discretize the observer solution in the time domain. By rewriting the solution in equation 11 we can use it to represent the evolution of the estimate a(t) between any two discrete time instants, - and ty (for convenience, we assume the discretization timesteps hy lx fk-1 to be approximately constant with k). Moreover, possible to rewrite this solution such a form that we can identify two terms, prediction and
correction (see Appendix A for details), leading to:
a(tx) = T(t - 44-14 (4-4))((-1)+ 1d T(4=
- G)G((()),
T( - -CE)))
1k-1
Prediction
Correction
(12)
The prediction term depends only the past solutions 5(4x-1), predicting the evolution of the PDE states, whereas the correction term is based on the reconstruction error. This structure will be used in the design of the proposed learning-based solution.
Moreover, since most differential operators A generate a semigroup operator T which has an in- tegral form (Gerlach 2014), this motivates the of NO framework to approximate the terms in equation [2] Thus, we propose a -earning-based approach that leverages the structure of (12) to perform both prediction and data assimilation in flexible manner. Denote the estimated solution at time te and at some spatial discretization level (typically in a regular grid) by &t, (z )), where II a discretization operator which maps from the original solution space Z, to some finite- dimensional Euclidean space. Specifically, we separate the prediction and correction terms in (12), and propose architectures, described in the following, to learn each of these operations separately. The prediction operation is applied at every time instant tg. generating the predicted solution spred Then, to adapt the method to the assimilation setting, measurements are available (i.e., v(tx) is observed) the predicted solution is further processed by the correction operation to generate the final estimated solution بيg- Otherwise, there are no measurements, the estimated solution is set as
the predicted one as 2t =pred This process, which constitutes the NODA framework, is further described in the following.
which predicts of the solution at time tx given the estimate at time tx-1, with a learnable NO in a residual formulation, given by:
Prediction:
We approximate the first term in equation 12
apred 2t
- + W(
(13)
where Wi is neural operator. Note that this operation inherits the discretization invariance from the NO W, such as in the FNO (Li et al. 2020a). In this work, we use as W an FNO layer as defined in equation 4 with the Fourier transform computed only over the spatial dimension, a.. If there are no measurements available for assimilation at time lg, then we set the estimated solution as the predicted one, 2, = 2preed If a measurement is available, however, the predicted solution is further added to the correction term as described in the following.
5

---

Published as a conference paper at ICLR 2024
We approximate the second term in equation 12 which corrects the predicted solu- tion "tx based on the available measurement y(tx), using the following model, inspired from the observer operator form in 12 as
Correction:
Lt. apred (( [v(()) F((E
(14)
where E is a learnable operator which approximates the (possibly unknown) measurement operator C' in equation g over a given discretization level. In this work, we parametrized E using a one- hidden layer fully connected ReLU NN, which can also tackle cases where the measurement model might be nonlinear.
To parametrize the observer gain K, generalized the gating architecture used in (Guen & Thome 2020 to cope with general measurement operators C' and the observer design:
K ((pred) (
[u] tanh (W ( Wyz((a) + b) (Ó*u),
(15)
where @ denotes the Hadamard product and Wa, Wy and b are learnable linear operators. Ö* is the conjugate operator of Ć C Hp. č & + y(t) being the discretized version of C. Note, however, that O* can be a learnable linear operator in (15). particularly when C' is unknown.
Learning criterion: Given a set of training data with S different realizations of trajectories of discretized solutions to the PDE (6 each generated from a different random initial condi- (1) (to) and containing N time samples, which we denote by {=p (tx), (tw)) for i i.,...., S, NODA is learned by minimizing the following loss function:
tions 2D
s
N
H
J(6) = SN
-00 (1+)22+ Ptz
(16)
SH
(0) where "D (tx) = (2 (tx)) denotes the the discretization of the PDE trajectory, and $ denotes
the recovered solution to the PDE, for the i-th realization. To address both the prediction and data assimilation objectives during training, NODA is initialized with = =D (4o). then, we supply it with measurements up time index t S tN. i.e., (o (tx)) which are used to perform data assimilation in the interval [4,,4r] after which it performs only predictions between ( and In.
generate the estimated trajectory {*t. N k= The loss function (() contains two terms. The first term measures the reconstruction of the PDE solution, while the second term measures the ability of the model reconstruct the measurements, and the balance between these two terms is controlled by the weight X 2 0. The second term is particularly important when learning the operator E. The learnable parameters are denoted by d, and contains the parameters needed for the prediction (i.e., the NO W) as well as the correction step (i.e., E and the parameters of K).
4 EXPERIMENTS
In this section, investigate the use of the proposed NODA framework on ID and 2D semilinear PDEs. We present experiments with the intent of showing the efficacy of the proposed model for: a) Accurately predicting trajectories over long temporal horizons using only prediction steps; b) The efficacy of the correction step and the influence of measurement noise in the estimation performance; and c) The influence of the amount of sparsely measured data in the assimilation performance. For b) and c), the snapshot measurements y(4)) are generated by synthetically adding white Gaussian noise the clean trajectories 2D (tx), with signal-to-noise (SNR) ratios varying between 10dB and 30dB. Morcover, the noiseless case is also considered in some examples. For the third experiment, c), we define as a E [0, 1] the ratio between the amount of available (noisy) measurement data (sampled at random instants) and the length of the prediction horizon Tp [f fy]. To exploit the potential NODA during test time, a warm-up procedure where prediction and update steps are performed at every time step. We define the warm-up horizon window Tw [f1, L Hr). and assume that data is available for every time step t4 E Tw. In all experiments, we evaluate NODA's predictive performance the prediction horizon Tp for different values of a% (including only the prediction performance), of tyt. and for different SNRs, depending on each experiment. All the competing algorithms are initialized with the true solution after the warm-up period, 2n (In--1)
0% to evaluate

---

Published as a conference paper at ICLR 2024
Table 1: Averaged ReIMSE (x (x109) for prediction (a: = 0%) on the Kuramoto-Sivashinsky equation
as function of the sequence length y and of the SNR.
SNR
20 dB
30 dB
4
60
90
120
60
90
120
60
90
120
MNO
144.77
403
143116
108.16
379.115
521116
13.11
20418
410.177
FNO
29819
412:112
715114
2591100
324.1100
659113
223:1111
643.119
892.118
LSTM
303+:15
130=115
33090+:33
13
483+10
981
14
688:221
953+119
NODA
123+:10
129=13
493:t10
63499
310=9
483+-:0
1944
1854:1
405+:12
Table 2: Averaged ReIMSE (x10²) for prediction (a 0%) on the Navier-Stokes equation as function of the sequence length ły and of the SNR.
SNR
20 dB
30 dB
QC
500
750
1000
500
750
1000
500
750
1000
MNO
48:17
72:6
96
3914
51=7
87:16
1843
22:133
26:
ENO
192::10
294::1
351=14
189419
2654:21
317416
136:11
1774:14
265::6
MWNO
172.110
22619
278
147.99
192.19
240.110
78.15
104.16
188.111
C-LSTM
487=15
522+=10
604=18
397#20
426+117
481+19
366#11
417#:13
529+11
NODA
10:2
18:14
32+-6
8+:1
13=2
26
7+1
8+1
13#2
Table 3: Averaged ReIMSE (x109) of NODA for the Navier-Stoker equation as a function of a.
0%
10%
20%
30%
Averaged ReIMSE
26:14
1811
1313
943
We benchmark our results on 2D PDE for the Navier-Stokes (NS) equation and on 1D PDE for the Kuramoto-Sivaan Sivassiiski (KS) equation. We also performed experiments on the Kortweg-
Vries (KdV) equation (Wazwaz 2010). but due to space limitations the results are only included in Appendix C These data sets represent a family of chaotic semi-linear PDEs.
Benchmark Methods: We compare NODA with the NO-based methods FNO (Li et al. 2020a
MNO (Li et al.) 2022). and MWNO (Gupta al., (02T) (for the NS), and with the auto-regressive neural network architectures Conv-LSTM (Shi et al. 2015) (for the NS) and LSTM (Chung) ctt al. 2014) (for the KS). The MNO and LSTM-based methods were used as baselines due to their recurrent nature, which is also one of the main aspects of NODA, while FNO selected since it is one of the main building blocks used in the NODA architecture. To evaluate the ac- curacy of the results of the different methods, we use the average Relative Mean Squared Error (ReIMSH) to compare the accuracy of the trajectories estimated by each method, which is defined as ReIMSE = B E% (+x) - ((t) ΣV. [I= (1)) where dD (tx) is the discretized ground truth sequence, =(4+) is the sequence estimated by the method, and E{- denotes the ex- pected value operator, which computed over the random initializations that generate the trajecto- ries (t&). Details on how the loss function is optimized, and on architecture and hyperparameter selections are given in Appendix B
4.1 DATASETS AND SETUP
Srramoto-Sivashinste (KS) equation: The one dimensional KS equation can be represented as
8t d where =(l) is defined over the spatial domain E [0, J]. with initial = 2g. For the experiments we selected J = 64T and periodic boundary conditions on [O. We generate 1200 sequences with t, E {60, 90, 120} seconds each and a timestep of h = 0.25 seconds from independent initializations are generated by randomly sampling according to the procedure described in (Li et al.) 2022) We used 1000 sequences for training the different methods, and the remaining 200 for the evaluation of their performance. The resolution was fixed at 512 samples. To test NODA we used a warm-up period of tH = 40 seconds.
82
--
82
82
condition 2(0)
Navier-Stokes (NS) Equation: The 2D Navier-Stokes equation represents a viscous, incom- @z(2.t)
pressible fluid in vorticity form on the unit torus, which can be represented u(x, V2 t) = tA + f(x) and V u(x,t)
5t + 0, where u is the velocity field, Re: Re > is the Reynolds number, the forcing term is
2 V u is the vorticity, f(z) sin(2 + v)) + cos(27(x + y)), and the initial condition is 2(x 0) = 2u(z). Note that the NS equation can be written in semilinear form (see Section 3.3 of (Temam) 1982

---

Published as a conference paper at ICLR 2024
We generated 200 independent tra- jectories with y E (500, 750, 1 seconds each and timestep of h second by randomly sampling 2p according to the procedure described
Groundd mnth
in (Li et al., (2222.. We use 180 trajectories for training and evalu- ate on the remaining 20. We fix the resolution at 64x64 for both train- ing and testing. We select Re 40, which presents a non-turbulent case for NS, for benchmarking and comparisons. To test NODA we used warm-up period of ty 50 seconds. Additional experiments for with higher Reynolds numbers (Re & 5000)) are available in Appendix B
(Padd SIYSE
Figure 1: Samples of realization of a true trajectory of the Navier-Stoke equation for t {150, 300, 450} (Top Row), and elementwise enror plots for the corresponding predic- tions by NODA (Bottom Row).
4.2 RESULTS AND DISCUSSION
a) Prediction performance: We evaluate the prediction of the solutions over the test interval [y] (excluding the samples that were used as warm-up). compare NODA with no data assimilation (i.e., using
0%) to the remaining methods, for different sequence lengths L,. SNRs, and for the KS and NS equations (additional results for KdV equation can be found in Appendix C We compute the averaged ReIMSE results, which is presented for the KS and NS equations in Tables [ [nd respectively.
2
Figure 2: Prediction error for realization of the KS equation. Warm-up with 40 was performed. Plots depict the prediction period (40, 200] seconds. Ground truth evolution (Lefft). Error plots for NODA's solutions with 0% (Center) and 30% (Right).
can seen that the averaged errors obtained by NODA are general significantly lower than those achieved by the competing methods, with the MNO being the second-best performing ap- proach, followed by the MWNO (for the NS example), by the FNO and by the LSTM- meth- ods. This behavior consistent across the equations, although MNO achieves slightly better results for SNR 10dB and ty 60 in the KS equation. However, we note that MNO uses the Sobolev norm and dissipativity regularizations in the training objective instead of the Euclidean norm as in the other methods. This illustrates the ability of NODA to outperform competing meth- ods both clean and noisy data well short and long rollouts. A visualization of samples of the ground truth trajectory, as well as error plots for the NODA method and the NS equation can be seen in Figure [ The plots demonstrate the high accuracy obtained with NODA predictions, which emror that show no considerable increase between / 300 and L. 450. Prediction results for the NS and KS equations are shown in Appendix
We also evaluate the influence of the amount of samples tu used in the warm-up step of NODA on its predicting performance by computing the averaged ReIMSE as a function of tu for the NS equation for 0%, 1, 1000 and two different SNRs. It can be seen that the error decreases considerably between tu = 1 and tu = 150, after which is stabilizes. This shows the benefit of the warm-up step in the performance of NODA, and also demonstrates that relatively small values of tu suffice to provide considerable performance improvements. Moreover, comparing these plots to the results the other NO-based approaches (i.e., MNO and FNO), that NODA provides com-
8

---

Published as a conference paper at ICLR 2024
petitive performance even for tu 1 (i.e., performing predictions based on single data snapshot), particularly for noisy data, and significantly better for larger values of ty-
b) Data assimilation performance: We evaluate the performance of NODA while performing data assimilation in the test interval [tu+1.tf t (i.e., when a > 0). We evaluate the averaged ReIMSE for the NS equation with SNR = 30dB and f = 1000 with different data assimilation sampling rates a E (0%, 10%, 20%, 30%}. The results are shown in Table 3 from which it can be seen that the data assimilation provides an important improvement the estimation performance of NODA, which increases consistently with a. Moreover, in the right panel of Figure 3 we also show the averaged ReIMSE for the KS equation as function of the both the sequence length ty E [60, 140] and of a. for an SNR = 10dB. It can be seen that the error increases only moderately with the length of the prediction interval for all tested values of a, which indicates that NODA performs well when estimating both short and long trajectories. Moreover, the averaged ReIMSE decreases consistently with a, demonstrating the improvements of data assimilation on the results.
A realization of the ground truth and elementwise prediction errors of NODA for the KS equation, both without data assimilation and for a = 30% shown in Figure 2 can be seen that the prediction errors tend to increase moderately with I, which is consistent with what was observed in Figure 3. Morcover, using data assimilation with NODA considerably reduces the prediction errors compared to the case where performance of NODA on the KdV equation are shown in Appendix C
0%. Additional results showcasing the data assimilation
e) Computation overhead: To assess NODA's computational overhead in comparison with other NO-based algorithms, we measured the average time required for the prediction plus assimilation per sample for the NS (Re 40, a 30%) example. The fastest algorithm was the MNO, which took 8 X 10-5 seconds. FNO was the slowest method, taking 2.25 times as long as MNO, while NODA had an intermediate performance, taking 1.5 times as long as MNO, illustrating its small computational overhead compared to MNO.
10
NODMdb
MNDwch INO+db
Average ReiMSB
Lo-1
NCDA 10cb
Average ReiMSS
10db
PNO 10dh
0%
10-3
0 20%
30%
350
100
Figure 3: Average ReIMSE as function of the warm-up length Lr for the NS equation for different SNRs (Left); and average ReIMSE function of , for different values of a for the KS equation, where a warm-up period of H = 40 being used for NODA (Right).
5 CONCLUSIONS
In this work, we proposed NODA, a recursive neural operator framework for prediction and data assimilation. To do so, we extended the NO concepts by leveraging the structure of semilinear PDE systems and the correction-biie based state estimation paradigm from the theory of infinite dimensional observer design. As result, NODA incorporates prediction and update steps, allowing it to perform both prediction and data assimilation depending on the availability of noisy measurement data. Ex- tensive experiments demonstrate the capability of the proposed methodology to provide significantly more accurate estimations of the solutions of the systems even when only prediction is performed, and especially when doing data assimilation. We also show that even without the warm-up phase (i.e., when performing predictions using only a single snapshot of data), NODA solutions are still far better than the ones provided by its direct counterpart, the FNO. Nevertheless, NODA's framework flexible and can be directly extended to incorporate other NOs and training losses proposed in the literature, such the approach adopted by the MNO.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGMENTS
This work was supported in part by the French National Research Agency, under grants ANR- 23-CE23-0024, ANR-23-CE94-0001, and by the National Science Foundation, under grant NSF 2316420, and the National Institute of Health, NIH OT2OD030524.
REFERENCES
Sepideh Afshar, Fabian Germ, and Kirsten Morris. Extended Kalman filter based observer design for semilinear duimite-dimensiona systems. IEEE Transactions on Automatic Control, 2023.
Mark Asch, Marc Bocquet, and Maëlle Nodet. Data assimilation: methods, algorithms, and appli- cations. SIAM, 2016.
Ricardo Augusto Borsoi, Tales Imbiriba, José Carlos Moreira Bermudez, Cédric Richard, Jocelyn Chanussot, Lucas Drumetz, Jean-Y Yees Tourneret, Alina Zare, and Christian Jutten. Spectral vari- ability in hyperspectral data unmixing: A comprehensive review. IEEE geoscience and remote sensing magazine, 94))223-2770 2021.
Johannes Brandstetter, Rianne den Berg, Max Welling, and Jayesh K Gupta. Clifford neural layers for PDE modeling. arXiv preprint arXiv:2209.0e4 04434,, 2022a.
Johannes Brandstetter, Max Welling, and Daniel Worrall. Lie point symmetry data augmenta- tion for neural PDE solvers. In International Conference on Machine Learning, pp. 2241-2256. PMLR, 2022b.
Richard Buxton. The physics of functional magnetic resonance imaging (IMRI). Reports on Progress in Physics, 76(9):096601, 2013.
Gengxiang Chen, Xu Liu, Yingguang Li, Qinglu Meng, and Lu Chen. Laplace neural operator for complex geometries. arXiv preprint 2.00::2302.08/600 2023.
Sibo Cheng, César Quilodrán-Casas Said Ouala, Alban Farchi, Che Liu, Pierre Tandeo, Ronan Fablet, Didier Lucor, Bertrand looss, Julien Brajard, al. Machine learning with data assimilation and uncertainty quantification for dynamical systems: review. IEEE/CAA Journal of Automatica Sinica, 1066):1311-138,, 2023.
Junyoung Chung, Caglar Gulcehre, KyungHyun Cho, and Yoshua Bengio. Empirical evaluation of gated recurrent neural networks sequence modeling. arXiv preprint arXiv:1412.3555 2014.
Ruth Curtain and Hans Zwart. Introduction diffestedimmensional systems theory: a state-space approach, volume 71. Springer Nature, 2020.
Lawrence Evans. Partial differential equations, volume 19. American Mathematical Society, 2022.
Alban Farchi, Patrick Laloyaux, Massimo Bonavita, and Mare Boequet. Using machine learning to correct model error in data assimilation and forecast applications. Quarterly Journal of the Royal Meteorological Society, 147(739):30673388 2021.
Anthony Frion, Lucas Drumetz, Mauro Dalla Mura, Guillaume Tochon, and Abdeldjalil Aissa El Bey. Neural Koopman prior for data assimilation. arXiv preprint arXiv:2309.05317, 2023.
Moritz Gerlach. Semigroups of kernel operators. PhD thesis, Universität Ulm, 2014.
Somdatta Goswami, Aniruddha Bora, Yue Yu, and George Em Karniadakis, Physies-informed neu- ral operators. arXiv preprint arXiv:2207.05742 2022.
Vincent Le Guen and Nicolas Thome. Disentangling physical dynamics from unknown factors for unsupervised video prediction. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 11474-11484, 2020.
John Guibas, Morteza Mardani, Zongyi Li, Andrew Tao, Anima Anandkumar, and Bryan Catan- zaro. Adaptive Fourier neural operators: Efficient token mixers for transformers, arXiv preprint arXiv:2111. 2021.
10

---

Published as a conference paper at ICLR 2024
Gaurav Gupta, Xiongye Xiao, and Paul Bogdan. Multiwavelet- operator learning for differen- tial equations. Advances in neural information processing systems, 01224088-24066 2021.
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. International Conference on Machine Learning, pp. 12556-12569, PMLR, 2023.
Peishi Jiang, Nis Meinert, Helga Jordão, Constantin Weisser, Simon Holgate, Alexander Lavin, Björ Lütjens, Dava Newman, Haruko Wainwright, Catherine Walker, et al. Digital twin earth- coasts: Developing a fast and physics-informed surrogate model for coastal floods via neural operators. arXiv preprint arXiv:2110.07100, 2021.
Peishi Jiang, Zhao Yang, Jiali Wang, Chenfu Huang, Pengfei Xue, TC Chakraborty, Xingyuan Chen, and Yun Qian. Efficient super-resolution of near-surface climate modeling using the Fourier neural operator. Journal of Advances in Modeling Earth Systems, 15())=2003350038000 2023.
Sebastian Kaltenbach, Paris Perdikaris, and Phaedon-Stelie Koutsourelakiis. Semi-supervised in- vertible neural operators for Bayesian inverse problems. Computational Mechanics, pp. 1-20, 2023.
Aly-Khan Kassam and Lloyd N Trefethen. Fourth-order time-stepping for stiff pdes. SIAM Journal on Scientific Computing, 26(4):1214-1233, 2005.
Patrick Kidger, James Morrill, James Foster, and Terry Lyons. Neural controlled differential equa- tions for irregular time series, Advances in Neural Information Processing Systems, 33:6696- 6707, 2020.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint 1412,6690 2014.
Nikola Kovachki, Zongyi Li, Burigede Liu, Kamyar Azizzadenesheli, Kaushik Bhattacharya, An- drew Stuart, and Anima Anandkumar. Neural operator: Learning maps between function spaces. arXiv preprint arXiv:2108.066 08481, 2021.
Miroslav Krstic and Andrey Smyshlyaev. Boundary control of PDEs: course on backstepping designs. STAM, 2008
Nathan Kutz, Steven Brunton, Bingni Brunton, and Joshua Proctor. Dynamic mode decomposition: data-driven modeling of complex systems. SIAM, 2016.
Mats Larson and Fredrik Bengzon. The finite element method: theory, implementation, and appli- cations, volume 10. Springer Science & Business Media, 2013.
Matthew Levine and Andrew Stuart. A framework for machine learning of model error in dynamical systems. Communications of the American Mathematical Sociery, 2(07):283-344 3344, 2022.
Matthew Emanuel Levine. Machine Learning and Data Assimilation for Blending Incomplete Mod- els and Noisy Data. PhD thesis, California Institute of Technology, 2023.
Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, An- drew Stuart, and Anima Anandkumar. Fourier neural operator for parametric partial differential equations. arXiv preprint arXiv:2010.08899 2020a.
Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, An- drew Stuart, and Anima Anandkumar. Neural operator: Graph kernel network for partial differ- ential equations. arXiv preprint S00i::2003.03488 2020b.
Zongyi Li, Miguel Liu-Schiaffini, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Learning chaotic dynamics in dissipative systems. Advances in Neural Information Processing Systems, 35:16768-16781, 2022.
Lu Lu, Pengzhan Jin, and George Em Kamiadakis. Deeponet: Learning nonlinear operators for iden- tifying differential equations based on the universal approximation theorem of operators. arXiv 2019.
preprint
11

---

Resolution- Published as a conference paper at ICLR 2024
Emilia Magnani, Nicholas Krämer, Runa Eschenhagen, Lorenzo Rosasco, and Philipp Hennig. Ap- proximate Bayesian neural operators: Uncertainty quantification for parametric PDEs. arXiv preprint arXiv:2208.01556 2022.
James Morrill, Cristopher Salvi, Patrick Kidger, and James Foster. Neural rough differential equa- tions for long time series. In International Conference on Machine Learning, pp. 7829-7833 PMLR, 2021.
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. 12 2022.
Md Ashiqur Rahman, Zachary Ross, and Kamyar Azizzadenesheli. U-NO: U-shaped neural op- erators. Transactions on Machine Learning Research, 2023. ISSN 2835-8856. URL https: / //openreview. net. forts/foummico d= 3oQF 9coJd
Samuel M Rankin. Semilinear evolution equations in Banach spaces with application to parabolic partial differential equations. Trrassactions of the American Mathematical Society, 336(2):523- 535, 1993.
Guy Revach, Nir Shlezinger, Xiaoyong Ni, Adria Lopez Escoriza, Ruud Van Sloun, and Yonina Eldar. KalmanNet: Neural network aided Kalman filtering for partially known dynamics. IEEE Transactions on Signal Processing, 70:1532-1547, 2022.
Michael Rotman, Amit Dekel, Ran Ilan Ber, Lior Wolf, and Yaron Oz. Semi-supervised learning of partial differential operators and dynamical flows. In Uncertainty in Artificial Intelligence, pp- 1785-1794. PMLR, 2023.
Cristopher Salvi, Maud Lemercier, and Andris Gerasimovies. Neural stochastic PDEs: invariant learning of continuous spatiotemporal dynamics. Advances in Neural Information Pro- cessing Systems, 35:1333-1344, 2022.
Simo Särkkä and Lennart Svensson. Bayesian filtering and smoothing, volume 17. Cambridge university press, 2023.
Xingjian Shi, Zhourong Chen, Hao Wang, Dit-Yan Yeung, Wai-Kin Wong, and Wang-chun Woo. Convolutional LSTM network: A machine learning approach for precipitation nowcasting. Ad- vances in neural information processing systems, 28, 2015.
Ashutosh Singh, Christiana Westlin, Hedwig Eisenbarth, Elizabeth A Reynolds Losin, Jessica R Andrews-Hanna, Tor Wager, Ajay B Satpute, Lisa Feldman Barrett, Dana H Brooks, and Deniz Erdogmus. Variation is the norm: Brain state dynamics evoked by emotional video clips. In 2021 43rd Annual International Conference of the IEEE Engineering in Medicine & Biology Society (EMBC), pp. 6003-6007. IEEE, 2021.
Gerald Smith, Stanley Schmidt, and Leonard McGee, Application of statistical filter theory to the optimal estimation of position and velocity board circumlunar vehicle, volume 135. National Acronautics and Space Administration, 1962.
Stefanie Sonner. An introduction partial differential equations. Radboud Universiy, Nijmegen, IMAPP- Mathematics, 2022.
Roger Temam. Behaviour time t=0 of the solutions of semi-linear evolution equations. Journal of Differential Equations, 43(1):73-92, 1982.
Bilal Thonnam Thodi, Sai Venkata Ramana Ambadipudi, and Saif Eddin Jabari. Fourier neural operator for learning solutions to macroscopic traffic flow models: Application to the forward and inverse problems. arXiv preprint arXiv:2308.0705/. 2023.
Majud Wazwaz. Partial differential equations and solitary waves theory. Springer Science & Business Media, 2010.
12

---

Published as a conference paper at ICLR 2024
Giulio Weikmann, Claudia Paris, and Lorenzo Bruzzone. Timesen2crop: A million labeled samples dataset of Sentinel 2 image time series for crop-type classification. IEEE Journal of Selected Topics in Applied Earth Observations and Remote Sensing, 14:4699-4700 2021.
Fred Weissler. Semilinear evolution equations in Banach spaces. Journal of Functional Analysis, 32(3):277-296 1979.
Gege Wen, Zongyi Li, Kamyar Azizzadenesheli, Anima Anandkumar, and Sally M Benson. U- FNO-An An enhanced fourier neural operator-based deep-learning model for multiphase flow, Ad- vances in Water Resources, 163:104180, 2022.
APPENDIX
A DERIVATION OF EQUATION (12)
First, let that as explained in (Afshar al. 2023) equation (11) is solution the observer
system in equation (10) when (t) is strongly continuous, due to proposition 2.4 in (Afshar et al., 2023 To improve clarity, we have made precise the need for the strong continuity of (t) for the solution in equation (11) to be valid right before equation (11) in the revised paper.
Equation (12) consists of a discreticetinnnionnn of equation (11). It can be derived by using the semigroup property of (()) (Curtain & Zwart 2020), which satisfies
T(t + +') = (t))
for any pair t, t'. This way, identifying in equation (11) with the discrete time t5 we can express a(k)) as:
a(te) Tlx) + 1-' T(t ()(((())) (()))) - C=(+))
(17)
T(h t- 4a 14-1080 (// T(t 4- +4 (((((s),))) (((s))))))((())(())))))) CCas) ds
** / T(t 5) ((B()), + ( - [[(()(s))C((s]])) ds
(18)
T(h 44- 1I-1)T((R 1)00 T(l 4-1) J. th T(k=11 T(4-11-)) (C(((s)))) -Cte)))( K()lv(s)))))))) C+(ss)
th T(4 - 3) (()()), + -Ks))l(s)-
(19)
T(k 1e-1) T( + (' T( - -)
(c(((s)))) + ((>)(o(+)) C=(+)] ds
=dith 1)
+f"h T(ta -))((()), s)ds f T(t. - -)) Ca(a)]ds,
(20)
where we used the linearity of T()). This now in the same form as equation (12).
B IMPLEMENTATION DETAILS AND THE TRAINING PROCEDURE
For all the experiments. for the training of NODA we optimized the loss function (()) using the Adam (Kingma & Ba 2014) optimizer, using learning rate of 10-4 with multiplicative learning rate scheduler which decays 0.5 every 50 epochs. We train the model for a total of 300 epochs, using batch size of 32. The weighting parameter À was fixed as 0.5. The backbone of NODA uses an FNO block for W. For details on the FNO architecture, please see (Li al. 2020a). The total number of parameters in the model with a width of 64 and 20 frequencies per channel. The increase in learnalle parameters for NODA compared to the FNO only depends the choices made for the additional parameters FF, W W..
comparable that used in the original FNO work,
13

---

perturbation. We gen- Published as a conference paper at ICLR 2024
b} and possibly C+ In the current setting, yy W. are single linear layers, b is bias, and choose C* to be the identity matrix since C' was known a priori and equal to the identity operator in the experiments. To parametrize E in a way that is flexible enough to approximate the measurement process for different data acquisition models, we considered a 2-layer fully connected neural network with ReLU activation function. Thus, in the current setting of NODA, its additional parameters constitute only light increase to the number of learnable parameters in the base FNO architecture. Hence, the training time of NODA remains comparable to the original FNO (2D+time). For the experiments on the KS equation, we chose a ID FNO as our base model, composed of four Fourier layers, whereas for the experiments Fourier layers. To implement the MNO (Li et al. 2022), the FNO (Li et al., 2020a) the Conv- LSTM (Shi et al., 2015), and the LSTM (Chung et al. 2014) baselines, considered the setup described in the original papers.
the NS equation, we considered 2D FNO, also with four
For the NS equation, we train NODA using the first 500 time samples of the 180 trajectories con- tained in the training dataset. We choose Itrain = 300 for training of the model on the NS dataset. This allows the model to use both the prediction and the correction terms of the loss function in the first 300 samples, while only using the prediction part for the remaining 200 snapshots. As for the KS dataset, we selected strain as 500. We trained NODA on 1000 trajectories contained in the training dataset, each containing a total of 800 snapshots. We generate the ground truth trajecto- ries for the KS equation using the exponential time-differencing fourth-order Runge-Kutta method from (Kassam & Trefethen 2005). To generate the trajectories for the NS equation, we used the Crank-Nicolson scheme as described (Li al. 2020a)..
All the experiments Nvidia Tesla V100 GPUs Codes used to implement NODA can be found at: http::::p//// / com/ com/singhl7ar www.sngghhhsshhaaoam www.shhwwoooaNNNaaa-aaaaa -pprratt--pperttrrrr -withhhhhwwwwhhhhhhhh--h Data Assimilation
performed computation nodes with Intel Xeon Gold 6132 CPUs and
C RESULTS FOR THE KORTEWEG-DE VRIES (KDV) EQUATION
Korteweg-de Vries (KdV) equation: The KdV PDE is given by
0=
0z
8²=
(21)
0t
Or
83
with =(1) E [0, J]. where = 128, and F is a nonlinear operator. periodic boundary condition is considered for this equation, and initial conditions =(0) 20. This is third order semilinear PDE (Sonner 2022). For more details on this PDE, please refer to (Wazwaz) 2010). This equation represents the propagation of waves on fluid surface when subjected
erated 1200 trajectories from random initializations 2g sampled as described in (Brandstetter et al.] 2022b), containing I E {100, 150,200) seconds each, with timestep of h 0.5 seconds. The first 1000 trajectories were used for training NODA and the competing methods, while the remain- ing 200 were used for evaluating their performance. The resolution was fixed 128 samples. To test NODA we used a warm-up period of tu = 40 seconds.
Benchmark Methods: For the KdV equation, we compare NODA with the FNO Li ct al. 2020a) and with the GRU (Chung et al. 2014) methods.
The average ReIMSE results of all methods for prediction (a 0%) are shown in Table 4 where one can see clear improvement on the performance of NODA when compared to the benchmark algorithms. Regarding the data assimilation performance, Figure O shows the Average ReIMSE for different values of L and for different values of a. The results corroborate the findings stated in the main manuscript with the amount of available data directly impacting NODA's performance.
Plots containing samples of the ground truth solution, of NODA'S estimation, and of elementwise estimation errors are shown in Figure both with (a = 30%) and without (a = 0%) data assim- ilation. These results show that the estimation errors increase only moderately over time, and that data assimilation can significantly improve the estimation results. Different samples of the ground truth trajectories and of the corresponding predictions by NODA for a 0% are shown in Fig- ure 6 where it is shown that NODA provides accurate predictions of the solution for different initial conditions.
14

---

Published as a conference paper at ICLR 2024
10
Average ReIMSE
10
a 0%
10%
a 20%
30%
10-4
Figure 4: Average ReIMSE as function of ty for different values of a for the KdV equation, where warm-up period of tg 40 being used for NODA.
Table Averaged ReIMSE (x 10 for prediction (a 0%) on the Korteweg-de Vries equation as function of the sequence length ( and of the SNR.
SNR
20 dB
30 dB
ts
100
150
200
100
150
200
100
150
200
FNO
21
129
164
17
95
134
13
79
112
MWNO
18
107
149
15
88
114
10
59
92
GRU
207
296
421
190
219
258
122
178
254
NODA
14
89
102
12
92
9
48
D RESULTS DIFFERENT MEASUREMENT MODELS
In this section we evaluate the performance of NODA for two different choices of the discretized measurement operator Ć. We present results for the NS example (R, = 40, by = 1000, SNR 30dB), where the first choice of measurement operator is Ó (the identity operator, also consid- ered in the main body of the paper), and Crand, which is full matrix generated randomly, with each element sampled uniformly in the interval Gij E [0, 1]. In Table 5 we present the averaged ReIMSE results for the two choices of the measurement operator and for different choices of It can be seen that the example with the random operator Crand is noticeably more challenging, as evidenced by the consistently higher averaged ReIMSE values. Nevertheless, the increase in the amount of data available for assimilation (i.e., or) consistently increases the performance of NODA for both choices of measurement operator.
Table Averaged ReIMSE (x10²) of NODA for the Navier-Stokes equation as a function of for two different discretized measurement operators Ć.
0%
10%
20%
30%
Č =
18
13
Crand
42
39
33
25
15

---

Published as a conference paper at ICLR 2024
64
0.15
0.10
40 80
120
160 200
0.05
0.00
Figure 5: Sample of the predictionresults for the Korteweg-de Vries equation with L 200 seconds.
Ground truth trajectory (Leff Image). Solution estimated by NODA with a = 0% (Top left). Solution estimated by NODA with 30% (Top Right). Elementwise error plot for NODA's estimate without data assimilation (Bottom Left). Elementwise error plot for NODA's estimate with data assimilation (Bottom Right).
Table 6: Averaged ReIMSE (x 103) of NODA, MNO and FNO for the Navier-Stokes equation function of for Re E {500, 5000].
α
0%
10%
20%
30%
Re
500
196
179
152
145
NODA
5000
258
226
202
189
500
183
MNO
5000
221
500
247
FNO
5000
310
E RESULTS FOR THE NAVIER-STOKES EQUATION WITH DIFFERENT VORTICITIES
In this section, we present results for NS equation with Re E (500, 5000). corresponding flows with higher turbulence. We compare the results of NODA, MNO and FNO. For NODA we show the results for different values of a E (0%, 10% 20%, 30%). while for the remaining methods only 0% is shown as they do not perform data assimilation. For both choices of Re, we considered and ty 500. The averaged ReIMSE results are shown Table From these results, it can be seen that when
SNR 30dB, tи
0% (for prediction only), NODA outperforms INO and has performance that is competitive with MNO, which provided the lowest average ReIMSE values. However, more data is used for assimilation (particularly for E {20%, 30%)). the average RelMSE of NODA decreases significantly, becoming lower than that those provided by both MNO and FNO. This highlights the benefits of data assimilation in the performance of NODA.
16

---

Published as a conference paper at ICLR 2024
Figure 6: Sample of the prediction results for the Korteweg-de Vries equation with L = 200 sec- onds. Three different realisations of the ground truth trajectory generated from different initial conditions (Top Row) Corresponding trajectories predicted by NODA (Bottom Row)
F ADDITIONAL RESULTS ON THE NAVIER-STOKES AND KURAMOTO-SIVASHIISEE EQUATIONS
this section, we present complementary results for the KS and NS equations. These consist of visualizations of samples of the ground truth and reconstructed trajectories plots for predictions of NODA (using a 0%). These can be seen in Figures and
17

---

Published as a conference paper at ICLR 2024
Figure 7: Different realizations of ground truth trajectories of the Suramoto-Sivasinnse equation with 4 200 seconds (Top Row), and the corresponding predictions obtained by NODA without data assimilation (Bottom Row).
Ground Truth
Pred.
t-150
t=300
t=450
Figure 8: Sample of the ground truth sequence (Top Row) and predictions provided by NODA without data assimilation (Bottom Row) for the Navier-Stokes, equation for f (150, 300, 450}.
18


Reviews:
[{"review_id": "EhSvltJICs", "reviewer": "Reviewer_ozZs", "summary": "The paper presents NODA, a novel means of solving a particular class of PDEs with Neural Operators.\nBy exploiting the structure of semilinear PDEs, NODA benefits from increased accuracy on suitable tasks, and the important ability to incorporate noisy data to assist in making its predictions. Experiments are run on two different PDEs, Kuramoto-Sivashinsky and Navier-Stokes, in one and two spatial dimensions respectively.", "questions": "My suggestions are linked to the weaknesses highlighted above:\n\n- Including timings for the experiments, to show how fast the algorithm is compared to competing methods. \n\n- To introduce more variety in the experiments. Would it be impractical to explore higher Reynolds numbers, and do we expect NODA to cope less well with turbulent flow?\n\n- I imagine the uncertainties in the numerical experiments are quite small but it would still be a strong recommendation to include them.\n\n- Either some justification needs to be added for the particular choices of the benchmark methods, or more suitable alternatives should be introduced in their place.\n\n- Figures 1 and 2 are lacking a colour scale, making it challenging for the reader to assess the nature of the uncertainties. The reader is left to guess values based upon the Table but it would be preferable to be explicit and quantitative. \n\nOverall I would stress that I think the paper is very promising, so I hope these suggestions help the authors to address my concerns.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "Originality: NODA is, to the best of my knowledge, a highly novel approach to leveraging the structure of semilinear PDEs to create more powerful Neural Operators. \n\nMotivation: The authors make a strong case for NODA's theoretical motivation, by leveraging the structure of semilinear PDEs, and introducing an interesting prediction/correction mechanism.\n\nClarity: The paper is well written, with a clear exposition of the concepts and methodology.\n\nSignificance: This is certainly a highly impactful area of machine learning, one which could lead to improvements in the speed and accuracy of predicting complex dynamical systems.", "weaknesses": "- The abstract claims that the proposed method makes fast predictions, and has “little computational overhead”. Yet no experimental evidence is provided to support this statement. It’s important the authors rectify this by including quantitative timing metrics for the experiments, and compare it against the other methods presented.\n\n- The selection of experiments is very limited – only looking at one 1D and one 2D PDE. It would be advisable to at least present results for more than one Reynolds number.\n\n- The  experimental results are well presented, but the tables lack uncertainty estimates. The reader is therefore left unable to gauge whether the performance differences between different methods are statistically significant.\n\n- While there are a good number of alternative methods in the benchmarks, I’m not sure they were the most appropriate choices. The two LSTM-based models are taken from papers that are over eight years old. Meanwhile the MNO is specifically designed for dissipative dynamics, yet the Navier-Stokes experiment is chosen to be in the non-dissipative regime. Furthermore, none of the other methods selected for benchmarking are able to make use of the noisy data, so we are unable to test the key feature of NODA. Alternative works which may offer more suitable benchmarks include ‘Approximate Bayesian Neural Operators’ and ‘Multiwavelet-based Operator Learning for Differential Equations’.", "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": "aaSwX3AwLY", "reviewer": "Reviewer_2zWc", "summary": "The paper introduces a learning-based state-space approach to address the challenges associated with solving complex systems governed by spatio-temporal Partial Differential Equations (PDEs) over long time scales. It highlights the limitations of existing Neural Operators (NOs) when dealing with data assimilation and correction operations based on noisy measurements. The proposed framework leverages the structure of semilinear PDEs and nonlinear observers to develop a recursive method that combines prediction and correction operations efficiently. The paper presents promising results from experiments on various PDEs, showcasing the model's robustness and ability to correct predictions with irregularly sampled noisy measurements.", "questions": "1. Can you elaborate on the specific scenarios or application domains where the proposed learning-based state-space approach is expected to have the most significant impact?\n\n2. How does the model handle variations in the amount of available measurements, and are there limitations in terms of the minimum number of measurements required for effective data assimilation?\n\n3. Could you provide more details about the computational resources needed for implementing and running the proposed framework in practical applications, particularly in scenarios involving large-scale systems?\n\n4. In the context of the experiments conducted, are there specific parameters or settings that were found to be critical for achieving the robustness of the model to noise?\n\n5. Are there plans or ongoing research aimed at simplifying the implementation and improving the user-friendliness of the proposed approach for researchers and practitioners in related fields?\n\nThese questions aim to gain further insights into the practical applicability and potential impact of the proposed learning-based state-space approach for solving complex systems described by spatio-temporal PDEs.", "limitations": "", "rating": 5, "confidence": 3, "soundness": 1, "presentation": 1, "contribution": 1, "strengths": "1. **Innovative Approach**: The paper presents a novel learning-based state-space approach to address a critical issue in the context of solving spatio-temporal PDEs over long time scales, providing a fresh perspective on solving complex problems in science and engineering.\n\n2. **Efficiency and Accuracy**: The proposed framework aims to produce fast and accurate predictions while handling irregularly sampled noisy measurements, enhancing the accuracy of PDE solutions.\n\n3. **Real-World Relevance**: The need for data assimilation and correction in dynamical systems, as discussed in the paper, has significant real-world relevance, particularly in fields such as Earth science, remote sensing, traffic analysis, and medical imaging.\n\n4. **Robustness to Noise**: The experiments demonstrate the model's robustness to noise, which is a crucial consideration in real-world applications where measurements are often affected by noise and uncertainties.", "weaknesses": "1. **Complexity**: While the proposed approach appears promising, the complexity of the model and the methods discussed could make implementation and practical application challenging for researchers and practitioners without expertise in this specific field.\n\n2. **Limited Application Scenarios**: The paper primarily focuses on solving PDEs and addressing issues related to data assimilation and correction. It would be beneficial to discuss broader applications and practical scenarios where this approach could have a significant impact.\n\n3. **Computational Overhead**: Although the paper suggests that the proposed model has little computational overhead, it would be valuable to provide more specific information regarding the computational resources required for practical implementation.", "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": "d1g8dHv2jh", "reviewer": "Reviewer_M6oo", "summary": "Acknowledging the recent limitations in current Neural Operators, namely, data assimilation and lack of forecast error correction, authors propose learning semilinear operators instead, for which a body of mathematics literature exists already and allows for data assimilation and error correction through the use of the *observer* operator. Another feature of this framework is the fact that it assumes that the measurements are inherently noisy and also assumes in general that the dynamics themselves could be noisy. The authors base their architectural design on the semilinear operator by proposing a time-discretization that they later parametrize using Fourier Neural Operators (FNO). An essential facet of their research lies in their emphasis on learning correction operators, and this is substantiated by comprehensive ablation studies that illustrate their efficacy. In summary, the method presented in this work yields notably robust outcomes, particularly in the context of the considered datasets.\n\nThe first weakness that I pointed to in the weakness section is a huge one in my opinion and is reflected in my score. I am ready to change my score if that weakness is addressed since I think it will make the whole paper much better.", "questions": "- Is it possible to have a clear demonstration of why Navier-Stokes and KdV equations are semi-linear PDEs?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Well-written paper and method is well motivated with respect to the literature.\n- Considers the snapshots as inherently noisy instead of assuming their ground-truth nature.\n- Proposed framework is baked in sound mathematical theory which is reflected in the design choice of the architecture, furthemore, it proposes a way of mitigating the long-standing challenge of long-term forecast error.\n- Proposed framework beats the baselines when $\\alpha=0$ suggesting that even in the prediction-only phase (Tables 1, 2 and 4), the architecture is well justified. Moreover, the authors demonstrate the soundness of the correction operator through (Table 3, Figures 3 and 4) in which they show that the error decreases as more updates steps are considered.", "weaknesses": "- Unclear how Eq. 12 was derived. I tried going through the effort of deriving it myself but couldn't get the same results, so it's critical that this step is justified (you can include the proof in the appendix if it burdens the main text) but a proof must be provided since the whole method is based on it.\n- This is a \"weak\" weakness but given that the setting considered in the paper is that of $y(t)$ being just noise added to $z(t)$, it seems perhaps too much trying to learn a linear operator ($I_d$ in this case) using a non-linear neural net. While I understand that this was done to accomodate general case scenarios, it would be nice to see some ablation experiments showing the superiority of using a neural net for learning the operator $C$ as compared to using either: A learnable projection matrix or using not learning it at all and assuming $C=I_d$. Also, having an additional dataset where $C$ is not trivial could help motivate the choice of the architecture for it.\n- In the main text and the paper from Afshar et al. (2022), The *observer gain* $K(t)$ needs to satisfy the condition that in the absence of noise from the system dynamics $z(t)$ and noise from the measurements $y(t)$, $\\hat{z}(t)\\rightarrow z(t)$, yet nothing is said about that when the architecture for it was introduced in Eq. 15. Perhaps a pretraining period where no noise is introduced in the samples could be used to make sure that the learned $K$ satisfies the aforementioned condition (at least approximately) and then finetuning further using your training scheme.", "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": "aSej2N45jY", "reviewer": "Reviewer_o6QX", "summary": "The paper proposes the NODA method for learning semilinear neural operators that aloows both prediction and data assimilation. The method is motivated by the observer operator gain in [1], by which the evolution of the estimates are designed to consist of two terms, one for prediction and one for correction. Experiments on  the Kuramoto-Sivashinsky, NS, and Korteweg-de Vries equations are conducted to demonstrate the performance.\n\n[1] Afshar, Sepideh, Fabian Germ, and Kirsten Morris. \"Extended Kalman filter based observer design for semilinear infinite-dimensional systems.\" IEEE Transactions on Automatic Control (2023).", "questions": "See weaknesses.", "limitations": "", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- The motivation is clear and applying data assimilation to neural operator learning is very relevant.\n- The paper is well-written and easy-to-follow. The derivation of the method is intriguing.\n- The integration of the observer design of semilinear PDEs seems novel.\n- Strong experimental results.", "weaknesses": "- Lack of explanation of the observer design of semilinear PDEs. On page 5, the authors claim that under “mild” additional conditions, the solution converges to that of the real solution. In the experiments, the measurements are obtained by injecting Gaussian noises to the real solutions. It might be beneficial to provide some details of the claim so that readers may learn about in which data assimilation scenario this framework may fit.\n- While the adoption of the FNO architecture is reasonable for the prediction term, it is not clear why the observer gain $K$ should/ are sufficient to be parametrized as in (15). A further explanation of the intuition here might be necessary.\n- As for the experiments, it seems only the proposed NODA method needs a warmup. While I totally understand the methodology here and the need for the warmup, I just wonder if this would cause unfair comparison with other methods that have no access to this part of the data and thus less strong performance.\n- It might be better to include colorbars for all the figures, not only for one in Figure 5.", "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 (poster)