Learning Space-Time Continuous Latent Neural PDEs from Partially Observed States
Authors: Valerii Iakovlev
Markus Heinonen
Harri Lähdesmäki
Venue: NeurIPS 2023 poster (2023)
Keywords: neural
PDEs
neural PDEs
partial observations
space time continuous

Abstract: We introduce a novel grid-independent model for learning partial differential equations (PDEs) from noisy and partial observations on irregular spatiotemporal grids. We propose a space-time continuous latent neural PDE model with an efficient probabilistic framework and a novel encoder design for improved data efficiency and grid independence. The latent state dynamics are governed by a PDE model that combines the collocation method and the method of lines. We employ amortized variational inference for approximate posterior estimation and utilize a multiple shooting technique for enhanced training speed and stability. Our model demonstrates state-of-the-art performance on complex synthetic and real-world datasets, overcoming limitations of previous approaches and effectively handling partially-observed data. The proposed model outperforms recent methods, showing its potential to advance data-driven PDE modeling and enabling robust, grid-independent modeling of complex partially-observed dynamic processes across various domains.

Learning Space-Time Continuous Neural PDEs from Partially Observed States
Valerii Iakovlev
Markus Heinonen Harri Lähdesmäki
Department of Computer Science, Aalto University, Finland
(valerii iakovlev, markus. heinonen, harri lahdesmaki) eaaltto...aaaa...........o fa
Abstract
We introduce a novel grid-independent model for learning partial differential equations (PDEs) from noisy and partial observations on irregular spatiotemporal grids. We propose space-time continuous latent neural PDE model with an efficient probabilistic framework and novel encoder design for improved data efficiency and grid independence. The latent state dynamics are governed by a PDE model that combines the collocation method and the method of lines. We employ amortized variational inference for approximate posterior estimation and utilize a multiple shooting technique for enhanced training speed and stability. Our model demonstrates state-of-the-art performance complex synthetic and real-world datasets, overcoming limitations of previous approaches and effectively handling partially-obsery data. The proposed model outperforms recent methods, showing its potential to advance data-driven PDE modeling and enabling robust, trid-independent modeling of complex partially-observe dynamic processes.
1 Introduction
Modeling spatiotemporal processes allows to understand and predict the behavior of complex systems that evolve over time and space (Cressie and Wikle, 201 Partial differential equations (PDEs) are a popular tool for this task they have solid mathematical foundation (Evans, 2010) and can describe the dynamics of wide range physical, biological, and social phenomena (Murray, 2002; Hirsch, 2007). However, deriving PDEs can be challenging, especially when the system's 's underlying mechanisms are complex and not well understood. Data-driven methods can bypass these challenges (Brunton and Kutz, 2019). By learning the underlying system dynamics directly from data, can develop accurate PDE models that capture the essential features of the system. This approach has changed ability to model complex systems and make predictions about their behavior in data-driven manner.
While current data-dríven PDE models have been successful modeling complex spatiotemporal phenomena, they often operate under various simplifying assumptions such as regularity the spatial or temporal grids (Long et al., 2018; Kochkov et al., 2021; Pfaff et al., 2021; Li et al., 2021; Han et al., 2022; Poli et al., 2022), discreteness in space or time (Sco et al., 2020; Pfaff et al., 2021; Lienen and Günnemann, 2022: Brandstetter et al., 2022), and availability of complete and noiseless observations (Long et al., 2018; Pfaff et al., 2021; Wu et al., 2022). Such assumptions become increasingly limiting in more realistic scenarios with searce data and irregularly spaced, noisy and partial observations.
We address the limitations of existing methods and propose space-time continuous and grid- independent model that can learn PDE dynamics from noisy and partial observations made on irregular spatiotemporal grids. Our main contributions include:
Source code and datasets can found our github repository.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Development of an efficient generative modeling framework for learning latent neural PDE data;
models from noisy and partially-o
Novel PDE model that merges two PDE solution techniques the collocation method and the method of lines achieve space-time continuity, grid-independence and data efficiency;
Novel encoder design that operates on local spatiotemporal neighborhoods for improved data data-efficiency and grid indeeeedeedenee.
Our model demonstrates state-of-the-art performance on complex synthetic and real-world datasets, opening up the possibility for accurate and efficient modeling of complex dynamic processes and promoting further advancements data-driven PDE modeling.
2 Problem Setup
In this work we are concerned with modeling of spatiotemporal processes. For brevity, we present method for single observed trajectory, but extension to multiple trajectories is straightforward. We observe a spatiotemporal dynamical system evolving over time on a spatial domain Q. The observations made at M arbitrary consecutive time points 11.M (f1..... tM) and N arbitrary observation locations T1:N = (æ1,.... 2N), where F; E Q. This generates a sequence obser- vations U1:M (u1...., U.y) ) where u, IR NxD contains D-dimensional observations at the N observation locations. We define ui as the observation time t, and location æy- The number of time points and observation locations may vary between different observed trajectories.
We assume the data is generated by dynamical system with latent state z(t,, x)) E RRd, where t is time and x E Ω is spatial location. The latent state is governed by an unknown PDE and is mapped to the observed state utt, E R by an unknown observation function g and likelihood model p: @z x) Ot
F (z 2), (t, x), (t.
(1)
utt, x) ~ p(1(2(1, x)));
where 2t, æ) denotes partial derivatives wrt x.
In this work we make two assumptions that are highly relevant in real-world scenarios. First, we assume partial observations, that is, the observed state u(t, x) does not contain all information about the latent state (t, 2) (e.g., (t,, x) contains pressure and velocity, but u(t, x) contains information only about the pressure). Second, we assume outt-f-ddistiiutiinn time points and observation locations, that is, their number, positions, and density can change arbitrarily at test time.
3 Model
Here describe the model components (Sec. 3.1) which are then used to construct the generative model (Sec. 3.2),
dr Po...)
x
z(t,, σ)
z(tz, x)
(t,,
3.1 Model components
Sham
Our model consists of four parts: space-time continuous latent state z(t, x) and observed state (t,, x), dynamics function Fe... governing the temporal evolution of the latent state, and an observation function g0., mapping the latent state to the observed state (see Figure 1). Next, we describe these components in detail.
(1 σ
uut2, x))
TA (fa, 2
Figure 1: Model sketch. Initial latent state 2(1,, æ) is evolved via Fe,, to the following latent states which then mapped to the observed states by gaw.
Latent state. To define a space-time continuous latent state z(t, æ) E Rd, we introduce x(t) := (z (t),.., (,.).,,)))) E R xd, where each 2²()) RRd corresponds to the observation location Z;- Then, we define the latent state z(t, æ)) spatial interpolant of z(t):
2(1, ar) Interpolate(((())())
where Interpolate(-) maps z(t) to an interpolant which can be evaluated at any spatial location E Ω (see Figure 2). We do not rely on particular interpolation method, but in this work we use linear interpolation as it shows good performance and facilitates efficient implementation.

---

Latent state dynamics. Given a space-time continuous latent state, one can naturally define its dynamics in terms of PDE:
(+)
z (t)
(z(1, x) Ot
z (t)
(.) (2= ..
(4)
z (t)
023
22
where Fe.. is dynamics function with parameters Oäyn. This is viable approach known the collocation method (Kansa, 1990;
Cheng, 2009), but it has several limitations. It requires us to de- cide which partial derivatives to include in the dynamics function, and also requires an interpolant which has all the selected partial derivatives (c.g., linear interpolant has only first order derivatives). To avoid these limitations, we combine the collocation method with another PDE solution technique known as the method of lines (Schiesser, 1991; Hamdi et al., 2007), which approximates spatial derivatives 0² 2) using only evaluations of z (tt, x)), and then let the dynamics function approximate all required derivatives in data-driven manner. To do that, we define the spatial neighborhood of æ (x) which is set containing and its spatial neighbors, and also define =(1,Ass(<:)) (2)). which is a set of evaluations of the interpolant = (1, a:) at points (r)):):
x2
z(t, 2))
Figure 2: Latent state z(t, x)
defined as an interpolant of
=(t) == (z (t), 2 (t)).
(m) == (0' E Q x' or x' is spatial neighbor of x), (t, Ns())) = {z(t, x²) x' E (N)
(5)
(6)
and assume that this information is sufficient to approximate all required spatial derivatives at æ. This is reasonable assumption since, e.g., finite differences can approximate derivatives using only function values and locations of the evaluation points. Hence, we define the dynamies of z(1, a) as
0z(1, x)) Feeg Ōt
which is defined only in terms of the values of the latent state, but not its spatial derivatives.
One way define the spatial neighbors for is in terms of the observa- tion locations I1.N (e.g.. use the nearest ones) as was done, for example, in (Long et al., 2018; Pfaff et al., 2021; Lienen and Günnemann, 2022). Instead, we utilize continuity of the latent state (1,, x) and define the spa- tial neighbors in a grid-independent manner fixed number of points arranged in a predefined patter around æ (see Figure 3). This allows to fix the shape and size of the spatial neighborhoods in advance, making them independent of the observation locations. In this work we use the spatial neighborhood consisting of two concentric circles of radius r and r/2, each circle contains 8 evaluation points as in Figure 3. In Appendix we compare neighborhoods of various shapes and sizes.
x
x
Figure Example of Ns (mz). Instead of us- ing the observation lo- cations (dots) to define spatial neighbors, we use spatial locations ar- ranged in fixed prede- fined pattern (crosses).
Equation 7 allows to simulate the temporal evolution of z(t, at any spatial location. However, since z(1, .) is defined only in terms of spatial interpolant of ztt) (see Eq. 3), with 2²(t) = 2t,, a.:). it is sufficient to simulate the latent state dynamics only at the observation locations Ø1.N- Hence, we can completely characterize the latent state dynamics in terms of system of N ODEs:
de (t)
(z((,2)) St
Fegm
dz(t)
dt
(8)
dt
@z(t,ax) Dt
Fege
dt
For convenience, we define 2(t; 11:1112.0 21, (ayn)) ODESolvo(t: f1, 21, 0yyn) the solution of the ODE system in Equation 8 at time t with initial state xx(1) = 21 and parameters Guyn- We also define (1,, a;; L1, 21, Oayn) as the spatial interpolant of 2(1; L 21, Oayn) as in Equation 3. We solve the ODEs using off the shelf differentiable ODE solvers from torchdiffeq package (Chen, 201 8). Note that we solve for the state 2(l) only at the observation locations I1.N so to get the neighborhood values x (s(x.) we perform interpolation at every step of the ODE solver.
3

---

We define the mapping from the latent space to the observation space as parametric function gi with parameters R
Observation function.
u(1, x)) (go (z(1, (1)), TDD,)
where N is the Gaussian distribution, o is noise variance, and Ip is D-by-D identity matrix.
3.2 Generative model
Training models of dynamic systems is often challenging due to long training times and training instabilities (Ribeiro et al., 2020; Metz et al., 2021). To alleviate these problems, various heuristics have been proposed, such progressive lengthening and splitting of the training trajectories (Yildiz et al., 2019; Um et al., 2020). We use multiple shooting (Bock and Plitt, 1984; Voss al.. 2004), simple and efficient technique which has demonstrated its effectiveness ODE learning applications (Jordana et al., 2021; Hegde al., 2022). We extent the multiple shooting framework for latent ODE models PDE modeling setup by introducing spatial dimensions in the latent state and designing an encoder adapted specifically to the PDE setting (Section 4.2).
1
82
f2
13
14
[0]
1(2)
b=1
b=2
presented in (lakovlev al., 2023)
Iy= [1,21,,3)
(1,5)
Figure 4: Multiple shooting splits trajectory with one ini- tial state (top) into two sub- trajectories with two initial states (bottom) and tries to minimize the gap between sub- trajectories (orange arrow).
Multiple shooting splits a single trajectory (=(t)))=I....d with initial state 21 into B consecutive non-overlapping sub- trajectories (=(t,)
= 1....... with initial states $1.B == (5 (S1....,81) while imposing continuity penalty between the sub- trajectories (see Figure 4). The index set I, contains time point indices for the b'th sub-trajectory. We also denote the temporal posi- tion 8, as 110) and place 8, at the first time point preceding the b'th sub-trajectory (except 8 which placed at t1). Note that the shooting states 8,, have the same dimension as the original latent state =()) i.c., 8, E IR xd Multiple shooting allows to parallelize the simulation over the sub-trajectories and shortens the simulation intervals thus improving the training speed and stability. Appendix D we demonstrate the effect of multiple shooting on the model training and prediction accuracy.
We begin by defining the prior over the unknown model parameters and initial states:
p(81.8, Oayn: Osec) = p( p(81.1 (p (0ay)) )((ee))))
(10)
where p(Oay)) and p(Oacc) are zero-mean diagonal Gaussians, and the continuity inducing prior 11 (dyn) is defined as in (lakovlev al., 2023)
p($1. (dya)) = p(s)) II bm2
p(s) $6-1, Oaya).
(11)
Intuitively, the continuity prior (s 9dyn) takes the initial latent state $6-1, simulates it forward from time ti--i to fip to get [b = ODESolve(tm) f/b- 1]: Sh Oayy), and then forces μ4 to approximately match the initial state 8; of the next sub-trajectory, thus promoting continuity of the full trajectory. We assume the continuity inducing prior factorizes across the grid points, i.e..
(0ya))
II p(s{) j=1
II b-2j=1
p(s 16- Adyn)
(12)
II
p(s
II N ( [s]l=(t()) Zz tz:tto-n sb-1, Ooyn), o2.d)
(13)
where p(s() is a diagonal Gaussian, and parameter o² controls the strength of the prior. Note that the term 2(1ny Zyi tp- $6-1, Gaya) in Equation 13 equals the ODE forward solution ODESolve(tigj 1|0-1], $b-1, Ooya) at grid location wy.

---

Finally, we define our generative in terms of the following sampling procedure:
O Oayn: Odce $1.B p(Odyn)pl (0 (p)))))))))))))))))))))))))))))) (Oan))
(14)
<(1) = (1,;; [(b) 8b, fyyn);
b (1, ,.... B), E Ib
(15)
un ~ p( ( gom. (z( mz)),
i = 1, M,
N,
(16)
with the following joint distribution (see Appendix A for details about the model specification.):
B
N
pl U1 81.B; Oaya: Odce)
II
1I
p(u! (8. Oaya, doc)
(17)
2 ==1
4 Parameter Inference, Encoder, and Forecasting
4.1 Amortized variational inference
We approximate the true posterior over the model parameters and initial states pl p(81.1 Odyn: Osed (11:M) using variational inference (Blei et al., 2017) with the following approximate posterior:
B
N
a(aayn Odees 81.B) =
= Qahtyo (Oayn Qypan (Odec)
II b=1j=1
Dez (o%)
(18)
where Thaag Thbss and Dug are diagonal Gaussians, and 4/ayn. thete and wag are variational parameters. the local variational parameters v/y
To avoid direct optimization inference (Kingma and Welling, 2013) and train an encoder ho... with parameters Oenc which maps observations U1:M to ubd (see Section 4.2). For brevity, we sometimes omit the dependence of approximate posteriors variational parameters and simply write e.g., q(s))-
amortized variational
In variational inference the best approximation of the posterior obtained by minimizing the Kullback-Leibler divergence: KL [g(Oayn: Oáee $1.B) ())((((( Odee 81.B [u],,]]] which is equivalent to maximizing the evidence lower bound (ELBO), defined for our model as:
= Σ
[lo p(u; (sb, Odyns dasz)
KL(a(s))ll(ss))
Ag( 90ss,00sn: Osc))
1=1 IEI, j=1
( observation model
(ii) initial state prior
B N
BygBea 1b-1)
((s))(())) 86-1, 0dym)]
KL. [g ((50y) (pp ((lp(((a)))
KT. [T((q( (((Odee)))))) (1p((scc)
(iv) dynamics
decoder prior
(Hit) continuity prior
The terms (ii), (év), and (v) are computed analytically, while terms (i) and (1i1) are approximated using Monte Carlo integration for expectations, and numerical ODE solvers for initial value problems. See Appendix and B approximate posterior details and derivation and computation of the ELBO.
4.2 Encoder
Here we describe our encoder which maps observations u1:M to local variational parameters t// required to sample the initial latent state of the sub-trajectory b at time point fp. and observation location x.- Similarly to our model, the encoder should data-efficient and grid-independent.
Similarly to our model (Section 3.1), we enable grid-inddpenddece by making the encoder operate on spatial interpolants of the observations ui.ar (even if they are noisy):
u (a:) := Interpolatc(u.) (2:),
1....,,MM M,
(19)
where spatial interpolation done separately for each time point We then use the interpolants (x(x) to define the spatial neighborhoods (z( in grid-independent manner.
5

---

de
b7
4-2
4--1
flu f2
{wu-1 (x (aelhessste) (ansss) [te-- (ec
Repstial
Repatial
Repasial
lxcempoall
18
Ryead
Figure 5: Spatiotemporal neighborhood of a multiple shooting time point 40j = 4 and location Zj, ultps æı (denoted by green, blue and orange crosses and the dots inside), is mapped to the variational parameters v via the encoder.
To improve data-efficiency, we assume v/? does not depend on the whole observed sequence M1:Ms but only some local information in spatiotemporal neighborhood of [μy and æz. We define the temporal neighhorhood of (1y) as
(t)+ = {k: It, tayl s Er, 1, M).
(20)
where 5T is a hyperparameter controlling the neighborhood size, and then define the spatiotemporal neighborhood of (p and æ, as
u [l æy} = {ux (w) E (T E (xxx))
(21)
Our encoder operates such spatiotemporal neighborhoods ult my] and works in three steps (see Figure 5). First, for each time index k E Nt((tt)) it aggregates the spatial information
{ux (xx) (mEN(()) into a vector a Then, it aggregates the spatial representations a across time into another vector of ai which is finally mapped to the variational parameters vi as follows:
4/. = hom. (u (u(tp), æ,]) Aread (Meemporaal (Asguaiial
[p] az)))))
(22)
Spatial aggregation. Since the spatial neighborhoods are fixed and remain identical for all spatial locations (see Figure 5), implement the spatial aggregation function hspatial an MLP which takes elements of the set (u) =ENs((e) stacked in a fixed order as the input.
Temporal aggregation. We implement Aemporal stack of transformer layers (Vaswani et al., 2017) which allows it to operate on input sets of arbitrary size. We use time-aware attention and continuous relative positional encodings (Iakovlev et al., 2023) which were shown to be effeettvee on data from dynamical systems observed at irregular time intervals. Each transformer layer takes layer-specific input set {Ei IRENT()))) where <in is located at tt. and maps to output set {Sx )keNe(tt))) where each < computed using only the input elements within distance ô7 from te. thus promoting temporal locality. Furthermore, instead of using absolute positional encodings the model assumes the behavior of the system does not depend on time and uses relative temporal distances to inject positional information. The first layer takes {aa keWr((t))))) the input, while the last layer returns single element at time point f1j, which represents the temporal aggregation T "[i)-
Variational parameter readout. Since a. is a fixed-length vector, we implement hraad as MLP.

---

SHALLOW WATER dataset contains observations of the wave height a pool of water. Middle: NAVIER-STOKER dataset contains observations of the concentration of a species transported fluid via buoyancy and velocity field. Bottom: SCALAR FLOW dataset contains observations of smoke plumes raising in warm air.
Figure Top:
4.3 Forecasting
Given initial observations ü1.m time points f1:m, predict the future observation ü, time point ta > fm as the expected value of the approximate posterior predictive distribution:
p Ui:, u1:M)) f p( &m, Oayn. dec ( ( (d,
(23)
The expected value is estimated via Monte Carlo integration (see Appendix C.4 for details).
5 Experiments
use three challenging datasets: SHALLOW WATER, WAVIER-STOKES, and SCALAR FLOW which contain observations of spatiotemporal system at N ss 1100 grid points evolving over time (see Figure 6). The first two datasets are synthetic and generated using numeric PDE solvers (we use scikit- fdiff (Cellier, 2019) for SHALLOW WATER, and PhiFlow (Holl et al., 2020) for NAVIER-STOKES), while the third dataset contains real-world observations (camera images) of smoke plumes raising in warm air (Eckert et al., 2019). In all cases the observations are made at irregular spatiotemporal grids and contain only partial information about the true system state. In particular, for SHAL.LOW WATER observe only the wave height, for HAVIER-STOKE the species, and for SCALAR FL.W only pixel densities are known. All datasets contain 60/20/20 training/validaaiioneessggg trajectories. See Appendix for details.
observe only the concentration of
We train our model for 20k iterations with constant learning rate of 3e-4 and linear warmup. The latent spatiotemporal dynamics are simulated using differentiable ODE solvers from the torchdiffeq package (Chen, 2018) (we use dopri5 with rtol=le-3, atol=le-4, no adjoint). Training is done on a single NVIDIA Tesla V100 GPU, with single run taking 3-4 hours. We use the mean absolute error (MAE) on the test set as the performance measure. Error bars are standard errors over random seeds. For forecasting we use the expected value of the posterior predictive distribution. See Appendix C for all details about the training, validation, and testing setup.
Latent state dimension. Here show the advantage of using latent-space models on partially observed data. We change the latent state dimension from to 5 and measure the test MAE. Note that for d we effectively have data-space model which models the observations without trying to reconstruct the missing states. Figure 7 shows that in all cases there is improvement in performance as the latent dimension grows. For SHALLLOW WATER and NAVIER-STOKES the true latent dimension is 3. Since SCALAR FLOW is a real-world process, there is no true latent dimension. As benchmark, we provide the performance of our model trained on fully-observed versions of the synthetic datasets (we use the same architecture and hyperparameters, but fix d to 3). Figure also shows examples of model predictions (at the final time point) for different values of d. We see huge difference between
and d = 3, 5. Note how apparently small difference in MAE at d = 1 and 5 for SCALAR FLOW corresponds a dramatic improvement in the prediction quality.

---

Shallow Water Navier-Stokes
Scalar" Flow
0.080
Data
0.060
Shallow Water
0.040
0.020
0.000
Test MAE
0.120
d-5
0.090
Stoker-Stokes
0.060
0.030
0.000
0.070
0.060
Scalar Flow
0.050
0.040
0.030
d=1
3
4
5
Latent State Dim.
Figure 7: Left: Test MAE vs latent state dimension Black lines are test MAE on fully-observed versions of the datasets (# standard error). Right: Model predictions for different d.
Grid independence. In this experiment we demon- strate the grid-independence property of our model by training it on grids with & 1100 observation lo- cations, and then testing on a coarser, original, and finer grids. We evaluate the effect of using differ- ent interpolation methods by repeating the experiment with linear, k-nearest neighbors, and inverse distance weighting (IDW) interpolants. For SHALLOW WA- TER and TAVIER-STOKE, the coarser/finer grids con- tain 290/4200 nodes, while for SCALAR FL.OW we have 560/6420 nodes, respectively. Table shows the model's performance for different spatial grids and in- terpolation methods. see that all interpolation meth-
Shallow Water Novier-Stokes
Scolar Flow
Data
Finer
Coarser Original
ods perform rather similarly on the original grid, but linear interpolation and IDW tend to perform better on finer/coarser grids than k-NN. Performance drop on coarse grids is expected since we get less accurate information about the system's 's initial state and simu- late the dynamics on coarse grids. Figure also shows examples of model predictions (at the final time point) for different grid sizes and linear interpolant.
Figure Predictions on spatial grids of dif- ferent density (linear interpolant, test data).
Comparison to other models. We test our model against recent spa- tiotemporal models from the litera- ture: Finite Element Networks (FEN) (Lienen and Günnemann, 2022), Neu- ral Stochastic PDEs (NSPDE) (Salvi al., 2021), MAgNet (Boussif et al., 2022), and DINo (Yin et al., 2023). We also use Neural ODEs (NODE) (Chen et al., 2018) as the baseline. We use the official implementation for all models and tune their hyperparame- ters for the best performance (see App. C for details). For SHALLOW WATER and LAVIER-STOKER we use the first 5 time points to infer latent state and then predict the next 20 time points, while for SCALAR FLOW we use the first 10 points for inference and predict the next 10 points. For synthetic data, we consider two settings: one where the data fully observed (i.e., the complete state recorded) a setting for which most models are designed and where the data is partially observed (i.c., only part of the full state
Table 1: Test MAE for different interpolation methods.
Datuset
Grid
k-NN
Linear
IDW
Coarser
0.046 + 0.002
0.034 + 0.001
0.038
0.002
Shallow Water
Original
0.017
0.016
0.017
Finer
0.031 +:t 0.003
0.017 : 0.003
0.030 -: 0.002
Coarser
0.087 + 0.006
0.069 + 0.009
0.066 +: 0.006
Navier Stokes
Original
0.048 + 0.009
0.041 + 0.003
0.045 + 0.010
Finer
0.054
0.044 0.004
0.049 0.002
Coarser
0.041 0.021
0.032 0.009
0.035 0.012
Scalar Flow
Original
0.019 + 0.001
0.018 + 0.000
0.018 + 0.001
Finer
0.040 + 0.016
0.026 + 0.006
0.028 0.007

---

is given, as discussed at the beginning of this section). The results are shown in Table 2. We see that some of the baseline models achieve reasonably good results on the fully-observed datasets, but they fail on partially-observed data, while our model maintains strong performance in all cases. Apart from the fully observed SHALLOW WATER dataset where FEN performs slightly better, our method outperforms other methods on all other datasets by a clear margin. See Appendix C for hy hyperparameter details. In Appendix E we demonstrate our model's capability to learn dynamics from noisy data. In Appendix F we show model predictions on different datasets.
Table Test MAE for different models.
Shallow Water (Full)
Shallow Water (Partial)
Navier Stokes (Full)
Navier Stokes
Model
Scalar Flow
(Partial)
NODE
0.036 0.000
0.084 0.001
0.053 0.001
0.109 0.001
0.056 0.001
FEN
0.011 0.002
0.0641 = 0.005
0.031
0.001
0.108 0.002
0.062 0.005
SNPDE
0.019 + 0.002
0.033 = 0.001
0.042
0,004
0,075 + 0.002
0.059 + 0.002
DINo
0.027 + 0.001
0.063 = 0.003
0.047 = 0.001
0.113 + 0.002
0.059 + 0.001
Net
NA
0.061 0.001
0.103 0.003
0.056 + 0.003
Ours
0.014 0.002
0.016 0.002
0.024
0.003
0.041 0.003
0.042 0.001
6
Related Work
Closest to our work is Ayed et al. (2022), where they considered the problem of learning PDEs from partial observations and proposed a discrete and grid dededdeenddntt model that is restricted to regular spatiotemporal grids. Another related work is that of Nguyen et al. (2020), where they proposed a variational inference framework for learning ODEs from noisy and partially-observed data. However, they consider only oow-dimensional ODEs and are restricted to regular grids.
Other works considered learning the latent space PDE dynamics using the "encode-process-ded process-decode" approach. Pfaff et al. (2021) use GNN-based encoder and dynamics function and map the observations the same spatial grid in the latent space and learn the latent space dynamics. Sanchez et al. (2020) use a similar approach but with CNNs and map the observations to a coarser latent grid and learn the coarse-scale dynamics. Wu al. (2022) use CNNs to map observations vector and learn the latent dynamics. However, all these approaches are pridddependent, limited to regular spatial/temporal grids, and require fully-observed data.
low-dimensional latent
Interpolation has been used in numerous studies for various applications. Works such as (Alet et al., 2019; Jiang et al., 2020; Han al., 2022) use interpolation to map latent states coarse grids to observations on finer grids. Hua et al. (2022) used interpolation as a post-processing step to obtain continuous predictions, while Boussif et al. (2022) used it to recover observations missing nodes.
Another approach for achieving grid grid-independence was presented neural operators (Li et al., 2021; Lu et al., 2021), which learn mapping between infinitediimensional function spaces and represent the mapping in a grid-indeppndent manner.
Conclusion
We proposed novel space-time continuous, grid-independent model for learning PDE dynamics from noisy and partial observations on irregular spatiotemporal grids. Our contributions include an efficient generative modeling framework, novel latent PDE model merging collocation and method of lines, and a data-efficient, grid-independent encoder design. The model demonstrates state-of-the-an art performance on complex datasets, highlighting its potential for advancing data-driven PDE modeling and enabling accurate predictions of spatiotemporal phenomena in diverse fields. However, our model and encoder operate on every spatial and temporal location which might not be the most efficient approach and hinders scaling to extremely large grids, hence research into more efficient latent state extraction and dynamics modeling methods is needed.

---

References
Ferran Alet, Adarsh Keshav Jeewajee, Maria Bauza Villalonga, Alberto Rodriguez, Tomas Lozano- Perez, and Leslie Kaelbling. Graph element networks: adaptive, structured computation and memory. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference Machine Learning, volume 97 of Proceedings of Machine Learning Research, pages 212-222. PMLR, 09-15 Jun 2019. URL https: wwtpt::stppppeeeeeeeeeeeeeeeeeee http::///ooceeeeeeeee mlr, press/ /97/aleta html.
Ibrahim Ayed, Emmanuel de Bézenac, Arthur Pajot, and Patrick Gallinari. Modelling spatiotemporal dynamics from earth observation data with neural differential equations. Machine Learning, 111 (6):2349-223) 2380, 2022. doi: 11.00//s1994-002-00000 s 111999-022--61111
URL htttps:/dde org/ 1000.007
David M. Blei, Alp Kucukelbir, and Jon D. McAuliffe. Variational inference: A review for statisticians. Journal of the American Statistical Association, 112(518):859-87) Apr 2017, ISSN 1537-274X. doi: 10. 111000006214599201118888888 1285773.
URL http: www.dooooo.oo..ooooo.o org/ 1000/0000001100 2017.
H.G. Bock and K.J. Plitt. A multiple shooting algorithm for direct solution of optimal control problems*. IFAC Proceedings Volumes, 17(2): 1603-1680, 1984, ISSN 1474-6670. doi: https://doi. 009/00000166111444-6677000666666 atric1le/p11//4444477017722222
URL https:/ /ttps://ww//r sciencedirect. cwm/sciicccccccccccccceecceeecce /cccee//////////////// 9th IFAC World Congress: Bridge Between Control Science and Technology, Budapest, Hungary, 2-6 July 1984.
Oussama Boussif, Yoshua Bengio, Loubna Benabbou, and Dan Assouline. MAgnet: Mesh agnostic neural PDE solver. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022. URL https::/openreviie hott/forumiid-bx220 8hca8.
Johannes Brandstetter, Daniel E. Worrall, and Max Welling. Message passing neural PDE solvers. In International Conference on Learning Representations, 2022. URL forum?id=vii SrruumiddvSSixxPPSSS3
net/
Steven L. Brunton and J. Nathan Kutz. Data-Driven Science and Engineering: Machine Learning, Dy- namical Systems, and Control. Cambridge University Press, 2019. doi: 0011017/981108888669999
Nicolas Cellier. scikit-fdiff, 2019. URL [KL.attpp//ggtaaacc/////eeelllep
Ricky T. Q. Chen. torchdiffeq, 2018. URL
Ricky T. Q. Chen, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud. Neural ordinary differential equations, 2018.
Alexander H.-D. Cheng. Radial basis function collocation method. In Computational Mechanics, pages 219-219, Berlin, Heidelberg, 2009. Springer Berlin Heidelberg. ISBN 978-3-540-7599997,
N. Cressie and C. K. Wikle. Statistics for Spatio-Temporte Data. Wiley, 2011. ISBN 9780471692744 URL https::///// /bbokks.//..ooooooo google. wwwbookks/iiiccoooooooo ww0.DDDDDOO..
Marie-Lenat Eckert, Kiwon Um, and Nils Thuerey. Scalarflow: A large-scale volumetric data set of real-world scalar transport flows for computer animation and machine learning. ACM Transactions on Graphics, 38(6):239, 2019.
L. C. Evans. Partial Differential Equations. American Mathematical Society, 2010. ISBN 9780821849743.
Xavier Glorot and Yoshua Bengio. Understanding the difficulty of training deep feedforward neural networks. Yee Whye Teh and Mike Titterington, editors, Proceedings of the Thirteenth International Conference on Artificial Intelligence und Statistics, volume 9 of Proceedings of Machine Learning Research, pages 249-256, Chia Laguna Resort, Sardinia, Italy, 13-15 May 2010. PMLR. URL htttp::/procceeeeeee
mlr prees/v9/glorotiie www.oooia......oo.ooooooo00....o html.
10

---

S. Hamdi, W. E. Schiesser, and G. W Griffiths. Method of lines. Scholarpedia, 2(7):2859, 2007. doi: 10.4249/scholarpcdiaa2253 revision #26511.
Xu Han, Han Gao, Tobias Pfaff, Jian-Xun Wang, and Liping Liu. Predicting physics in mesh-reduced space with temporal attention. In International Conference on Learning Representations, 2022. URL https: http///pppeeeeeeeeeeeeee/eeeeeee wtwp://iwww..eeeooonnnnnaaaaaaaa
Pashupati Hegde, Cagatay Yildiz, Harri Lähdesmäki, Samuel Kaski, and Markus Heinonen. Vari- ational multiple shooting for bayesian ODEs with gaussian processes. The 38th Conference on Uncertainty in Artificial Intelligence, 2022. URL https: /ttp::///ppeeeeieeee http:///pppreeeeeeeeeeteeeeeee net/forumidd= INNuh
Charles Hirsch. Numerical computation of internal and external flows: The fundamentals of compu- tational fluid dynamics. Elsevier, 2007.
Philipp Holl, Vladlen Koltun, Kiwon Um, and Nils Thuerey. phiflow: A differentiable pde solving framework for deep learning via physical simulations. In NeurIPS Workshop on Dif- ferentiable vision, graphics, and physics applied machine learning, 2020. URL https:
//montrealrobotics ca/diffcvgp/ assets patsss//ppppp/..//// pdf.
Chuanbo Hua, Federico Berto, Michael Poli, Stefano Massaroli, and Jinkyoo Park. Efficient contin- uous spatio-temporal simulation with graph spline networks. In ICMI. 2022 2nd Al for Science Workshop, 2022. URL
Valerii lakovlev, Cagatay Yildiz, Markus Heinonen, and Harri Lähdesmäki. Latent neural ODEs with sparse bayesian multiple shooting. In The Eleventh International Conference on Learning Representations, 2023. URL https::/openrrwiieet wtwpeerrwweee..ee/ooummmeiiaii http://orrmmmddmmmo11111..11...
Chiyu IMaxr Jiang, Soheil Esmaeilzadeh, Kamyar Azizzadenesheli, Karthik Kashinath, Mustafa Mustafa, Hamdi A. Tchelepi, Philip Marcus, Mr Prabhat, and Anima Anandkumar. Mesh- freeflownet: A physics-constraiie International Conference for High Performance Computing, Networking, Storage and Analysis, URL httt:://xx.dde doi wttg/...10/////////...... 00119//SC44000
deep continuous space-time super-resolution framework. SC20:
Nov 2020. doi: 105100999441405.2200000 2020.00013.
Armand Jordana, Justin Carpentier, and Ludovic Righetti. Learning dynamical systems from noisy sensor measurements using multiple shooting, 2021.
E.J. Kansa. Multiquadrics- scattered data approximation scheme with applications computational fluid-dynamics-il solutions to parabolic, hyperbolic and elliptic partial differential equations. Computers and Mathematics with Applications, 1976111771161 1990. ISSN 0898- 1221. doi: http://bbo.og//11100060o//111111 sccaate/aarticaaa////////2111111
URL https: https:///wwww sciencedireett. com/
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization, 2017.
Diederik P Kingma and Max Welling. Auto-encoding variational bayes, 2013.
Dmitrii Kochkov, Jamie A. Smith, Ayya Alieva, Qing Wang. Michael P. Brenner, and Stephan Hoyer. Machine learning-accelerated computational fluid dynamics. Proceedings of the National Academy of Sciences, (1821)):2101188111 2021. doi: 00.073//pa.22001888111 URL https: //www. pnas atgg/diiabs///// 1073 /pas. 2101784118.
Zongyi Li, Nikola Borislavov Kovachki, Kamyar Azizzadenesheli, Burigede liu, Kaushik Bhat- tacharya, Andrew Stuart, and Anima Anandkumar. Fourier neural operator for parametric partial differential equations. In International Conference on Learning Representations, 2021. URL
https://openrevier wwt.rriwwwweeeeeeeeee wwt. http://www.eneetooiiiiiidd-do... www.ooiidd-dPPNgeeam...
Marten Lienen and Stephan Günnemann. Learning the dynamics of physical systems from sparse ob- servations with finite element networks. International Conference on Learning Representations, 2022. URL https: // /pwprevieeeeeeeeee/eeee/ee
11

---

Zichao Long, Yiping Lu, Xianzhong Ma, and Bin Dong. PDE-net: Learning PDEs from data. In Jennifer Dy and Andreas Krause, editors, Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pages 3208-3216 URL httph:t///// cttt:p:://pocceeding
PMLR, 10-15 Jul 2018.
mlr. www.erreessseeee press/v80/long html.
Lu Lu, Pengzhan Jin, Guofei Pang. Zhongqiang Zhang, and George Em Karniadakis. Learning nonlin- ear operators via deeponet based on the universal approximation theorem of operators. Nature Ma- chine Intelligence, 3(3): Mar 2021. ISSN 2522-5839. doi: 10. 10.108/642256-22-0000355..
URL http::////// /ttp:://xx.. doi..ddoiooo org/ /ttg//..1// 1038/ 240.://42256--0 022255021000000-5-5-555
Luke Metz, C. Daniel Freeman, Samuel S. Schoenholz, and Tal Kachman. Gradients are not all you need, 2021.
James D. Murray, Mathematical Biology I. An Introduction, volume 17 of Interdisciplinary Applied Mathematics. Springer, New York, 3 edition, 2002. doi: 10. 10.1007//98888
Duong Nguyen, Said Ouala, Lucas Drumetz, and Ronan Fablet. Variational deep learning for the identification and reconstruction of chaotic and stochastic dynamical systems from noisy and partial observations. ArXiv, abs/2009.02296. 2020.
Tobias Pfaff, Meire Fortunato, Alvaro Sanchez-Gonzalez, and Peter Battaglia. Learning mesh-based simulation with graph networks. In International Conference on Learning Representations, 2021. URL
net. wttp://www..eeeetorrrr//tttooooo
XP.
Michael Poli, Stefano Massaroli, Federico Berto, Jinkyoo Park, Tri Dao, Christopher Re, and Stefano Ermon. Transform once: Efficient operator learning in frequency domain. In ICML 2022 2nd Al
for Science Workshop, 2022. URL https: httpps//tppppee http::///ppeeeieeeeeee
Antônio H. Ribeiro, Koen Tiels, Jack Umenberger, Thomas B. Schön, and Luis A. Aguirre. On the smoothness of nonlinear system identification. Automatica, 121:109158, 2020. ISSN 0005-1098. doi: https://doi.co http:///oooorg/110/ .rrg/1.. att. com/sciemss
URL https: : //ww://wwwww/// sci wcc.enere..........
article/fa weapiiceelpiiiio/e /p01///000000882003566666
Cristopher Salvi, Maud Lemercier, and Andris Gerasimovics. Neural stochastic pdes; Resolution- invariant learning of continuous spatiotemporal dynamics, 2021.
Alvaro Sanchez, Dmitrii Kochkov, Jamie Alexander Smith, Michael Brenner, Peter Battaglia, and Tobias Joachim Pfaff. Learning latent field dynamics pdes. 2020. We are submitting to Machine Learning and the Physical Sciences workshop with submission deadline on October 2nd.
William E. Schiesser. The numerical method of lines: Integration of partial differential equations. 1991.
Sungyong Seo, Chuizheng Meng, and Yan Liu. Physics-aware difference graph networks for sparsely- observed dynamics. In International Conference on Learning Representations, 2020. URL.
https://opperevvi/ w heet/forummddrrtte dorue??ddrrlllllytt...
Kiwon Um, Robert Brand, Yun Fei, Philipp Holl, and Nils Thuerey, Solver-in-the-Loop:: Learning from Differentiable Physics to Interact with Iterative PDE-Solvers. Advances in Neural Information Processing Systems, 2020.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, E ukasz Kaiser, and Illia Polosukhin. Attention all you need. In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, edi- tors, Advances Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URL https: https::/proceeee httpa//////ooco/sss/e/ee pwourips.c/ppapee (ttp///444476111100000111.e.....
wttpc//ppe..eeeeeppppeeeeeeeeeee wttps///aaa/////////111 Paper .dff.
Henning Voss, J. Timmer, and Juergen Kurths. Nonlinear dynamical system identification from uncertain and indirect measurements. International Journal of Bifurcation and Chaos, 14, 01 2004.
Tailin Wu, Takashi Maruyama, and Jure Leskovec. Learning to accelerate partial differential equations via latent global evolution. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022. URL https::///// http///pppeeeiieeee//eeeeee
12

---

Cagatay Yildiz, Markus Heinonen, and Harri Lähdesmäki. Ode2vae: Deep generative second order odes with bayesian neural networks. ArXiv, 0bbs/1905.10944 2019.
Yuan Yin, Matthieu Kirchmeyer, Jean- Yves Franceschi, Alain Rakotomamoniy, and patrick gallinari. Continuous PDE dynamics forecasting with implicit neural representations. In The Eleventh International Conference on Learning Representations, 2023. URL https: http:::s//ppeeevieee net/ forum/id-B73n1ljjbrr
13

---

A
Appendix A
A.1 Model specification.
Here provide all details about our model specification. The joint distribution for model is pl M $1:B, Oaya dse) p(u1:N 81.B, Odyn: (ss.)p((1.1100lll))
(24)
Next, we specify each component in detail.
Parameter priors. The parameter priors are isotropic zero-mean multivariate normal distributions:
p(Oayn) (ddyn 00 I). p (ddee) = (Ode 00, /);
(25)
(26)
where N is the normal distribution, O is zero vector, and 7 is the identity matrix, both have an appropriate dimensionality dependent on the number of encoder and dynamics parameters.
Continuity prior.
We define the continuity prior
B
p(81.n (ayn) = p(81)
II
p(s)) 1h Odyn),
(27)
b=2
N
B N IIII b=2
II p() 1=1
pl 15-1 Odpa)
(28)
N II (s 1)
B N IIII N (0]1=(t10) Ty; 11. 86 Odyn), o2.1) ] 0=2
(29)
where N is the normal distribution, 0 E Ed is zero vector, Rdxd the identity matrix, and Te E IR is the parameter controlling the strength of the prior. Smaller values of ae tend to produce smaller gaps between the sub-trajectories.
Observation model
p(uINN(1.. Ooyn: O
p(u: (sb, Odyns Oave)
(30)
6=1 IE j=1
N
pl ( Igew (2(1,, my; 401, 8b. 00m)))
(31)
b=1 < j=1
II KEI j=1
N 1oom (=(te, 20; 8b. (lyn)),
o
(32)
where N the normal distribution, oZ is the observation noise variance, and I E DxD is the identity matrix. Note again that 2t,, 2t; f(p), 8b, Odya) above equals the ODE forward solution ODESolve(1;; (nj, 8b, Oayn) at grid location Zj.
A.2 Approximate posterior specification.
Here we provide all details about the approximate posterior. We define the approximate posterior as
B
(fdyyn Osec, $1.B) = ((0aym) (alaayn)((aaaa ( Prhayo (Bayn) Iatan (farc)
IIII
(s/)
(33)
b=1j=1
Next, specify each component in detail.
14

---

Dynamics parameters posterior.
We define Qehenn (Odya) as
Arhann (Odyn) = ((dynn Yaya, diag Tdya)),
(34)
where Tayn and Tay are vectors with an appropriate dimension (dependent on the number of dynamics parameters), and diag( (yaa)) is matrix with Tayn on the diagonal. We define the vector of variational parameters as dayn (7aya Thyn). We optimize directly over Dayn and initialize "ayn using Xavier (Glorot and Bengio, 2010) initialization, while Tdya is initialized with each element equal to 9 10
Decoder parameters posterior.
We define arbua (O as
ae (Oace) = N (0 Ydees diag )
(35)
where Yitec and Tee are vectors with appropriate dimension (dependent the number of decoder parameters), and diagg(Tae) is matrix with Tee on the diagonal. We define the vector of variational parameters as daee (Yeec. Tdec). We optimize directly over dace and initialize "áce using Xavier (Glorot and Bengio, 2010) initialization, while Tece is initialized with each element equal to 9 10
Shooting variables posterior.
We define Duz (s/) as
Ong = (ss mp. diag((T){3))),
(36)
where the vectors 7/. Ti. E I2d are returned by the encoder he and diag([T]]] is matrix with [Ty]] on the diagonal. We define the vector of variational parameters as v = (7g [Tg]). Because the variational inference for the shooting variables is amortized, our model is trained w.r.t. the parameters of the encoder network, denc-
B Appendix B
B.1 Derivation of ELBO.
For our model and the choice of the approximate posterior the ELBO be written
/ 9 (00yn Oac. $1:B) In
p( ((U1::d S1:B, Odya, Bose)
C
(37)
q(Odyn: Odec, $1.B)
/ (0ayn Otse, 81.B) In pl [1 81.B, Bayn: Bac pl( $ ayy (p0ynn (0l)(((((ae)ee)) d0 d0 de da 1:1
Q(81.B )g( 1yy (fsec)
(38)
-f g(aayn Ouee $1:B) In pu $1:B, Oayn: Gsc d000
(39)
f (Dddyn Odees 81.B) In ((81.B)
d0cc
(40)
p(81.1 O4yn)
f ((aay: Odees $1.B) ) In (0yya) (00am)
dl ds
dl
(41)
q(Osucs Odves $1.B) ) In (0ax) ploccc)
dlsyn
d8
de
(42)
L L2 Lg L4.
(43)
15

---

Next, we will look at each term L, separately.
C1 f /000mm dece $1.B) In p( (s.., Oöyn: P dan
(44)
N II p(uel 86, Oayn, Osce) dDayndaaddsi..
/ g Oddyn Odees 81.B) In II
(45)
/ (00ya, Bdse, $1.B) In [p( so, Oaya: Baxe) dodya
(46)
b=1
1ET,
B
N
/ (00yyn Odee sb) In pluce] 1s,, Ouyn, Boos) docyadlasooss
(47)
b=
IETs
B
N
In [p Odyn 0oce)
(48)
1ET
C2 / q( (laayn Pdcer $1.B) In (($1.B)
d0 dyn
(49)
P[81.B/ [dyn)
/ (Day:: Odee, $1:53 In
((8))
((8))
dRosndBaecd$1::
(50)
p(s1)
p(sb 88 8dyn)
b=2
= / (Day: Odces 81.B) In
N
g ((s²)
d0ayn d8
(51)
j=1 pl ({)
+ / g(Oayn: Osec, $1:B) In
BB N
((5)))
II
(52)
p( s5)] 10 Osyn)
b=25=1
N
(((5()
((0dyns Odces 81.8) In
d0 d0 do 1:B
(53)
p(s{)
[ g(Odym: Odec $1.B)
N
(($)))
Σ In
dosndlaaods1...
(54)
p(s0] 186-1.8-11
1. Odyn)
1=1
N
{(5{) (((())
((s1)
In
dsf
(55)
Σ
B
/ 9(0dym: 86-1, 86)
N
In
((5)))
d0 dOayndss-
(56)
p p(1/1 11-1 1, Odyn)
6-2
1=1
N
((0(1) [p(sf)
Σ
((0))
do;
(57)
1=1-
N
/ ((s)) In
((x))
Σg
g(Baym: 88-1) 1)
Σ
dsp
dday
(58)
p(s/1sb- 0dym)
N
B
N
= Σ KL (a(s))ll(<)))
+
E EgtOan 80 1)
Σ j=1
KL
(59)
where KL is Kullback-Leibler (KL) divergence. Both of the KL divergences above have closed form but the expectation w.r.t. g(Dayn; 86-1) does not.
L3 = KL(g(0asn)) (l
Li = (L(((()
(60)
16

---

B.2 Computation of ELBO.
We compute the ELBO using the following algorithm:
1. Sample Odyn: Ode from dohu. (Dayn), OePrce (Oaes).
2. Sample $1.B by sampling each sy from 44g (s()) with 2 = he (u(tt6)) x,]).
3. Compute Wi:M from 8:B 3 as in Equations 14-16.
4. Compute ELBO C (KL terms are computed in closed form, for expectations we use Monte Carlo integration with one sample).
Sampling is done using reparametrization to allow unbiased gradients w.r.t the model parameters.
C Appendix C
C.1 Datasets.
SHALLOW WATER. The shallow water equations are a system of partial differential equations (PDEs) that simulate the behavior of water shallow basin. These equations are effectively depth-integrated version of the Navier-Stokes equations, assuming the horizontal length scale is significantly larger than the vertical length scale. Given these assumptions, they provide model for water dynamics in basin or similar environment, and are commonly utilized in predicting the propagation of water waves, tides, tsunamis, and coastal currents. The state of the system modeled by these equations consists of the wave height h(t, v, y). velocity in the x-direction u(t, x. v) and velocity in the y-direction 2(,,2, Given an initial state (ho: Un to), we solve the PDEs on spatial domain & over time interval [0, TJ. The shallow water equations are defined as:
0h (hu)
a(hv)
+
0,
(61)
0t
dr
dy
du +
du
du
0h
u Dx
0, v Oy
+ 9 Dr
(62)
0t
du
du
Ou
0h
+
+ 9 Oy
0,
(63)
0
0x
Oy
where is the gravitational constant.
The spatial domain Ω is a unit square with periodic boundary conditions. We set 7 = 0.1 sec. The solution is evaluated at randomly selected spatial locations and time points. We use 1089 spatial locations and 25 time points. The spatial end temporal grids the same for all trajectories. Since are dealing with partially-observer cases, we assume that we observe only the wave height lu(t, x, y).
For each trajectory, we start with zero initial velocities and the initial height ho(a, 8) generated as:
ho(a,y)
Σ
Aut cos((2 (ker ly)) + Tha sin(2 + ly)),
(64)
k.1=-N
ho(2,y) 1 hola, v) min(ho)
(65)
max(ho) min(ho)
where 3 and Aua Th ~ 1).
The datasets used for training, validation, and testing contain 60, 20, and 20 trajectories, respectively. We use scikit-fdiff (Cellier, 2019) to solve the PDEs.
NAVIER-STOKES. For this dataset we model the propagation of a scalar field (e.g., smoke concen- tration) in a fluid (e.g., air). The modeling is done by coupling the Navier-Stokes equations with the Boussinesq buoyancy term and the transport equation to model the propagation of the scalar field. The state of the system modeled by these equations consists of the scalar field c(t, x, y), velocity in --dirrction u(t, y). velocity in y-direction u(t, x, v). and pressure p(t, x. y). Given an initial state
17

---

(co, No. vo, Pu), we solve the PDEs on spatial domain Ω over time interval [0, T). The Navier-Stokes equations with the transport equation are defined as:
Iu
du
= 0, Dy
(66)
Dz
Du
du
du
Op
Puu
Pu
+
+
+
(67)
0t
0x
Oy
de
0x2
dy²
dr
Ou
Du
Op
02v
16
dr
0
0y
Oy
0.2
01y²
+
(68)
dc
0c
dc
8
0
(69)
Ot
dx
dy
dx²
0y²
where v 0.002.
The spatial domain
unit square with periodic boundary conditions. We set T 2.0 sec, but drop the first 0.5 seconds due to slow dynamics during this time period. The solution is evaluated at randomly selected spatial locations and time points. We use 1089 spatial locations and 25 time points. The spatial and temporal grids are the same for all trajectories. Since we are dealing with partially-observed cases, we assume that we observe only the scalar field c(t, x, y).
For each trajectory, we start with zero initial velocities and pressure, and the initial scalar field co co v) generated as:
N
80(x, = Σ
As cos((2 + Iy)) + The sin(2 Iy)),
(70)
k11=-N
C0(2, y) min(Co)
co(x, u) (1ax(7x) min(Cn)
(71)
where = 2 and λu, Tku ~ (0,
The datasets used for training, validation, and testing contain 60, 20, and 20 trajectories, respectively. We use PhiFlow (Holl et al., 2020) to solve the PDEs.
SCALAR FLOW. This dataset, proposed by Eckert al. (2019), consists of observations of smoke plumes rising in hot air. The observations are post- processed camera images of the smoke plumes taken from multiple views. For simplicity, we use only the front view. The dataset contains 104 trajectories, where each trajectory has 150 time points and each image has the resolution 1080 x 1920. Each trajectory was recorded for T = 2.5 seconds.
To reduce dimensionality of the observations we sub-sample the original spatial and temporal grids. For the temporal grid, we remove the first 50 time points, which leaves 100 time points, and then take every 4th time point, thus leaving 20 time points in total. The original 1080 1920 spatial grid is first down- sampled by factor of 9 giving a new grid with resolution 120 213, and then the new grid is further sub-sampled based on the smoke density at each node. In particular, we compute the average smoke density at each node (averaged over time), and then sample the nodes without replacement with the probability proportional the average smoke density (thus, nodes that have zero density most of the time not selected). See example of final grid in Figure 9, This gives a new grid with 1089 nodes.
300
50
Figure 9: Spa- tial grid used for SCALAR FLOW dataset.
We further smooth the observations by applying Gaussian smoothing with the standard deviation of 1.5 (assuming domain size 120 213).
We use the first 60 trajectories for training, next 20 for validation and next 20 for testing.
In this case the spatial domain is non-periodic, which means that for some observation location some its spatial neighbors (a;) might be outside of the domain. We found that to account for such cases it is sufficient to mark such out-of-domain neighbors by setting their value to -1.
18

---

Time grids used for the three datasets are shown in Figure 10.
60
No malized Time
0.1
No Nomallize Time
0.0
0.
U4
0/
10
Normalized Time
Figure 10: Time grids used for SHALLOW WATER (top), NAVIER-STOKES (middle), and SCALAR
FLOW (bottom).
C.2 Model architecture and hyper-parameters.
For all datasets we define Fe... as an MLP. For SHALLOW WATER/NAVIER- STOKES/SCALAR FLOW we use 1/3/3 hidden layers with the size of 1024/512/512, respectively, We use ReLU nonlinearities.
Dynamics function.
Observation function. For all datasets define Gos as selector function which takes the latent state (t,, E Rd and returns its first component.
Encoder. Our encoder ha... consists of three function: ha B and he The spatial aggregation function higgual a linear mapping to 128, The temporal aggregation function he stack of transformer layers with temporal attention and continuous relative positional encodings (Iakovlev et al., 2023). For all datasets, we set the number of transformer layers to 6. Finally, the variational parameter readout function ho is mapping defined
Linear(att)
Ab howw (ATT)
=
(Linear((
(72)
where Linear is linear layer (different for each line), and og and T/ are the variational parameters discussed in Appendix A.
Spatial and temporal neighborhoods. We the same spatial neighborhoods As( for both the encoder and the dynamics function. We define N's (x) as the set of points consisting of the point æ and points on concentric circles centered at æ, with radii and r/2, respectively. Each circle contains 8 points spaced 45 degrees apart (see Figure II (right)). The radius is set to 0.1. For FLOW the size of temporal neighborhood (6T) is set
SHALLOW to 0../0.1/0.2, respectively.
Multiple Shooting. For SHALLOW SATTE//NAIIERSSTOKESSSSSSS FLOW split the full training trajectories into 4/4/19 sub-trajectories, or, equivalently, have the sub-trajectory length of 6/6/2.
C.3 Training, validation, and testing setup.
Data preprocessing. We scale the temporal grids, spatial grids, and observations to be within the interval [0, 1].
Training. train our model for 20000 iterations using Adam (Kingma and Ba, 2017) optimizer with constant learning rate 3e-4 and linear warmup for 200 iterations. The latent spatiotemporal dynamics are simulated using differentiable ODE solvers from the torchdiffeq package (Chen, 2018) (we use dopri5 with rtol=1e-3, atol=1e-4, no adjoint). The batch size is 1.
19

---

Validation. We use validation set to track the performance of our model during training and save the parameters that produce the best validation performance. As performance measure we use the mean absolute error at predicting the full validation trajectories given some number of initial observations. For SHALLOW
FL.Ow we use the first 5/5/10 observations. The predictions are made by taking one sample from the posterior predictive distribution (see Appendix C.4 for details).
Testing. Testing is done similarly to validation, except that as the prediction we an estimate of the expected value the posterior predictive distribution (see Appendix C.4 for details).
C.4 Forecasting.
Given initial observations 11:m at time points 1.mm we predict the future observation ü, at a time point in > fm as the expected value of the approximate posterior predictive distribution:
p(ūn fimm u1:M) 20 f p( (8m, Oayn: Iee )g( (m (O(y (d8 dl0 df0
(73)
The expected value is estimated via Monte Carlo integration, so the algorithm for predicting ü, is:
1. Sample 8ayn: dec from ((day)), ((ddcc)
2. Sample 8m from g( = IT Iwi. (8%)). where the variational parameters v are given by the encoder ha... operating the initial observations 11: as tby ho... (ū (ü(tm- æFy)).
3. Compute the latent state (t,,,) 2 (fm mm 8m: fyyn)-
4. Sample ü, by sampling each dz from N (u, (E(tns By xy))), o..
5. Repeat steps 1-4 n times and average the predictions (we use 10).
C.5 Model comparison setup.
NODE. For the NODE model the dynamics function was implemented as fully connected feedforward neural network with hidden layers, 512 neurons each, and ReLU nonlinearities.
FEN. We use the official implementation of FEN. We use FEN variant without the transport term as we found it improves results on our datasets. The dynamics were assumed to be stationary and autonomous in all cases. The dynamics function was represented by a fully connected feedforward neural network with 3 hidden layers, 512 neurons each, and ReLU nonlincarities.
NSPDE. We use the official implementation of NSPDE. We set the number of hidden channels to 16, and set modesl and modes2 to 32.
DINo. We use the official implementation of DINo. The encoder an MLP with 3 hidden layers, 512 neurons each, and Swish non non-linearities. The code dimension is 100. The dynamics function is an MLP with 3 hidden layers, 512 neurons each, and Swish non-linearities. The decoder has 3 layers and 64 channels.
MAgNet. We use the official implementation of MAgNet. We use the graph neural network variant of the model. The number of message- steps is 5. All MLPs have 4 layers with 128 neurons each in each layer. The latent state dimension is 128.
D Appendix D
D.1
Spatiotemporal neighborhood shapes and sizes.
Here we investigate the effect changing the shape and size of spatial and temporal neighborhoods used by the encoder and dynamics functions. We use the default hyperparameters discussed in Appendix C and change only the neighborhood shape or size. neighborhood size of zero implies no spatial/tmporal aggregation.
20

---

Initially, we use the original circular neighborhood displayed in Figure 11 for both encoder and dynamics function and change only its size (radius). The results are presented Figures 12a and 12b. In Figure 12a, it is surprising to see very little effect from changing the encoder's spatial neighborhood size. A potential explanation that the dynamies function shares the spatial aggregation task with the encoder. However, the results in Figure 12b are more intuitive, displaying a U-shaped curve for the test MAE, indicating the importance of using spatial neighborhoods of appropriate size, Interestingly, the best results tend to be achieved with relatively large neighborhood sizes. Similarly, Figure 12c shows U-shaped curves for the encoder's 's temporal neighborhood size, suggesting that latent state inference benefits from utilizing local temporal information.
We then examine the effect of changing the shape of the dynamics function's spatial neighborhood. We use ncircle neighborhoods, which consist of n equidistant concentric circular neighborhoods (see examples Figure 11). Effectively, we maintain a fixed neighborhood size while altering its density. The results can be seen in Figure 13. We find that performance does not significantly improve when using denser (and presumably more informative) spatial neighborhoods, indicating that accurate predictions only require relatively sparse neighborhood with appropriate size.
Figure 11: Left: original circular neighborhood (Icirele). Center: circular neighborhood with increased size. Right: circular neighborhood of different shape (2circle).
0.040
Shallow
0,050
Shallow Water
0.060
Shallow Water
0.025
0.035
0.025
0.010
0.010
Test MAE
0.080
Test MAR
0.160
-avier-Slokes Stokes
Ntvier-Stokes
Stokes
0.055
0.090
0.030
0.050
0
Flow
0.060
Scalar Flow
0.080
Scalar Flow
0.045
0.050
0.040
0.020
0.0
0.05
0.1 0.2
0.0
0.05
0.1
0.1
0.2
Spotial Neighb. Size (Enc.)
Spatial Neighb. Size (Dyn.)
Temporal Nelghb. Size (Eno.)
(a)
(h)
(c)
Figure 12: (a),():: Test MAE vs spatial neighborhood sizes of the encoder and dynamics function, respectively. (c): Test MAE temporal neighborhood size of the encoder. Note that the spatial and temporal domains are normalized, so their largest size in any dimension is 1.
0.05
Shallow Water
Test MAE
0.00
0.05
Navier-Slokes
0.00
Scalar Flow
0.00
Icincle
2circle
4circle
Baircle
Spatial Neighb. Shape (Dyn.)
Figure 13: Test MAE vs spatial neighborhood shape.
21

---

D.2
Multiple shooting.
Here we demonstrate the effect of using multiple shooting for model training. Figure (left), we vary the sub-traj uub-trajectory length (longer sub-trajectories imply more difficult training) and plot the test errors for each sub-trajectory length. We observe that all cases, the best results are achieved when the sub-trajectory length is considerably smaller than the full trajectory length. In Figure 14 (right) we further show the training times, and can be seen multiple shooting allows to noticeably reduce the training times.
Training Time Training Timo Training Time
0.025
Shallow Water
(hours)
Shallow Walet
0.020
0.015
5
13
25
Test MAE
0.075
Navier-Stoken
(bours)
Navier-Stokes
0.050
5
13
25
0.075
Scaler Flow
(hours)
3
Scular Flow
0.050
2
4
10
20
10
20
Sub Sub tub-trajectory length
Sub trajecaajcryy length
Figure 14: Test MAE vs training sub-trajectory length.
E
Appendix E
Noisy Data. Here show the effect of observation noise model and compare the results against other models. We train all models with data noise various strengths, and then compute test MAE noiseless data (we still use noisy data to infer the initial state test time). Figure 15 shows that our model can manage noise strength up to 0.1 without significant drops in performance, Note that all observations are in the range [0, 1].
MAgNet
DINo
Ours
Shallow Water
Navier-
Scalar Flow
0.30
0.40
0.30
Test MAE
0.23
Test MAE
0.30
Test MAE
0.23
0.15
0.20
0.15
0.07
0.10
0.07
0.00
0.00
0.00
0.0
0.1
0.2
0.5
0.0
0.1
0.2
0.5
0.0
0.1
0.2
0.5
Noise Std
Noise Std
Noise Std
Figure 15: Test MAE observation noise d..
F
Appendix F
F.
Model Predictions
We show (Fig. 16) predictions of different models trained on different datasets (synthetic data is partially observed).
F.2
Visualization of Prediction Uncertainty
Figures 17, 18, and 19 demonstrate the prediction uncertainty across different samples from the posterior distribution.
22

---

Shallow Water
Navier-Stakee
Scalar Flow
Data
Ours
DINo
MAgNet
Time
Figure 16: Predictions from different models. Only forecasting is shown.
Cats Moon Peodicion and Soyydard
2odie
Figure 17: Left: Test prediction for single trajectory SHALLOW WATER dataset. Show data (top), mean prediction (middle), and standard deviation of the predictions (bottom). Columns show predictions consecutive time points. Right: Ground truth observations (dashed black) and predictions (colored) with standard deviation plotted over time at two spatial locations.
03
T24e Joden
Test prediction for a single trajectory on NAVIER-STOKES dataset. Show are data (top), mean prediction (middle), and standard deviation of the predictions (bottom). Columns show predictions at consecutive time points. Right: Ground truth observations (dashed black) and predictions (colored) with standard deviation plotted over time at two spatial locations.
Figure 18: Left:
23

---

Date, Mean Prediction, and Standard Deviation
Tiaee Sodes
Test prediction for single trajectory on SCALAR FLOW dataset. Show are data (top), mean prediction (middle), and standard deviation of the predictions (bottom). Columns show predictions at consecutive time points. Right: Ground truth observations (dashed black) and predictions (colored) with standard deviation plotted over time at two spatial locations.
Figure 19: Left:
24


Reviews:
[{"review_id": "7uVJvzp93a", "reviewer": "Reviewer_1peH", "summary": "Modeling systems with spatiotemporal evolutions, such as the ones arise in problems governed by PDEs, is challenging. This is more pronounced when the system's underlying mechanisms are too complex or unknown. The authors propose a grid-independent generative model from noisy and partial observations on irregular domains. The latent state dynamics is constructed by merging the ideas inspired by the classical numerical PDE analysis such as collocation method and the method of lines, which is discussed in section 3.1 and 3.2. They also propose a novel encoder design that operates on local spatiotemporal neighborhoods for improved data-efficiency and grid-independence. The authors apply that their model three use cases from two synthetic analysis and one from real-world datasets. ", "questions": "The fact that the number and positions of sensor (observation locations) change with time may actually help the problem at the time of inference (Reconstruction). Have authors considered a fixed setting for which the observation locations are the same? Is there a minimum number of sensors required for a certain accuracy? Are there any insights about using this approach for sensor placement problems?\n\nIn equation 3, authors propose using a linear interpolation: what about other choices of interpolant? Any studies show performance of other choices? Seems to me that a linear interpolant may require a larger number of observation locations? Any insights?\n\nIn equation 7, what is the order of such analysis? Can authors prove this converges to actual dynamics? Seems heuristic. By order, I mean the order of accuracy of the approximation of the derivative in terms of the distance to neighborhood points (big O s^n).\n\nIn equations 12 and 13, what is the implication of assuming that the continuity inducing prior is Gaussian? Such assumption makes the solution of the problem computationally tractable but what if the actual process is non-Gaussian? What are limitations?\n\nIn the results section, for shallow water the model for d=5 (and maybe even 4) outperforms fully observed MAE which is based on true d=3. Are there any insight for this?\n\nMinor comment: the notation for $\\partial z/\\partialx_x$ on line 56 and elsewhere with a dot is a bit obscure. Maybe more common choice for notation is to use n as the degree of differentiation. \nAlso, A is capitalized on line 236.\n", "limitations": "N/A", "rating": 3, "confidence": 5, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "Multiple shooting analysis for training datasets from dynamical systems with long time simulations is a creative idea to enable reasonable cost of learning and avoid learning instabilities. \nIn Appendix D, they analyzed the impact of radius and multiple shooting on the overall performance of the model to further complete the analysis.\n\nThe examples provided are sophisticated enough to show that the framework can be applied to even more realistic scenarios.  \n", "weaknesses": "I have two major concerns about this work: \n1.\tThere are various heuristic choices/assumptions with no concrete proofs. For example, the choice of linear interpolant to map the values at the grid points of z(t) to z(t,x) is arbitrary. Moreover, when authors write the dynamical system in form of equation 7, while they are inspired by collocation methods, they are not actually utilizing such method. Instead, they argue that the functions and its (sufficiently smooth) derivatives are a function of neighborhood locations. This is another heuristic assumption which prevents a careful convergence analysis. Can authors prove if the radius tends to zero, such solution converges to the “real” solution? In what order the error decreases? The analysis in Appendix D is in fact anecdotal and does not mathematically guarantee a desired performance. Finally, regarding multiple shooting method, while I think this is a very nice idea, it’s not novel. For example, see \"50 years of time parallel time integration.\" By Martin, where several other methods other than multiple shooting are proposed to accelerate ODE integration. While I don’t see any fundamental error, I am not convinced that the analysis is mathematically sounds for computational adoption.\n2.\tThe topic of the paper is closely related to model discovery, which has been investigated in PDE community for quite some time now. A more appropriate comparison would be the one with SINDy (\"Discovering governing equations from data by sparse identification of nonlinear dynamical systems.\" By Brunton et al), Neural ODE (Neural Ordinary Differential Equations by Chen et al), or auto-encoder-based discussions (Data-driven discovery of coordinates and governing equations by Champion et al). Please also note, recently people also used neural ODE with collocation points (Physics-informed neural ODE (PINODE): embedding physics into models using collocation points by Sholokhov et al). To better position this work in the larger community, a much more detailed comparison is needed.  \n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your detailed review. We address the raised concerns below.\n\n**Heuristics**\n\nWe would be glad to rely on theory to guide all our choices, but unfortunately it is not always available for neural network parameterized models, and what remains is heuristics and empirical evidence.\n\nAn interpolation method is often chosen based on a modeling decision or empirical evidence on its performance. We chose linear interpolation since for our case it provides sufficient accuracy and computationally efficiency. However, your comment inspired us to investigate other interpolation techniques. We added two new interpolation methods: k-NN and inverse distance weighting, and tested them on spatial grids with different resolution. Please see the global rebuttal for description of the setup and results.\n\nAs you correctly noted, we do not use the classical form of the collocation method. Instead, we modify it by combining it with the method of lines (see Section 3.1 for motivation). Based on intuition from finite difference formulas, using a universal neural network based function approximator and decreasing the spatial neighborhood size should improve the accuracy of derivative estimation and hence the predictive performance. However, in practice the best performance is achieved with relatively large neighborhood sizes (see Appendix, Fig. 13). The optimal neighborhood size depends on the system being modeled. Overall, this is an important part of our proposed model that guarantees good performance and efficient computation.\n\nPlease note that we do not use classical multiple shooting and do not optimize over the shooting parameters. Instead, we introduce an encoder that maps observations directly to shooting states, and optimize over parameters of the encoder. As the result, the number of optimization parameters does not grow with the dataset size and spatiotemporal resolution, and no separate optimization loop is required at test time to infer the shooting states. Furthermore, Fig. 15 in Appendix shows that adoption of multiple shooting is worthwhile as it reduces training time and improves predictive performance.\n\n**Comparison to other methods**\n\nWe agree that a more detailed comparison to a larger number of relevant methods is required. We added three more methods, a simple baseline NeuralODE [3], and two SOTA methods as suggested by the reviewers: Neural Stochastic PDEs (NSPDE) [4], and Finite Element Networks (FEN) [5]. Please see the global rebuttal for description of the setup and results.\n\n**Questions**\n\n**Q1:** Number and positions of sensors.\n\n**A1:** Please note that we always use fixed observation locations (we will clarify this point in the revision). Studying optimal sensor placement is outside the scope of our work. However, moving observation locations could indeed be beneficial for the model's performance, especially if they are moved to parts of the space where more accurate resolution is required.\n\n**Q2:** Equation 3, other interpolation methods and number of observation locations.\n\n**A2:** We added results for other interpolation methods (see above). Other study ([1], Fig. 6) compared different types of interpolants, including learned ones, and showed little difference relative to the linear interpolation. Given a fixed number of observation locations, different interpolation techniques have different trade-offs between accuracy and computational efficiency. Linear interpolation is more on the efficiency side, but as we show in Fig. 8 and Table 1 above it performs well even on very coarse spatial grids.\n\n**Q3:** Equation 7 and convergence to true dynamics.\n\n**A3:** The fact that our model makes accurate predictions implies that the dynamics it learns are close to the true dynamics of the system. Giving theoretical convergence guarantees is challenging due to complexity of the model.\n\n**Q4:** Implication of assuming that the continuity inducing prior is Gaussian.\n\n**A4:** The purpose of the continuity prior is to bring the shooting and the system states close to each other. Assuming Gaussian prior implies that closeness is measured in terms of the squared distance. Note that the underlying process that all neural PDE methods aim to learn is a continuous-time deterministic PDE, and the sole motivation of multiple shooting is to introduce an approximation that converts the original problem into a form that enables efficient and robust optimization. In other words, the choice of the continuity prior corresponds to a choice of a multiple shooting approximation. \n\n**Q5:** Model performance for different latent state dimensions.\n\n**A5:** Based on our experience, the best performance in latent-state dynamic models is achieved for latent space dimensions larger than the true system state dimension. Similar observations were made in other works, e.g., [2]. A possible explanation for this is that in larger latent spaces simpler dynamics might be sufficient to model the system's behavior.\n\n\n**References**\n\n[1] MAgNet: Mesh Agnostic Neural PDE Solver, https://arxiv.org/abs/2210.05495\n\n[2] Non-linear State-space Model Identification from Video Data using Deep Encoders, https://arxiv.org/abs/2012.07721\n\n[3] Neural Ordinary Differential Equations, https://arxiv.org/abs/1806.07366\n\n[4] Neural Stochastic PDEs: Resolution-Invariant Learning of Continuous Spatiotemporal Dynamics, https://arxiv.org/abs/2110.10249\n\n[5] Learning the Dynamics of Physical Systems from Sparse Observations with Finite Element Networks, https://openreview.net/forum?id=HFmAukZ-k-2"}, {"review_id": "BSHm3Qr1ry", "reviewer": "Reviewer_VsBY", "summary": "The authors propose a grid independent model for learning PDEs from noisy experimental data. The proposed framework is probabilistic with an encoder that handles data efficiency and makes the solution grid independent. I think better differentiating from Ayed et al. and DINO would be helpful for the novelty contribution even though experimentally in these settings, the proposed model performs better.", "questions": "1. I know the authors mention that they use linear interpolation in the results in this paper and I was wondering if they performed any ablations with different types of interpolation.\n2. It is a bit of strong assumption that the neighboring points are sufficient to compute all the derivatives. The authors mention finite differences here but higher order finite difference methods can use more grid points in the stencils than just the neighbors. Can higher-order stencils be considered here? It would also be good to add a reference for finite difference methods, e.g. LeVeque's numerical analysis book.\n3. The authors mention that they test neighborhoods of various shapes and sizes. I think this is an important ablation. How ere two concentric circles chosen? The shapes are very important in numerical methods such as finite elements and those that use a Voronoi tessellation. In this case, it's a bit hard for me to see how it is purely grid independent. The number of evaluation points must also be carefully chosen.\n4. Also in Equation 8, the ODEs are defined only at the gridpoints and then interpolation is performed at every time step in the ODE solver - how expensive is this?\n5. Why is the distribution modeled as Gaussian and can the model be extended to support other distributions?\n6. How is the temporal neighborhood size delta_T determined? Since the model performs extrapolation in time shouldn't this neighborhood only depend on the prior time steps?\n7. Were simpler architectures than Transformers tested for the temporal encodings?\n8. I understand that latent space dimension for the synthetic 2D datasets but for the scalar flow, how do we interpret the improvement in prediction quality for d=5 instead of d=1 despite similar metrics? Can you prove any relation between latent dimension d and the error since it seems experimentally to decay monotonically?", "limitations": "Limitations on the expense of the proposed method and future work are discussed.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The problem setting is important to be able to learn PDEs from noisy and sparse observation data.\n- The paper is well-written with a good introduction of PDEs in the intro section.\n- Good overview of the related data-driven methods, such as Neural Operators (Li et al., 2021) and MeshGraphNets (Pfaff et al., 2021).\n- Grid independent solvers are advantages of deep learning models over numerical methods. It is a nice idea to define a neighborhood rather than use the adjacent gridpoints.\n- The proposed method is a nice application of the methods of lines.\n- The problem statement is well-defined. \n- It is good that the authors are considering the real-world setting of extrapolation in time.\n- Adding probabilistic output is also important contribution.\n- Synthetic and real-world datasets are both tested including a nice set of benchmark challenging datasets, e.g., Shallow water and Navier stokes. The scalar flow real-world camera dataset is also interesting.\n- Nice experimental setting going from a given grid to a coarser one and vice vera to finer.\n- The data efficiency results are nice.", "weaknesses": "- There is no guarantee that the PDEs learned from data will be better than the classical numerical PDE modeling so the first paragraph of the introduction could be modified a bit.\n- The authors should better motivate the irregularly spaced noisy and partial observations in this context. In this case, Gaussian Processes or Attentive Neural Processes could also be considered as baselines. In particular \"Learning Physical Models that Can Respect Conservation Laws\" [Hansen et al., ICML 2023] uses a constrained generative ANP model to perform the interpolation task from noisy sparse data to a fine grid. This work also uses a Gaussian observation model as in Eqn. 9.\n- There are advantages to operator-based methods vs. interpolation. I think a comparison to those methods, which are resolution-independent would be beneficial as well and should also be added to the related work section rather than just considering interpolation based methods.\n- Small number of comparisons to only DINo and MAgNet as baselines. Neural Operator and MeshGraphNets could also be considered as baselines.\n- The use of Neural ODEs is similar to the approach in \"Continuous PDE Dynamics Forecasting with Implicit Neural Representations\" [Yin et al., ICLR 2023]. In this work, the authors parameterize the latent space using implicit neural representations and then solve the ODE using Neural ODEs. It is good that the authors benchmark against this work in DiNO\n- The GNN based models such as MeshGraphNet also assume the derivatives are computed using the neighboring points as mentioned on lines 97-99. While this is important to highlight, since the approach here uses a neighborhood rather than grid points, it may be better to move this to a related work section so the novelty and design of the proposed method is only discussed in 3.1.\n- The model seems to be connecting various methods from different papers such as Neural ODEs and Iakovlev et al., 2023. It may be good to highlight and clarify the novelty in particular with respect to Iakovlev et al., 2023.\n- It looks like several heuristics need to be applied to make Neural ODEs work here. Also see \"Learning continuous models for continuous physics\" [Krishnapriyan, 2022]. It may be simpler to use a discrete numerical time-stepping method as done in \"CROM: Continuous Reduced-Order Modeling of PDEs Using Implicit Neural Representations\" [Chen et al., ICLR 2023]. This is also used in the MAgNet baseline and I don't think the only changes in the time-stepping scheme from discrete Euler to Neural ODEs (also done in DiNO) is sufficient for the novelty for publication. Forward Euler is also unstable in many cases as shown in the aforementioned Krishnapriyan, 2022 and higher order schemes like Runge-Kutta 4 (RK4) may be beneficial.\n- The overall model has many components and is quite complicated. I was wondering if any ablations were run to see which component is the most important.\n- 3-4 hours depending on the problem size and architecture is a bit slow. I think cost-accuracy comparisons to numerical solvers in the experimental section would be beneficial.\n- Figures 7-8 take significant space and may be moved to the appendix as ablations especially since they also show the results of the proposed method. The error plots between the predicted and exact may be easier to visualize.The scalar flow results look diffused.\n- I think stating 'very high accurate predictions\" on line 254 is a bit strong. The magnitude of the error is still high 1e-2 in comparison to numerical solvers\nMinor\n- For notation $u_i^j$ in numerical methods, the superscript typically notates time and subscript notates spatial coordinate.\n- I would give Section 3 a more descriptive title than Method and also the name the proposed method as well.\n- Can clarify on line 226 that the true latent space dimension for Navier-Stokes and Shallow water is 3 due to the x,y components of the velocity and the scalar pressure.\n- Capital A on line 236\n- \"Similarly\" to similar on line 240\n- No parathesis in references on line 272.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your detailed review. We address the raised concerns below.\n\n**Comparison to numerical solvers**\n\nAs discussed in Section 2, we assume the dynamics are unknown, hence numerical solvers are not applicable since they require fully known system dynamics.\n\n**Comparison to other methods**\n\nWe added three more methods, a baseline NeuralODE [3], and two SOTA methods as suggested by the reviewers: Neural Stochastic PDEs (NSPDE) [4], and Finite Element Networks (FEN) [5]. Please see the global rebuttal for description of the setup and results. \n\n**Model complexity, importance of model components, and ablation studies**\n\nOur method consists of multiple components. Encoder is used to infer the latent state from observations. Dynamics function is used to propagate the latent state in time. Decoder maps the latent state to parameters of the observation distribution. And multiple shooting is a parameter inference technique which accelerates and stabilizes training. Considering the above, each component has its role and is important.\n\nRegarding various model parameters, we conducted multiple ablation studies throughout our work, please see the results in Fig. 7, 13, 14, and 15.\n\n**Discrete-time models and heuristics**\n\nIn our setting the temporal grids are irregular, thus discrete-time models are not applicable. Training dynamic models is hard and often requires heuristics (one-step training, progressive increasing the length of the training trajectory, multiple shooting, and other) to stabilize training.\n\n**Novelty relative to other methods.**\n\nAs outlined in the Introduction, the main advantage of our model are: (i) grid-independence, (ii) space-time continuity, (iii) data-efficiency, (iv) learning from partial observations, an (v) fast and stable training. No previous method has all these features.\n\nThe main differences wrt Iakovlev et al., 2023 is that they do latent ODE modeling, while we do latent PDE modeling. Also, their model is not applicable in our case since: 1) The encoder and decoder are grid-dependent and applicable only on uniform spatial grids, 2) Their dynamics function could not be adapted to the PDE setting as it assumes that every point affects every other point on the grid, which is not the case in PDEs, where each point is affected only by points within a small neighborhood.\n\n**Paper structure**\n\nThank you for your suggestions regarding the structure of the manuscript. We will improve our work by incorporating them into the revised version. Please see global rebuttal for details.\n\n**Questions**\n\n**Q1:** Different interpolation methods.\n\n**A1:** We investigated other interpolation techniques. We added two new interpolation methods: k-NN and inverse distance weighting, and tested them on spatial grids with different resolution. Please see the global rebuttal for description of the setup and results.\n\n**Q2:** Using neighborhood values to compute derivatives, and extension to higher-order stencils.\n\n**A2:** As shown in Fig. 3, we do not use only neighbor nodes' values, instead, we define fixed spatial neighborhoods whose values depend on other nodes as well (lines 97-107). In Appendix D we investigate the effect of using spatial neighborhoods of different shapes and sizes, so extension to higher-order stencils is straightforward. \n\n**Q3:** Spatial neighborhood shape and grid-independence.\n\n**A3:** Concentric circles shape was selected because it has better coverage than e.g., a cross or a square. As we show in Fig. 14, the choice of the neighborhood shape (number of circles) does not have a strong effect on the model's performance. Also note that the spatial neighborhoods have the same shape and do not depend on the observation locations, which makes our model grid-independent as we discuss in lines 97-107.\n\n**Q4:** How expensive is interpolation?\n\n**A4:** Interpolation amounts to matrix-vector multiplication where the matrix has dimensions Cn-by-n, where n is the number of nodes, and C is a small constant that depends on the spatial neighborhood shape. For our choices of the interpolation methods the matrix is extremely sparse, hence requires negligible memory to store and the matrix-vector product can be computed efficiently. Overall, it is one of the least expensive parts of the model.\n\n**Q5:** Why is the distribution modeled as Gaussian and can the model be extended to support other distributions?\n\n**A5:** The observation distribution is Gaussian as this is a reasonable choice in the absence of any system-specific requirements. Our model is agnostic to the choice of likelihood models, and our model can be easily extended to other observation distributions by asking the decoder to output parameters of that distribution (as in Eq. 9).\n\n**Q6:** Temporal neighborhood size and dependence only on prior time steps.\n\n**A6:** The temporal neighborhood size is a hyperparameter. Using only previous time points to infer the latent state would be similar to filtering, while using also the future time points is similar to smoothing. Smoothing tends to be more accurate (see e.g., Kalman filtering/smoothing).\n\n**Q7:** Were simpler architectures than Transformers tested for the temporal encodings?\n\n**A7:** We also considered continuous-time versions of RNNs, but they are much slower than Transformers.\n\n**Q8:** Improvement of the results for larger latent space dimension d.\n\n**A8:** For the scalar flow, improvement in prediction quality for larger d can be interpreted similarly to the synthetic datasets. Setting d=1 is not sufficient since the \"true state\" has higher dimension. Setting d > 1 allows the encoder to infer the missing states which allows the dynamics function to model the system more accurately. Proving any relationships, beyond simply extrapolating the points to larger d, is difficult due to complexity of the system.\n\n**References**\n\n[1] https://arxiv.org/abs/1806.07366\n\n[2] https://arxiv.org/abs/2110.10249\n\n[3] https://openreview.net/forum?id=HFmAukZ-k-2\n"}, {"review_id": "tcaFXNXdGl", "reviewer": "Reviewer_yPZT", "summary": "This paper introduces a new method for learning time-dependent PDE solutions with noisy, partially-observed data on irregular grids. This setting is quite challenging and is aligned with real-world data acquisition. They adopt a generative framework and combines two techniques for solving PDEs: the collocation method, which is used to approximate spatial derivatives, and the method of lines to propagate the PDE solution forward in time. At the core of their architecture lies a spatio-temporal encoder, that aggregates for every spatial/ temporal coordinate pair $(x_j, t)$ information over the spatial neighborhood $\\mathcal{N}_S(x_j)$ and the temporal neighborhood $\\mathcal{N}_T(t)$. They tackle the bayesian problem with a variational formulation, and learn the overall model by maximizing an ELBO objective. They use sparse bayesian multiple shooting to stabilize and accelerate the training. They validate their framework on three different datasets: Shallow-Water (SW), Navier-Stokes (NS) and Scalar Flow (SF).", "questions": "* Did you try interpolant functions other than linear interpolation ? Similarly, have you tried the model with spherical data to see how the interpolation on a sphere behaves ? \n* How does the model scale with the number of grid points ? ~1100 points for a 2D domain does not appear to be of really high-resolution. Is the method still applicable for instance with standard grids used in the community such as 64x64 or 64x128 points ?\n* Is the method able to capture high-frequency patterns that occur in turbulent flows ? \n* Is it possible that the proposed model overfitted the training horizon and as a result obtained better results than DINo and MAgNet ? This would be in line with the training-size analysis. How does the model extrapolate over the training horizon ?\n", "limitations": "There is no limitation section. The authors discuss briefly the difficulties for the method to scale with the numbers of points.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "This is overall a good technical contribution. The paper shows solid experimental results, outperforming DINo and MAgNet on the three different datasets. The idea of using a generative model for solving PDEs is interesting, and the combination of all the different blocks is far from trivial to train. The explanation of the architecture is clear, though the notations are sometimes a bit hard to follow. The authors also evaluate the capabilities of the model to adapt to coarser and finer grids to support their claims. \n", "weaknesses": "* The paper is quite difficult to read overall. The introduction and related work are very brief, especially for such a difficult topic. It would help the reader to add some background from the PDE Deep Learning literature in those sections. For instance, there is no mention of Neural Operators which has been a popular topic recently, and there is no clear explanation of the limitations of existing state-of-the-art methods. __N.B.__: DINo which tackles a very similar problem (except for noiseless data), is not cited in the introduction nor the related work but is still used as the main baseline. __N.B.2__: The Section 3.2 with the explanation of the multiple shooting framework is a bit confusing and does not add much in terms of explanation of the model. On the other hand, inference is only exposed in the Appendix and should be in the core of the paper. \n\n* Though the idea of spatiotemporal encoding is sound, its realization does not appear to be very elegant and some limitations that could contradict the original claim of spatial continuity come to mind. What if the interpolation in the spatial neighborhood is not applicable ? Let's say you have an obstacle in the domain or you want to know the solution at the boundaries of your system, how can this work ? In all the figures the frames seem to contain only the convex hull of the set of points. This is worrying in terms of possible application to domains with complex geometries which are of main interest for people working in computational dynamics.\n\n* The results shown in Figure 10 are very impressive, especially for Navier-Stokes (NS) and Shallow-Water (SW). However I wonder, how can the model generalize this well on new test initial conditions with only 2 training trajectories? In terms of MAE, the method is already at 0.03 and 0.07 with two training samples for SW and NS while it reaches 0.015 and 0.041 with 60 samples. Do you have a possible explanation for this phenomenon ?\n\n* I understand that the main motivation of the paper is to develop a framework for partially-observed, noisy, irregularly-spaced observations, and therefore there are not a lot of suitable candidate baselines. Still, I think it would help readers understand the difficulty of the tasks to have a comparison with more classical neural operators and pde solvers on regular settings.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your detailed review. We address the raised concerns below.\n\n**Paper structure**\n\nThank you for your suggestions regarding the structure of our manuscript. We will improve our work by incorporating them into the revision. We will move details about forecasting to the main text and will extend the introduction and related work to include neural operators and other relevant models. We will also elaborate on the limitations of existing SOTA methods.\n\n**Comparison to other methods**\n\nWe agree that a more detailed comparison to a larger number of relevant methods is required. We added three more methods, a simple baseline NeuralODE [1], and two SOTA methods as suggested by the reviewers: Neural Stochastic PDEs (NSPDE) [2], and Finite Element Networks (FEN) [3]. Please see the global rebuttal for description of the setup and results.\n\n**Applications of our method on non-convex spatial domains**\n\nOur method can be applied on non-convex spatial domains without any modifications. Existence of obstacles is not a problem as long as the interpolation method is applicable (which is the case for many interpolation techniques such as k-NN, IDW, linear interpolation, piece-wise polynomial interpolation).\n\nThe state of the system can be obtained at any point within the convex hull of the observation locations. One should expect that the number and positions of the observation locations are appropriate for the problem at hand. In that case it would be possible to obtain the state sufficiently close to the boundaries.\n\n**Data efficiency.**\n\nTo achieve outstanding data-efficiency shown in Fig. 10 we utilize spatiotemporal locality of dynamical systems. Namely, we use the fact that derivatives in a PDE are defined locally (hence we need only local information to define the time rate of change), and assume that the latent state depends only on a local spatiotemporal neighborhood (as discussed in Section 4.2). These ideas are reflected in the design of our model which operates on local spatial and temporal neighborhoods instead of working on the whole grid, such as DINo for example.\n\n**Questions**\n\n**Q1:** Spherical grids and other interpolation methods.\n\n**A1:** Spherical grids are out of the scope of our work, but our method can potentially be applied to such data assuming that an appropriate interpolation method for spheres is used and that the spatial neighborhoods are appropriately defined on the sphere.\n\nWe also investigated other interpolation techniques. We added two new interpolation methods: k-NN and inverse distance weighting, and tested them on spatial grids with different resolution. Please see the global rebuttal for description of the setup and results.\n\n**Q2:** Is the method able to capture high-frequency patterns that occur in turbulent flows ?\n\n**A2:** We did not consider that particular question, but as our experimental results suggest, our model is capable of learning complex flow-based phenomena.\n\n**Q3:** How does the model scale with the number of grid points ?\n\n**A3:** Since our model operates on each grid point, it scales linearly with the number of grid points. For reference, our model occupies around 6GB on GPU during training for shallow water dataset.\n\n**Q4:** How does the model extrapolate over the training horizon ?\n\n**A4:** To test this, we trained our model only on the first half of the training time grid (i.e., from 0 to 0.5T, where T is the observed time-horizon), and tested it on the full time grid (of length T). The test error for shallow water and scalar flow did not increase, while for Navier-Stokes it increased almost twice. This observation can be understood by the fact that velocity of the species transported in the Navier-Stokes dataset keeps increasing over time. Thus, the system behavior that the model was trained on (before time 0.5T), is different from the system behavior on which the model was tested (after time 0.5T), naturally leading to higher prediction errors.\n\n\n**References**\n\n[1] Neural Ordinary Differential Equations, https://arxiv.org/abs/1806.07366\n\n[2] Neural Stochastic PDEs: Resolution-Invariant Learning of Continuous Spatiotemporal Dynamics, https://arxiv.org/abs/2110.10249\n\n[3] Learning the Dynamics of Physical Systems from Sparse Observations with Finite Element Networks, https://openreview.net/forum?id=HFmAukZ-k-2"}, {"review_id": "myY2orDhzW", "reviewer": "Reviewer_mL6m", "summary": "This work proposes a latent variable model for PDEs with an encoder-decoder architecture that evolves the latent variables in time with a neural ODE. The model achieves independence of the locations of the evaluation points by linearly interpolating the data onto points distributed in a fixed pattern around each evaluation point.", "questions": "- What are the advantages of modeling PDEs in a latent space instead of data space? Do you have any experimental insight into this question?", "limitations": "The authors did not discuss any limitations of their model.", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "- The paper's main contribution is the grid-independence of the model by interpolating observations and latent states onto fixed neighboring points around each observation point\n- The experimental section investigates the effect of the latent state dimension and the size of the training dataset", "weaknesses": "- The collocation method is a function space method that finds solutions that solve the PDE at a given set of collocation points. I don't see how this is used in this paper. The authors refer to the general form of first-order-in-time PDE in Eq (4) as the collocation method, but Eq (4) is just a general form of a PDE and not a method.\n- The method of lines is only about eliminating spatial derivatives by discretizing the solution for each $t$ in a function space and does not contain any notion of data-driven models or being restricted to evaluations of $z(t, x)$ as claimed in lines 88-89\n- Generative modeling of PDE solutions (section 3.2) is neither motivated nor evaluated in experiments\n- Use of bayesian deep learning (variational inference of model parameters, section 3.2) is not motivated\n- Overall, the approach is quite similar to the one presented in (Lienen and Günnemann, 2022) in terms of method-of-lines and interpolation of data, though the authors only mention but don't compare to them. In contrast to the author's writing in their related work section, I would even see this as the most closely related work.\n- Application of method of lines is not novel as claimed in line 37-38, see (Iakovlev et al., 2020) and (Lienen and Günnemann, 2022)\n- This paper does not actually propose/apply a collocation method as claimed in line 37-38\n- As a result, novelty and significance of this work are limited\n\n(Iakovlev et al., 2020): https://arxiv.org/abs/2006.08956\n(Lienen and Günnemann, 2022): https://arxiv.org/abs/2203.08852", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your detailed review. We address the raised concerns below.\n\n**Motivating the use of generative modeling and Bayesian inference**\n\nWe agree that we can improve our presentation regarding these points. Generative modeling is a standard approach in state space modeling which allows to explicitly define the parameter, process, and observation models, allowing to make explicit modeling assumptions, see e.g., \"Statistics for Spatio-Temporal Data\" by by N. Cressie. The ability to incorporate uncertainty in the predictions is valuable when there is limited amounts of data and in challenging real-world scenarios, where it allows to account for inherent system stochasticity and model uncertainty. Please note that we use the mean of the posterior predictive distribution as our prediction, which implies that we fully utilize the Bayesian nature of our method. We will add these clarifications to the revision.\n\n**Comparison to other methods and to Lienen and Günnemann, 2022**\n\nIndeed, the method by Lienen and Günnemann, 2022 is related to ours as it also models data on irregular spatiotemporal grids. However, as we show in the new set of experiments (see the global rebuttal), it fails in more realistic and challenging setting that we consider in our work.\n\nWe agree that a more detailed comparison to a larger number of relevant methods is required. We added three more methods, a simple baseline NeuralODE [1], and two SOTA methods as suggested by the reviewers: Neural Stochastic PDEs (NSPDE) [2], and Finite Element Networks (FEN) from Lienen and Günnemann, 2022 [3]. Please see the global rebuttal for description of the setup and results\n\n**Collocation method and the method of lines**\n\nWe agree that our method does not utilize the classical form of the collocation method. But please note that our method closely follows the collocation method, and then modifies it using the method of lines. In more detail, we represent z(t, x) as an interpolant (Eq. 3) and substitute it in the PDE (Eq. 4), which is then represented as a system of ODEs (Eq. 8), where each ODE corresponds to an observation location. This is very similar to how the collocation method is applied to time-dependent PDE problems. However, instead of using the collocation method directly, we combine it with the method of lines. While application of the method of lines is not novel, its combination with the collocation method, to the best of our knowledge, is novel and, as we discuss in lines 80-92, leads to multiple advantages such as grid-independence.\n\nWe agree that the method of lines is not constrained to a particular type of spatial grids, but the way it was used in previous works makes their models grid-dependent. The main reason for grid-dependence, as discussed in lines 96-107, is reliance on the locations of the grid nodes, which our proposed methods avoids.\n\n**Questions**\n\n**Q1:** What are the advantages of modeling PDEs in a latent space instead of data space? Do you have any experimental insight into this question?\n\n**A1:** This is an important question as this is one of the main contributions of our work. We demonstrate the advantages of our latent PDE model in Fig. 7, where we compare a data-space variant of our model (latent state dimension d=1), and a latent space variant of our model (d > 1). As can be seen, modeling in the latent space results in drastic improvements in predictive performance. Setting d > 1 allows the encoder to infer the missing features of the state variable, which in turn allows the dynamics function to model the system more accurately. In other words, modeling a system in the data space sets a fundamental limitation to the model and its accuracy if the system states are only partially observed, such as in the Scalar Flow as well as in many real-world datasets.\n\n**References**\n\n[1] Neural Ordinary Differential Equations, https://arxiv.org/abs/1806.07366\n\n[2] Neural Stochastic PDEs: Resolution-Invariant Learning of Continuous Spatiotemporal Dynamics, https://arxiv.org/abs/2110.10249\n\n[3] Learning the Dynamics of Physical Systems from Sparse Observations with Finite Element Networks, https://openreview.net/forum?id=HFmAukZ-k-2"}, {"review_id": "c5IOLQxHCU", "reviewer": "Reviewer_9sfV", "summary": "The authors proposed a new method in learning PDEs by combining interpolation and NODEs. The difference from NODEs is that the authors uses spatial derivatives in hidden dynamics as PDE, compute the loss between shooting gaps with KL divergence, and compute the initial conditions with a transformer. The authors then support their claim by experimenting with three datasets, shallow water, navier-stokes, and scalar flow. ", "questions": "- In fig. 4 minimizing gaps, which norm do you use? Do you instead maximize the probability / likelihood in equation 17? In some commonly used PDEs, such as wave equations, derivatives of the PDE also matters. Did you take them into account? Or does your method only work on specific type of PDEs?\n- The experimental results in the paper is very different from the experimental results in DINO paper. Do you use a different experimental setup?\n- Are figure 8 and figure 9 partially the same?\n- What is a true latent dimension in line 225?\n- What are $h$ in figure 5? Are they intemediate states? But in equation 22 they look like functions. Are they the same $h$?\n- Also in figure 5, did you add extra observations around $x_j$? Or are they within the dataset? Did you assume all observations are on the same position but not grid-like? Then why don't we use NODEs with finite element meshes?\n- What types of problems are these experiments? Why are they interesting? Are they experiments from the SOTA / benchmark you compare with, or are they associated with any practical problems in application?\n- For those synthetic experiments, what are the noise level? What is the hidden equation? Navier-stokes might be self evident, but what about shallow water?\n- Do the models rely on initialization? How much runs did you run your experiments?\n", "limitations": "N/A", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 1, "contribution": 2, "strengths": "- The authors present a combination of various methods as a new solution to learning PDEs. \n- The performance of the method, on the particular datasets the authors choose, are quite well.", "weaknesses": "Overall, the proposed method might be nice, but the paper itself is very hard to read.\n- In section 3.2 generative modeling, there are two paragraphs discussing multiple shooting. There is no discussion on how multiple shooting is used on generative modeling.\n- There are too few captions for the figures. In figure 9, which parts are the time you interpolate? Which parts are extrapolation / forecasting? \n- The notations $q_\\psi(\\theta)$ is so confusing. \n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your detailed review. We address the raised concerns below.\n\n**Using multiple shooting in generative modeling, and gap minimization.**\n\nIndeed, the first two paragraphs of Section 3.2 only motivate the use of multiple shooting, but the rest of the section goes into details about how it is applied in our model. See especially Eq. 14-17 which describe our generative model and the role of multiple shooting in it.\n\nTo reduce the gaps we minimize the KL-divergence between the approximate posteriors of the shooting states and continuity prior (see part (iii) of the ELBO, Sec. 4.1). Note that the gap minimization via the KL-divergence is obtained directly from the variational lower bound (see the ELBO derivation in Appendix B.1) because the multiple shooting is formulated as part of our proposed generative model. During minimization we take the whole latent state into account, which means our method is not restricted to specific types of PDEs.\n\n**Do the models rely on initialization? How much runs did you run your experiments?**\n\nAs for any other deep-learning-based model, appropriate initialization helps our model to achieve better performance. As we discuss in Appendix A.2, we use the standard Xavier initialization. Also, as we mention in the experiments section, the results are the average over 4 runs with different random seeds. Overall, we observed that the performance of our method is very stable wrt to the random seed changes.\n\n**In Figure 9, which parts are interpolate/forecasting? Are Figures 8 and 9 partially the same?**\n\nIn Figure 9 we show only forecasting. We will clarify this point in the revision. Figures 8 and 9 both show predictions of our model, but the results in Figures 8 and 9 are for two different experiments.\n\n**What is a true latent dimension in line 225?**\n\nTrue latent state dimension refers to the dimension of the full system state, which is 3 for both shallow water and Navier-Stokes equations. For example, for the shallow water equation the state consists of the wave height (scalar) and velocity vector (2-dimensional vector). We will also clarify this point in the revision.\n\n**What are h in figure 5?**\n\nAs we discuss at the end of Section 4.2, h are functions that are parts of our encoder for amortized variational inference.\n\n**In figure 5, did you add extra observations around xj?**\n\nNo, we do not add any extra observations. We assume that we have observations only at the observation locations. As we discuss in Sec. 4.2, we use interpolation to obtain the system state outside of the observation locations (i.e., the locations marked with crosses around xj in figure 5). This is the core idea that is used both in the encoder and dynamics function which enables to make our proposed model continuous in space and grid-independent. We will clarify this aspect in the revised manuscript. \n\n**What types of problems are these experiments? Why are they interesting? Are they associated with any practical problems in application?**\n\nAs we discuss in Sections 1, 2, and 5, our experimental setup is highly-relevant for real-world applications, especially where the observations might be collected over irregular spatiotemporal grid, and the observed states might be incomplete. As our results (and extra comparisons show in the global rebuttal) demonstrate, this is a setting where currently available models fail, while our model demonstrates strong performance highlighting its utility in such challenging scenarios.\n\n**For those synthetic experiments, what are the noise level? What is the hidden equation?**\n\nData used for experiments in the main section is noiseless. However, in Appendix E we provide comprehensive results for noisy data with different levels of noise. As can be seen, our model is robust and maintains strong performance even under significant noise levels. Also, we provide all details about the systems used for our experiments in Appendix C, including the hidden equations. We will add this detail to the main text.\n\n**The experimental results in the paper is very different from the experimental results in DINO paper. Do you use a different experimental setup?**\n\nOur setup is more challenging than the one used in the DINo paper. They use fully-observed states and 512 training trajectories. In contrast, we use partially observed states and only 60 trajectories. To ensure fair comparisons, we conducted extensive hyperparameter tuning for DINo (as well as for all other baselines comparisons) using the default model parameters in the DINo paper as the starting point.\n"}, {"review_id": "CRexOTR3w6", "reviewer": "Reviewer_nYoy", "summary": "The authors present a grid-independent generative model to learn PDEs, which supports irregular grids. Amortized variational inference is employed for posterior approximation and multiple shooting, a recent method to train neural ODEs efficiently, is used and adapted to PDEs. \nResults show a quite significant dominance over two SOTA PDE learning methods. ", "questions": "Please refer to the weaknesses part for recommendations or modifications. \n\nRegarding other small questions or comments: \n\n- A general figure to describe the overall method would be very nice to have! Given, that the method is quite complex and gathers many different elements, it would really help grasping the whole framework, at least at a high level. \n\n- For the spatial aggregation, you use an MLP, like for example in MAgNet. The latter actually shows in the paper, in an ablation study, that the use of the MLP interpolator, is actually not so beneficial with respect to a cubic interpolator. Have you tried other interpolators, perhaps in an ablation study? Since this is the key to the grid-independent capabilities of the model, I think it is worth a part in the paper. \n\n- I think the C.4 appendix section on Forecasting should be put in the main text as it is quite crucial to the method. \n\n- I realized from the appendix training details that the batch size is 1. Was the data so big that it was impossible to fit a higher batch size in memory? Didn't it create training stability issues?\n\n- I think it is a bit unncessary to split the contributions in 3 the way you do it in the introduction : generative model, PDE model, encoder. To my understanding, the model you are proposing is exactly the mix of the 3, so I think it is better to simply do this one strong claim: you designed one PDE learning model, with all its characteristics. ", "limitations": "The authors addressed some limitations about the scalability of the method, but not much else; perhaps there are other points, like the training speed or the lack of special geometries taken into account within the method, or super-resolution capabilities still to be compared to SOTA? I do not see any potential negative societal impact. ", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The paper is well written and very well presented. The motivation is clear, the model formalism is well presented, the general flow is good. I enjoyed reading the paper. \n\n- The figures are well made and appreciated to have an intuition about the different elements. \n\n- The contribution of the paper is the model itself, which appears to be a natural extension to the Iakovlev et al., 2023, adpated for a PDE instead of an ODE setting. It is therefore not the most original methodologically, but a completely proper and natural contribution. \nNote that even though grid-independence and irregular grid prediction capabilities have been studied before for PDE learning, (notably by the SOTA the method is comparing to) which again restrains the novelty of the paper, the present methodology seems to significantly outshine the previous methods, which therefore justifies the contribution. ", "weaknesses": "\n- Relating to the previous comment regarding the contribution on the strength section: because it is not a new gap that the method is filling, but rather an improvement over SOTA methods, the contribution of the paper, besides being an extension of the multiple shooting for Neural ODE learning, relies on the performance in order to be strong. Therefore, it would be very useful to have more extensive experiments:\n   - of SOTA methods on finer and coarser grids, in order to test super resolution capabilities of the method compared to the other SOTA methods. An experiment is done on Figure 8, but not compared to SOTA methods\n   - to test prediction on long future horizon for all models including SOTA\n   - trying multiple grid sizes when training, instead of only multiple number of trajectories, as shown in Figure 10 \n\n- It seems that it lacks some important details\n    -  In the experiments regarding time specification, specifically the time horizon and the time resolution. It is mentioned in the appendix that \"25 time points\" are used for the two first data and 20 \"time points\" for the scalar flow, but it is not clear to me what they represent in real time. Maybe I missed it somewhere? If not, please mention it, as this is important when comparing to SOTA methods, as each method seems to have its strength, regarding spatial or time capabilities. \n\n    - In the figures: please add the time resolution on Figure 6 and 9 (how much is each time step?), and the horizon predicted on all the figures, to know what the MAE correspond to! If it is the same every time, please state it clearly somewhere (again, I apologize if I missed it...!)\n\n- How is the generative aspect of the method used? The fact that it is probabilistic is partly what makes it interesting as opposed to SOTA models (MAgNet and DINo, which are not generative), so it is a pity that no probabilistic outputs are used or shown in the results, but rather mean predictions.\n\nI think in general, the authors could move a bit of the formalism of the model to the appendix (which is also taken from Iakovlev et al., 2023 anyway, so that makes it easier), to let some space for more experiments and more details on the experiments, since, once again, they are in my view key to the contribution of the pure performance of the model!", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your detailed review. We address the raised concerns below.\n\n**Additional experiments**\n\nThank you for your suggestions. We agree that more experiments are required to fully demonstrate the advantages of our method. We conducted multiple additional experiments and present the results below.\n\nFirst, we extend our comparisons to a larger number of relevant methods. We added three more methods, a simple baseline NeuralODE [1], and two SOTA methods as suggested by the reviewers: Neural Stochastic PDEs (NSPDE) [2], and Finite Element Networks (FEN) [3]. Please see the global rebuttal for description of the setup and results.\n\nSecond, to test the performance of our model on predicting over longer time horizons than the ones seen during training, we trained our model only on the first half of the training time grid (i.e., from 0 to 0.5T, where T is the observed time-horizon), and tested it on the full time grid (of length T). The test error for shallow water and scalar flow did not increase, while for Navier-Stokes it increased almost twice. This observation can be understood by the fact that velocity of the species transported in the Navier-Stokes dataset keeps increasing over time. Thus, the system behavior that the model was trained on (before time 0.5T), is different from the system behavior on which the model was tested (after time 0.5T), naturally leading to higher prediction errors. Unfortunately, due to time constraints we were not able to conduct the time horizon experiments for other methods. But we are working on these experiments and will add them to the updated version of our manuscript.\n\nComparison of the super-resolution performance to other methods was not done due to poor performance of other methods even on the original training grids. However, this is a good suggestion and we will also conduct this experiment and add it to the updated version of our manuscript.\n\n**Time grids**\n\nIndeed, mentioning only the number of time points is not clear enough. For that reason we also provide information about the time horizons (T) in Appendix C.1. For Shallow Water T=0.1sec, for Navier-Stokes T=2sec, while for Scalar Flow T=2.5sec (we did not mention T for scalar flow, but we will add it in the revision). The system evolution over the mentioned time horizons is depicted in Fig. 6. As can be seen, the state of all systems changes a lot during these time intervals. For better intuition, we will visualize the time grids and add the figures into the revised version of our manuscript.\n\nThank you also for your suggestion for visualizing the forecasting intervals. We will rework the figures to include this suggestion. In our last experiment \"Comparison to other models\" we discuss the forecasting time intervals, but we agree that this detail should be discussed at the beginning of the Experiments section. We will modify our manuscript accordingly.\n\n**How is the generative aspect of the method used?**\n\nAs you correctly noted, we use the mean of the posterior predictive distribution as the prediction. As we discuss in Appendix C.4, the mean is computed by averaging over the predictions corresponding to different samples from the approximate posterior distribution. This is a common way in which uncertainty is incorporated into the model prediction.\n\n**Paper changes**\n\nThank you for your suggestions about improving our work! We will incorporate them in the revision by moving some of the multiple shooting formalism to the appendix and moving the appendix on forecasting to the main text. We will also work on providing a figure that gives a better overall idea about our method.\n\n**Using MLP for spatial aggregations**\n\nAs you correctly noted, we use an MLP as the spatial aggregation function. However, since the spatial aggregation is a part of the encoder, it is used only to infer the latent states, not for interpolation. For interpolation we use standard interpolation methods such as linear interpolation. \n\nIn addition to the linear interpolation, we investigated the performance of other interpolation methods: k-NN and inverse distance weighting. Please see the results and description of the experiment in the global rebuttal.\n\n**Using small batch sizes**\n\nIndeed, the batch size was set to 1 for all experiments. The amount of GPU memory used during training was around 4-6GB, so there was a possibility to increase the batch size. However, we did not face any issues with convergence, so we decided to leave it at 1.\n\n**References**\n\n[1] Neural Ordinary Differential Equations, https://arxiv.org/abs/1806.07366\n\n[2] Neural Stochastic PDEs: Resolution-Invariant Learning of Continuous Spatiotemporal Dynamics, https://arxiv.org/abs/2110.10249\n\n[3] Learning the Dynamics of Physical Systems from Sparse Observations with Finite Element Networks, https://openreview.net/forum?id=HFmAukZ-k-2"}]

Author rebuttal:
We would like to thank all reviewers for their careful reading and detailed comments. We believe that the suggested revisions have improved our manuscript. We believe our answers address all review comments, but if anything remains unclear we are happy to provide further clarifications.

We provided a response for each reviewer individually, and use this section to provide details and results of the extra experiments that the reviewers requested.

**Comparison to other models**

We agree that a more detailed comparison to a larger number of relevant methods is required. We added three more methods, a simple baseline NeuralODE [1], and two SOTA methods as suggested by the reviewers: Neural Stochastic PDEs (NSPDE) [2], and Finite Element Networks (FEN) [3]. We describe the setup and results below.

Setup: In addition to our partially-observed datasets, we created fully-observed versions of the synthetic datasets, where the whole system state is observed. This is a simplified setup where most models show good results. In the table below we show test MAE of all the methods.

|  Model | Shallow Water (Full) | Shallow Water (Partial) | Navier Stokes (Full) | Navier Stokes (Partial) | Scalar Flow (Real-World) |
|:------|:--------------------:|:-----------------------:|:--------------------:|:-----------------------:|:------------------------:|
| NODE   |   $0.036 \pm 0.000$  |    $0.084 \pm 0.001$    |   $0.053 \pm 0.001$  |    $0.109 \pm 0.001$    |     $0.056 \pm 0.001$    |
| FEN    |   $\boldsymbol{0.011 \pm 0.002}$  |    $0.064 \pm 0.005$    |   $0.031 \pm 0.001$  |    $0.108 \pm 0.002$    |     $0.062 \pm 0.005$    |
| NSPDE  |   $0.019 \pm 0.002$  |    $0.033 \pm 0.001$    |   $0.042 \pm 0.004$  |    $0.075 \pm 0.002$    |     $0.059 \pm 0.002$    |
| DINo   |   $0.027 \pm 0.001$  |    $0.063 \pm 0.003$    |   $0.047 \pm 0.001$  |    $0.113 \pm 0.002$    |     $0.059 \pm 0.001$    |
| MAgNet |          NA          |    $0.061 \pm 0.001$    |          NA          |    $0.103 \pm 0.003$    |     $0.056 \pm 0.003$    |
| Ours   |   $0.014 \pm 0.002$  |    $\boldsymbol{0.016 \pm 0.002}$    |   $\boldsymbol{0.024 \pm 0.003}$  |    $\boldsymbol{0.041 \pm 0.003}$    |     $\boldsymbol{0.042 \pm 0.001}$    |

We see that some of the baseline models achieve reasonably good results on the fully-observed datasets, but they fail on partially-observed data, while our model maintains strong performance in all cases. Apart from the fully observed shallow water dataset where FEN performs slightly better than ours, our proposed method outperforms other methods on all other datasets by a clear margin. We will add these results into the revised manuscript. 

**Investigating other interpolation methods**

We investigated other interpolation techniques and describe the setup and results below.

Setup: We used spatial grids with three different resolutions (Coarser, Original, and Finer). We trained our model only on the Original grid, and tested it on all grids to test how well it generalizes to grids with lower/higher resolution. We added two new interpolation methods: k-NN and inverse distance weighting. The results (test MAE) are shown in the table below.

|    Dataset    |   Grid   |        k-NN       |       Linear      |         IDW        |
|:-------------:|:--------:|:-----------------:|:-----------------:|:------------------:|
|               | Coarser  | $0.046 \pm 0.002$ | $0.034 \pm 0.001$ | $ 0.038 \pm 0.002$ |
| Shallow Water | Original | $0.017 \pm 0.002$ | $0.016 \pm 0.002$ |  $0.017 \pm 0.003$ |
|               | Finer    | $0.031 \pm 0.003$ | $0.017 \pm 0.003$ |  $0.030 \pm 0.002$ |
|               | Coarser  | $0.087 \pm 0.006$ | $0.069 \pm 0.009$ |  $0.066 \pm 0.006$ |
| Navier Stokes | Original | $0.048 \pm 0.009$ | $0.041 \pm 0.003$ |  $0.045 \pm 0.010$ |
|               | Finer    | $0.054 \pm 0.009$ | $0.044 \pm 0.004$ |  $0.049 \pm 0.002$ |
|               | Coarser  | $0.041 \pm 0.021$ | $0.032 \pm 0.009$ |  $0.035 \pm 0.012$ |
| Scalar Flow   | Original | $0.019 \pm 0.001$ | $0.018 \pm 0.000$ |  $0.018 \pm 0.001$ |
|               | Finer    | $0.040 \pm 0.016$ | $0.026 \pm 0.006$ |  $0.028 \pm 0.007$ |

We see that all interpolation methods perform rather similarly on the Original grid, but linear interpolation and IDW tend to perform better on finer/coarser grids than k-NN. Since our method can be combined with essentially any interpolation technique, this leaves a user the freedom to choose a particular technique that works for his/her application, although based on our experiments linear interpolation is often accurate. We believe this ablation is a valuable addition to our manuscript and we will include these experiments together with general discussion on choosing an interpolation method into the revised manuscript. 

**References**

[1] Neural Ordinary Differential Equations, https://arxiv.org/abs/1806.07366

[2] Neural Stochastic PDEs: Resolution-Invariant Learning of Continuous Spatiotemporal Dynamics, https://arxiv.org/abs/2110.10249

[3] Learning the Dynamics of Physical Systems from Sparse Observations with Finite Element Networks, https://openreview.net/forum?id=HFmAukZ-k-2

Meta-review: The paper presents a method for solving Partial Differential Equations (PDEs) using partially observed data on irregular spatio-temporal grids. This approach combines techniques inspired by collocation and the method of lines. Notably, it introduces a spatio-temporal encoder that aggregates neighbor information, coupled with a variational formulation. The method's performance is assessed through comparisons with various baselines on three datasets of trajectories, each corresponding to a distinct PDE generated from simulations.

All the reviewers agree on the importance of the problem addressed in the paper and acknowledge the strong performance of the proposed model. They conducted thorough analyses of the paper, including questions about its originality (as it seems to build upon [Iakovlev et al., ICLR 2023]), the initial experimental evaluation that relied on only two baselines, and the various heuristics employed in the model.

In response to these comments, the authors provided comprehensive answers and conducted additional experiments, including the inclusion of three new baselines. These new experiments reinforced the initial findings of the paper. The reviewers found these additions to be valuable and were convinced by the responses.

The authors are strongly encouraged to consider the reviewers' comments in order to enhance the organization and clarity of the paper
Decision: Accept (poster)