Evolve Smoothly, Fit Consistently: Learning Smooth Latent Dynamics For Advection-Dominated Systems
Authors: Zhong Yi Wan
Leonardo Zepeda-Nunez
Anudhyan Boral
Fei Sha
Venue: ICLR 2023 notable top 25% (2023)

Abstract: We present a data-driven, space-time continuous framework to learn surrogate models for complex physical systems described by advection-dominated partial differential equations. Those systems have slow-decaying Kolmogorov n-width that hinders standard methods, including reduced order modeling, from producing high-fidelity simulations at low cost. In this work, we construct hypernetwork-based latent dynamical models directly on the parameter space of a compact representation network. We leverage the expressive power of the network and a specially designed consistency-inducing regularization to obtain latent trajectories that are both low-dimensional and smooth. These properties render our surrogate models highly efficient at inference time. We show the efficacy of our framework by learning models that generate accurate multi-step rollout predictions at much faster inference speed compared to competitors, for several challenging examples.

Published as a conference paper at ICLR 2023
EVOLVE SMOOTHLY, FIT CONSISTENTLY: LEARN- ING SMOOTH LATENT DYNAMICS FOR ADVECTION- DOMINATED SYSTEMS
Zhong Yi Wan'
Leonardo Zepeda-Núñez"
Google Research
Google Research
Mountain View, CA 94043, USA
Mountain View CA 94043, USA
wanzy@gcogles com
1zepeddauuee72googgle.... com
Anudhyan Boral
Fei Sha
Google Research
Google Research
Mountain View, CA 94043, USA @oudhyan@google. com
Mountain View, CA 94043, USA
fsha@google.com
ABSTRACT
We present data-driven, space-time continuous framework to learn surrogate models for complex physical systems described by advection- differential equations. Those systems have alow-degaying Kolmogorov n-width that hinders standard methods, including reduced order modeling, from producing high-fidelity simulations low cost. In this work, construct hypernetwork- based latent dynamical models directly on the parameter space of compact representation network. We leverage the expressive power of the network and a specially designed consistency-inducing regularization to obtain latent trajectories that are both dineemsiiiinnll and smooth. These properties render our surrogate models highly efficient at inference time. We show the efficacy of our framework by learning models that generate accurate multi-step rollout predictions much faster inference speed compared to competitors, for several challenging examples.
partial
INTRODUCTION
High-fidelity numerical simulation of physical systems modeled by time-dependent partial differential equations (PDEs) has been the center of many technological advances in the last century. However, for engineering applications such as design, control, optimization, data assimilation, and uncertainty quantification, which require repeated model evaluation over a potentially large number of parameters, or initial conditions, high-fidelity simulations remain prohibitively expensive, even with state-of-art PDE solvers. The necessity of reducing the overall cost for such downstream applications has led to the development of surrogate models, which captures the core behavior of the target system but at fraction of the cost.
One of the most popular frameworks in the last decades (Aubry et al., 1988) to build such surrogates has been reduced order models (ROMs). In a nutshell, they construct lower-dimensional represen- tations and their corresponding reduced dynamics that capture the system's 's behavior of interest. The computational gains then stem from the evolution of lower-dimensional latent representation (see Benner et al. (2015) for comprehensive review). However, classical ROM techniques often prove inadequate for advection-dominntee systems, whose trajectories do not admit fast decaying Kolmogorov n-width (Pinkus, 2012), i.e. there does not exist well-approximating n-dimensional subspace with low n. This hinders projection-bas based ROM approaches from simultancously achieving high accuracy and efficiency (Peherstorfer, 2020).
Furthermore, many classical ROM methods require exact and complete knowledge of the underlying PDEs. However, that requirement unrealistic in most real-world applications. For example, in
"Equal contribution.

---

Published as a conference paper at ICLR 2023
Ambient space
((t)
u 2, t)
((t)
au((z,) Fu, u)
ui (zz
0)
u(ar t)
Aa
Aa
Encoder
Decoder
Aa
L ( (2)
00
(tt)
0 - (0)
00
u(2, 0)
00
Latent space evolution
Figure/ 1: (left) Diagram of the latent space evolution, where u represents the state variables, and 0 the latent space variables, which in this case correspond to the weights of neural network. (right) Sketch of two possible latent space trajectories, we seek to implicitly regularize the encode-decoder to obtain trajectories as the second one, which allows us to take very long time-steps latent space.
weather and climate modeling, many physical processes unknown or unresolved and thus approximated. In other cases, the form of the PDE is unknown. Hence, being able to learn ROMs directly from data is highly desirable. In particular, modern learning techniques offer the opportunities of using "black- box" neural networks to parameterize the projections between the original space and the low-dimensional space as well as the dynamics in the low-dimensional space Sanchez-Gonzaler et al., 2020; Chen et al., 2021; Kochkov et al., 2020).
On the other end, neural parameterization can be too flexible such that they are not necessarily grounded in physics. Thus, they do not extrapolate well to unseen conditions, or in the case of modeling dynamical systems, they do not evolve very far from their initial point before the trajectories diverge in nonphysical fashion. How can we inject into neural models the right types and amount of physical constraints?
In this paper, we explore a hypernetwork-tyne surrogate model, leveraging contemporary machine learning techniques but constrained with physics-informi
priors: (a) Time-Space Continuity: the trajectories of the model are continuous in time and space; (b) Causality: the present state of the model depends explicitly in the state of the model in the past; (c) Mesh-Agnostic: the trajectories of the model should be independent of the discretization both in space and time, it can be sampled at any given spacial-time point (d) Latent-Space Smoothness if the trajectories of the model evolve smoothly, the same should be true for the latent-space trajectories. The first three properties are enforced explicitly by choices of the architectures, whereas the last one is enforced implicitly at training via a consistency regularization.
Concretely, we leverage the expressive power of neural networks (Hornik et al.. 1990) to represent the state of the systems through an ansatz (decoder) network that takes cues from pseudo-spectral methods (Boyd, 2001), and Neural Radiance Fields (Mildenhall et al., 2020), whose weights encode the latent space, and are computed by an encoder hypeeeer-pprrreworkk. This hypernetwork encoder-decoder combination is trained on trajectories of the system, which together with a consistency regularization provides smooth latent-space trajectories. The smoothness of the trajectories renders the learning of the space dynamics effortless using neural ODE (Chen et al., 2018) model whose dynamics is governed by a multi-scale network similar to a U-Net (Ronneberger et al., 2015) that extracts weight, layer and graph level features important for computing the system dynamics directly. This allows us to evolve the system completely in latent-space, only decoding to ambient space when required by the downstream application (such process is depieted in Fig. 1 left). In addition, due to the smooth nature of the latent trajectories, we can take very large time-steps when evolving the system, thus providing remarkable computational gains, particularly compared competing methods that do not regularize the smoothness of the latent space (see Fig. I right for an illustration).
The proposed framework nonlinear and applicable to a wide range of systems, although in this paper we specifically focus on advection-dominatee systems to demonstrate its effecctivness. The rest of the paper is organized as follows. We briefly review existing work in solving PDEs in g2. We describe our methodology in $3, focusing on how to use consistency loss to learn smooth latent dynamics. We contrast the proposed methodology to existing methods on several benchmark systems in 54 and conclude in $5.
2

---

Published as a conference paper at ICLR 2023
2
RELATED WORK
The relevant literature is extensive and we divide it into six domains, ranging from PDE solvers for known equations, to purely data-driven learned surrogates, spanning wide spectrum of approaches.
Fast PDE solvers typically aim at leveraging the analytical properties of the underlying PDE to obtain low-complexity and highly parallel algorithms. Despite impressive progress (Martinsson, 2019), they are still limited by the need to mesh the ambient space, which plays an important role in the accuracy of the solution, as well as the stringent conditions for stable time-stepping.
Classical ROM methods seek to identify low-dimensional linear approximation spaces tailored to representing the snapshots of specific problems (in contrast with general spaces such as those in finite element or spectral methods). Approximation spaces are usually derived from data samples (Aubry al., 1988; Chinesta et al., 2011; Barrault et al., 2004; Amsallem al., 2012), and reduced order dynamical models are obtained by projecting the system equations onto the approximation space (Galerkin, 1915). These approaches rely on exact knowledge of the underlying PDEs and are linear in nature, although various methods have been devised to address nonlinearity in the systems (Willcox, 2006; Astrid et al., 2008; Chaturantabut & Sorensen, 2010; Geelen et al., 2022; Ayed et al., 2019).
Hybrid Physics-ML methods hybridize classical numerical methods with contemporary data-driven deep learning techniques (Mishra, 2018; Bar-Sinai et al., 2019; Kochkov et al., 2021; List et al., 2022; Bruno al., 2021; Frezat et al., 2022; Dresdner et al., 2022). These approaches learn corrections to numerical schemes from high-resolution simulation data, resulting in fast, low-resolution methods with accuracy comparable to the ones obtained from more expensive simulations.
Operator Learning seeks to learn the differential operator directly by mimicking the analytical properties of its class, such as pseuuo-differential (Hörmander, 2007) or Fourier integral operators (Hörmander, 2009), but without explicit PDE-informed components. These methods often leverage the Fourier transform (Li et al., 2021; Tran et al., 2021), the off-diagonal low-rank structure of the associated Green's function (Fan et al., 2019; Li et al., 2020a), or wwwpxiimaatiotttoreettte structures (Lu et al., 2021).
Neural Ansatz methods aim to leverage the approximation properties of neural networks (Hornik et al., 1990), by replacing the usual linear combination of handerafted basis functions by a more general neural network ansatz. The physics prior is incorporated explicitly by enforcing the underlying PDE in strong (Raissi et al., 2019; Eivazi et al., 2021), weak (E & Yu, 2018; Gao al., 2022), or min-max form (Zang al., 2020). Besides few exceptions, e.g., (de Avila Belbute-Peres et al., 2021; Bruna et al., 2022), these formulations often require solving highly non-convex optimization problems inference time.
Purely Learned Surrogates fully replace numerical schemes with surrogate models Iearned from data. A number of different architectures have been explored, including multi-scale convolutional neural networks (Ronneberger et al., 2015; Wang et al., 2020), graph neural networks (Sanchez- Gonzalez al., 2020), Encoder-Prooess--eeen ODEs (Ayed et al., 2019).
architectures (Stachenfeld al., 2022), and neural
Our proposed method sits between the last two categories. It inherits the advantages of both classes: our ansatz network is equally flexible compared to the free latent-spaces used in pure ML surrogates, and once trained there no need so solve any optimization problem at inference time, in contrast to many neural- ecural-ansatz-typ methods. In addition, our method is mesh-agnostic and therefore not subjected mesh-induced stability constraints, which are typical in classical PDE solvers and their ML-enhanced versions. Given that we seek to build surrogate models for trajectories of a potentially unknown underlying PDE, we not intend to replace nor compete with traditional PDE solvers.
METHODOLOGY
3.1
MAIN IDEA
We consider the space-time dynamics of physical system described by PDEs of the form
0,u(,, t) Flu(æ, t)],
(1)
u(x, 0) ugg
3

---

Published as a conference paper at ICLR 2023
Algorithm 1 Approximating u(
Input: initial condition: uo, time horizon: T
1. Construct uo following y sample interpolating un equi-spaced grid (xi)
2. Compute initial latent-space condition 0g- = Aa
3. Compute ((T) by integrating = (0))
4. Decode the approximation u(æ, T) Vo(T) (2)
5. Sample the approximation to the grid u" (0(T))
Output: u²
where u Ω R+ - R& denotes the state of the system with respect to spatial location and time FI.] a time-independent and potentially nonlinear differential operator dependent on a and its spatial derivatives {ØZut [dfu]ar and uo Uo is the initial condition where Uo is prescribed distribution. For time horizon T, we refer to {u(x, t), s.t. t E [0, 7]} as the rollout from ug, and to the partial function u(,, 1*) with fixed 1* snapshot belonging the function space U Mo- We call I the ambient space.
Encoder- Architecture For any given we approximate the corresponding snapshot of the solution
u(z, t) (z)))
where T is referred to as the ansatz decoder: neural network tailored to the specific PDE solve, it satisfies boundary conditions and general smoothness properties, ef. Eq. (10). The latent variables (() ④ are the weights of Y and time-varying quantities. θ is referred the latent space.
To compute the latent variables (()) we use hypernetwork encoder, which takes samples of u on grid given by x; iAr for i 0, ..., N. and outputs:
(tt) = A,, (fu(x,, o))
(3)
such that Eq. (2) satisfied. Here a encompasses all the trainable parameters the encoder.
Latent Space Dynamics We assume that satisfies an ODE given by
{ ((t) r')(O(t)), 0(0) = ((uo(()))
(4)
where I is another hypernetwork whose trainable parameters are aggregated as B.
Algorithm 1 sketches how solve PDEs through the latent space. For a given initial condition tgg we uniformly sample (or interpolate unstructured samples) on a grid and compute its latent representation using the encoder via Eq. (3). We numerically solve the ODE in Eq. (4) to obtain O(T). We then use the ansatz decoder to sample u(x, T') x Vocry at any spatial point This is reflected in the blue computation path in Fig. 1.
To simplify the presentation, consider only trajectory, and assume that the sampling grid, denoted by x = [x.) fixed. Then u(t) is used to represent the vector containing the samples of u(-,t) or [u(t)]; u(r,, Further, we conveniently denote u, u(t) and 0, ((t). Thus, following this notation we have
0. = A. (u:),
(5)
and (0) denotes the vector approximation of u sampled on x, i.e., [b))) = Vo, (x:).
Consistency Regularization from Commutable Computation Our goal is to identify the struc- ture of V, A and П, together with the parameters and so the diagram in Fig. commutes. In particular, we desire that the evolution in latent-space is equivalent to solving the PDE directly (denoted by the red path), but computationally efficient.
Thus, a pperfect' model of latent dynamics due to the encoder-decoder transformations should satisfy
" v(Auu))); == (O)O, and
= V,A(uúú
4

---

Published as a conference paper at ICLR 2023
channel vise projection
W
channel wits mixing
0.D.
20.
downsampling upsampling
Parumeters Rig
copy
weights
layers
global
Figure The Nonlinear Fourier Ansatz (left) featurizes the input with sin cos functions with specified frequency and trainable phases, and integrates small MLP to parametrize functions which are periodic [0. L). The Hyper U-Net (right) architecture partitions weights according to their layers (shown in different shades) and extract/aggregate features similarly to the original U-Net. More details on the transforms represented by arrows in the Hyper U-Net are included in Appendix A.3.
However, imposing the conditions entails the complete knowledge of the time derivative in the ambient space, which may not be feasible. To build the encoder/decoder independently of the time derivatives, note that
enforce the following constraints by combining the equations above:
ú V v (0) V Au)ú
and
Ô = Vo
These expressions are the time derivatives of
u, = (A(u,)) + Preumnstruct (u) = |lu (A((u)))
(8)
0, (0)) + conssstenccy (0) 110, A()().))
(9)
where the first equation leads to the typical reconstruction loss for autoencoders. The second equation, on the other end, leads to the consistency regularization which imposes the additional constraint ((V(O))) whenever possible.
This symmetric pair of loss functions is depicted in Fig. left. To enable transforming back and forth between ambient and latent spaces with little corruption, re-encoding u, gives back 0 is a must if we would take either red, blue or mixed path to rollout the trajectory from any point in time.
3.2 ARCHITECTURA, CHOICES
We provide a succinct overview the architectures used for the ansatz, encoder and dynamics.
The Nonlinear Fourier Ansatz (NFA) As example, in this work consider ansatz applicable to wide range of advection-dominated systems:
(x(x) = Σ (b(() sin (wkx + ax) + o (z) cos (wear + a-k),
(10)
where }wx} follow either a linear or dyadic profile (see Appendix A.3). This ansatz leverages low-frequency Fourier modes as envelop, corrected by shallow neural network in their amplitude and phase (the architecture depicted in Fig. 2). The ansatz periodic in [0, r.] by construction, and can be further tailored specific applications by adjusting the number of Fourier modes and their frequencies. this case, 0 aggregates the phases aş and the parameters inside MLPs in Φx-
Encoder Since assume that the input data of the encoder lies in one-dimensional equispaced grid, and that the boundary conditions are periodic, we use a simple convolutional encoder with periodic padding. This is achieved by sequence of convolutional ResNet modules (He et al., 2016) with batch normalization (loffe & Szegedy, 2015) and a sinusoidal activation function, which the input progressively downsampled in space. We consider either four five levels of downsampling. Each

---

Published as a conference paper at ICLR 2023
level consists of two consecutive ResNet modules and then downsample the spatial component by a factor of 2 while doubling the number of channels, thus maintaining the same amount information throughout the layers. A fully connected layer at the output generates the parameters of the ansatz network.
Hyper U-Net for Learning Dynamics Since the input and output of function I both follow the weight structure of the same neural network V we propose a novel neural network architecture, which we refer to as the Hyper U-Net, to model the mapping between them (illustrated in Fig. 2). Analogous to the classical U-Net (Ronneberger et al., 2015), which has 'levels' based on different resolutions of an image. the Hyper U-Net operates on 'levels' based on different resolutions of the computational graph of the ansatz network. In this paper, we use three levels which are natural to modern neural networks, namely individual weights, layer and graph levels, listed decreasing order of resolution. At each level, features are derived from the finer level in the downsampling leg of the U-Net, and later on mixed with the ones derived from coarser level in the upsampling leg. Differently from the U-Net, instead of convolutions use locally-connected layers that do not share weights.
3.3 MULTI-STAGE LEARNING PROTOCOL
Joint training of encoder, decoder and the dynamics U-Net is difficult, due to the covariate shift among these components (ablation study in Appendix A.5). To overcome this, we propose effective multi-stage learning procedure.
Learning to Encoder Snapshots For a given architecture of an ansatz network, we learn an encoder Aa that maps snapshots to the corresponding latent spaces, by optimizing the loss
(ene(a)) =
Σ (rrconssmaaaa (u.,) + ) Cooossstecyy (0t..)
(11)
The relative strength between the two terms in the loss function, is controlled by the scalar hyper- parameter ?. We found empirically that the second loss term injects inductive bias of preferring smoother and better-conditioned trajectories for the dynamic models to learn (see results section).
Learning to Evolve the Latent Representation After the encoder is trained, generate training latent space trajectories by computing Eq. (3) on all u. We then learn the latent dynamics operator T from sample latent rollouts Dg {to, t, (t1, Oz, ...) via the following loss function:
0 = 0co + /' f'r (s(((t)) dt,
CODE(()) = Σ 0r.. Ele.
Be.
(12)
where we use an explicit fourth-order Runge-Kutte Kutta scheme to compute the time integration. As our encoder identifies latent spaces with good smoothness conditions, we have the luxury to use relatively large, fixed step sizes, which serve as an additional regularization that biases towards learning models with efficient inference. Learning I' directly using this latent-space rollouts is empirically challenging, due to highly non-convex energy landscape and gradient explosions. Instead, we split the training process in stages: single-step pre-training, and fine-tuning using checkpointed neural ODE adjoints (Chen et al., 18). We describe them briefly in below:
steple-sepp pretraining We train the latent dynamical model by first optimizing single-step loss:
0 01per 6t., bn+1 In
Esingle (3) = Σ
(()(0..))
(13)
which simply attempts to match r3 with the first order finite difference approximation of the time derivative in 0. Note that cumulative errors beyond the (tnstn+)) window do not affect the learning of B.
Mtlti-step fine-tuning Although learning with single-step loss defined by Eq. (13) alone has been demonstrated to achieve success in similar contexts when optimized to high precision and combined empirically
with practical techniques such as noise injection (Pfaff et al., 2020; Han et al., 2022), observe that additional fine-tuning 3 using Eq. (12) as the loss function helps a great ideal in further improving the accuracy and stability of our inference rollouts, i.e. it is important to teach the model

---

Published as a conference paper at ICLR 2023
to "look few more steps head" during training. The fact that encoders generate smooth latent spaces, together with the use of fixed, large-step time integration, helps to reduce the cost of the otherwise expensive adjoint gradient compute, which had been one of the main drawbacks of such long-range neural-ODE training. We point out that there exists ample literature on multi-step training of differential equations, e.g. Keisler (2022); Brandstetter et al. (2022). For problems considered in this work, however, this simple two-stage training proved empirically effective.
To conclude this section, mention in passing that the recently proposed Neural Galerkin approach (Bruna al., 2022) can be used learn the dynamics by fitting 0:
0,u(z, -) 0 v (x) VoV(2) = Fz)
(14)
However, as demonstrated our experiments, this approach can at times prone to bad conditioning in the stiffness matrix, and relies on having complete knowledge about JF.
4 RESULTS
4.1 SETUP
Datasets We demonstrate the efficacy of the proposed approach on three exemplary PDE systems: (1) Viscid Burgers (VB) (2) Auramoto-Sivashinkk (KS) and (3) Kortweg-De Vries (KdV). These systems represent a diverse mix of dynamical behaviors resulting from advection, characterized by the presence of dissipative shocks, chaos and dispersive traveling waves respectively. Each is simulated on a periodic spatial domain using Explicit-Implice solver in time coupled with a pseudo-spectral discretization in space. We used the spectral code in jax-cfd (Dresdner et al., 2022) to compute the datasets. This solver provides very accurate and dispersion free ground- truth references. Detailed descriptions of the systems and the generated datasets can be found in Appendix A.I.
Benchmark Methods We compare to several recently proposed state-of-the-art methods for solving PDEs continuously in space and time: Neural Galerkin (NG) (Bruna et al., 2022) (at high and low sampling rates), DeepONet (Lu al.. 2021), Fourier Neural Operator (FNO) (Li et al., 2020b) and Dynamic Mode Decomposition (DMD) (Schmid, 2010). We also investigate variants of our model that uses the same encoder architectures but convolution-basee decoder and comparable latent space dimensions (abbreviated NDV neural decoder variant). Details of those methods are in Appendix A.4.
Metrics For each system, we use a training set of 1000 trajectories with at least 300 time steps each. For evaluation, trained models then used to generate multi-step rollouts 100 unseen initial conditions. The rollouts are compared to the ground truths, which are obtained by directly running a high-resolution pseudo-spectral solver the evaluation initial conditions. To quantitatively assess the predictive accuracy of the rollouts, we compute the point-wise relative root mean squared error (relRMSH)
rue) U
u true
reIRMSE(uPre
(15)
u
Irue
as a measure for the short-term prediction quality. We also compare the spatial energy spectra (norm of the Fourier transform) of the predictions to that of the ground truths via the log ratio:
EntrgyRatio, (ffe) = llog (E.reed (fa)/E, (f))))
(16)
E, k/L)
Σ u(x,, tr) exp (- (-j2 LL)
where L - NAr and f² -1. r denotes the range of the prediction. This metric (0 for perfect prediction and positive/negative values mean over-/under prrediccer-rrietingggg. energy in the corresponding frequencies) provides translationiinvaria. predictions. Lastly, we record the inference cost of each method in terms of the wall clock time (WCT) taken and the number of function evaluations (NFE) per simulation time unit.
assessment of the frequency features of the long-term

---

Published as a conference paper at ICLR 2023
Table Whether the benchmark methods are able to produce a work- ing inference model for each system. Working is defined as exceeding 100% reIRMSE in 40 rollout steps (twice the training range for our model).
109
102
narm
relative error
10"4
10
10
10-3
VB
KS
KdV
10
10" 10
Ours-NFA
10
10a
10
10a
Ours-NDV
NG-Lo
Figure 3: Relative reconstruction error (left, solid), consis- tency error (left, dashed) and total variation norm (right) vs. ? different learning rates for KS: x 10-1 -; 1 10-1 See metric definitions in Appendix A.5.
NG-Hi
DeepONet
FNO
= 3 x 10-5;
DMD
Sample Rollouts
ReIRMSE
EnergyRatio
10a
VB
-20
1.0
-30
10
307
10
timcc
[egq
KS
10a
-10
10-1
10"
10P
kav
10
-10
10-2
15-1
10'
101
freg
Ours-NFA: - NG-Lo: NG-Hi: - Ours-NDV:
DeepONet: FNO: -; DMD:
Figure 4: Sample long-term rollouts (left), point-wise error vs. time and long-range energy spectrum log ratios (right) plots. Rollouts are plotted in (z, t) plane with upper rows showing the ground truths and lower Ranges for energy ratios plotted: VB 3; KS 80; KdV 16. NG-Lo/NG-HHi are NG runs with a sampling rate of 198 and 600 respectively.
corresponding predictions. Columns correspond to different initial conditions.
4.2 MAIN RESULTS AND ABLATION STUDIES
Accurate and stable long-term roll-outs Fig. contrasts our method to the benchmarks. While each benchmarked method fails to produce working surrogate for at least of the systems (defined and summarized in Table 1), our method able to model all three systems successfully.
Qualitatively, shocks in VB, chaos patterns in KS, and dispersive traveling waves in KdV are all present in our model rollouts in ways that are visually consistent with the ground truths. Our models produce the lowest point-wise errors among working benchmarks for both KS and KdV (which are
8

---

Published as a conference paper at ICLR 2023
Table 2: Comparison of inference speeds in terms of WCT and NFE statistics. WCTs are measured in millseconds, NFEs are measured by an adaptive-step integrator. Numbers are normalized to cost per unit simulation time. Solver is run at the same resolution used to generate training data. High NFE count of NG-Hi results from running into stiff numerical regime before breaking down (around 120 rollout steps).
VB
KS
KdV
WCT
NFE
WCT
NFE
WCT
NFE
mean
std
mean
std
mean
mean
std
mean
std
mean
std
Solver
8.4
0.5
5.2
0.3
1041
13.1
Ours-NFA
3.8
0.9
16.0
4.1
1.7
2.2
5.3
0.2
52.1
18.4
164.7
40.3
Ours-NDV
1.1
0.6
7.1
2.7
18.0
5,4
71.4
19.5
NG-Lo
30.2
8.5
27.5
10.8
167.9
21.4
114.7
13.3
NG-Hi
49.6
7.7
27.5
10.8
319.4
33.5
111.7
11.7
35670
29035
11651
9502
FNO
18.1
14.5
100.0
0.0
0.4
0.4
5.0
0.0
DMD
2.4
0.4
100.0
0.0
the more complex examples) and come very close to the best- pest-performing NG model for VB. In the long-range prediction regime, our models are able to remain stable for the longest period of time and exhibit statistics which perfectly match the truth in low frequencies and are the best approximating one in high frequency components.
The importance of ansatzes The right choice of the ansatz plays two roles. Firstly, learning the encoder becomes easier when the key structures are already embedded in the decoder this is supported by the observation that our model easily learns encode the shocks in VB while NDV fails to do so even allowed to operate larger latent space. Secondly, the ansatz reduces the spurious fregerryy energy components in the state, which helpful in improving the stability of rollouts, as these energies are likely to accumulate in nonlinear fashion. This is the main reason why our model able have better errors than the closely related NDV model in KS and KdV.
Consistency regularization leads to smoothness The impact of y in Eq. (11) evident from the observation that (Fig. 3 left) as the strength of regularization increases the encoders resulting from minimizing the loss in Eq. (1I) have lower reconstruction error. also induces smoother latent weights trajectories, measured in the component-wise total variation norm (Fig. 3 right). The smoothness is beneficial for downstream training to learn efficient dynamical models that permit taking large integration time steps.
Inference efficiency Statistics the inference cost of each method are summarized in Table 2. Our models have the best efficiencies amongst the compared methods. Their low NFE counts confirm that generous time steps being taken during time integration. Note that the aforementioned trajectory smoothness is not a sufficient condition to guarantee such efficiency, supported by the observation that the PDE-induced dynamics found by NG typically require more NFES per unit time (can even run into prohibitively stiff regimes, like in KdV) despite starting from the same encoded initial conditions. The explicit large-step integration that we use to learn the latent dynamical models proves be an effective regularization to alleviate this issue the NDV models are also able to achieve efficient inference benefiting from the same factor. However, without the help of consistency regularization, this comes at compromise in accuracy and robustness compared to our main models.
CONCLUSION
In this work we have introduced method to learn reduced-order surrogate models for advection- dominated systems. We demonstrate that through specially tailored ansatz representations and effective regularization tools, method is able identify weight spaces with smooth trajectories and approximate the dynamics accurately using hyper networks. Experiments performed on challenging examples show that our proposed method achieves top performance in both accuracy and speed.

---

Published as a conference paper at ICLR 2023
REFERENCES
David Amsallem, Matthew J Zahr, and Charbel Farhat. Nonlinear model order reduction based on local reduced-order bases. International Journal for Numerical Methods in Engineering, 92(10): 891-916, 2012.
Patricia Astrid, Siep Weiland, Karen Willcox, and Ton Backx. Missing point estimation in models described by proper orthogonal decomposition. IEEE Transactions on Automatic Control, 53(10); 2237-2251, 2008. doi: 10.1109/AA.2008 20000 2006
Nadine Aubry, Philip Holmes, John L. Lumley, and Emily Stone. The dynamics of coherent structures in the wall region of a turbulent boundary layer. Journal of Fluid Mechanics, 192:115-173: 173,, 1988. doi: 10. 1017/00200 12088001 118.
Ibrahim Ayed, Emmanuel de Bezenac, Arthur Pajot, Julien Brajard, and Patrick Gallinari. Learning dynamical systems from partial observations, 2019.
Yohai Bar-Sinai, Stephan Hoyer, Jason Hickey, and Michael Brenner. Learning data-driven
discretizations for partial differential equations. Proceedings of the National Academy of Sciences, 116631):15344-15549 July 2019. ISSN 0027-8424, 1091-6490, doi: 10. 10.1073/pnas.881 1814pp88...11188111111 1814058116 URL
org/ /twp//ill/l/a/ee/eteeeeeotaltttt 1073 pwws..pass..... 18188144 11488116.
Maxime Barrault, Yvon Maday, Ngoc Cuong Nguyen, and Anthony T. Patera. An 'empirical
interpolation' method: application to efficient reduced-basis discretization of partial differential
equations. Comptes Rendus Mathematique, 339()):667-67) 2004. ISSN 1631-073X. doi:
200.000.000000000000000000000000 URL https: http::///wwwwweeeo http:///wsiieeedeeeeeeeeeee com/ scleence/artiille//1////1.. 24000044488
Peter Benner, Serkan Gugercin, and Karen Willcox. A survey of projection-based model reduction methods for parametric dynamical systems. SIAM Review, 57(4):483-531 2015. doi: 10.1 10.1137/
130932715. URL https: / //oo/i///////// /0og/0.0........../.00///001//// 1137/130992211
John P. Boyd. Chebyshev and Fourier Spectral Methods. Dover Books Mathematics. Dover Publications, Mineola, NY, second edition, 2001. ISBN (48641 1834 978048641 1835.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL
http://githut com/ hong//gole//////// jax.
Johannes Brandstetter, Daniel Worrall, and Max Welling, Message passing neural pde solvers, 2022.
URL / ort.//.rrr//aa.. abs /tt.ab///222.22 220. 03376.
Joan Bruna, Benjamin Peherstorfer, and Eric Vanden-Eijnden Neural galerkin scheme with active learning for high nigh-dimensional evolution equations. arXiv:2203.00 360 [math.NA]. May 2022. URL https:/ / arxiv.
att://b//22200 01360.
arXiv: 203.00360.
Oscar P. Bruno, Jan S. Hesthaven, and Daniel V. Leibovici. FC-based bhock-dynamice solver with neural-network localized antificial-viscosiiy assignment. arXiv:2111.01311 [cs, math], November 2021. arxiv: 211 1.01115.
Claudio G Canuto, M. Yousuff Hussaini, Alfio M. Quarteroni, and Thomas A. Zang. Spectral Methods: Evolution to Complex Geometries and Applications to Fluid Dynamics (Scientific Computation). Verlag, Berlin, Heidelberg, 2007. ISBN 3540307273.
Saifon Chaturantabut and Danny C. Sorensen. Nonlinear model reduction via discrete empirical interpolation. SIAM Journal on Scientific Computing, 32377,7777, 2010. doi: 10.1137/ 090766498. URL https: doii .rai.org/1100
Ricky T. Q. Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Beaiii Bianchi, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 31. Curran As- sociates, Inc., 2018 URL ht 69386f d fwd.eee6666ec 2
www.:eeeeediinn....nneee.e.. nwwrips..p... https/sn..s.......aa.eee////e//e //1///iie// pdf.
10

---

Published as a conference paper at ICLR 2023
Wenqian Chen, Qian Wang, Jan S. Hesthaven, and Chuhua Zhang.
Physics informed machine Physics-informed
learning for reduced-order modeling of nonlinear problems. Journal of Computational Physics, 446:110666, 2021. ISSN 0021-9991. doi: http://db..orr//0000061...gg2222 2002 URL https: / //www. www..scienced. rect. com/ sci wwcienceeaaaiii http://ceeeceeaaraaceeppaa/////s /tt.///0201999990009999
Francisco Chinesta, Pierre Ladeveze, and Elias Cueto. A short review on model order reduction based on proper generalized decomposition. Archives of Computational Methods in Engineering, 18(4):395-444 2011.
J. W. Cooley and J. W. Tukey. An algorithm for the machine calculation of complex Fourier series. Math. Comput., 19(90))297-301, 1965. ISSN 00255718 88 10886842. URL http: //ww://wwwwwww/// jstor. org/ wttssa/ebbe3335555
Filipe de Avila Belbute-Pere Yi fan Chen, and Fei Sha. HyperPINN: Learning parameterized differential equations with physics-informed hypernetworks. In The Symbiosis of Deep Learning and Differential Equations, 2021.
J.R. Dormand and P.J. Prince. A family of embedded runge-kutta formulae. Journal of Compu-
tational and Applied Mathematics, 6(1):19-26, 1980. ISSN 0377-0427. doi: https://doi.oo science/
10.100660771-0000080800000011 URL https: http://wwwwwccccc www.sciience.. axricle/ss1//07110609900030
Gideon Dresdner, Dmitrii Kochkov, Peter Norgaard, Leonardo Zepeda-Núñez, Jamie A. Smith, Michael P. Brenner, and Stephan Hoyer. Learning to correct spectral methods for simulating turbulent flows, 2022.
Weinan E and Bing Yu. The deep ritz method: A deep learning- numerical algorithm for solving variational problems. Communications in Mathematics and Statistics, 6(1):1-12, 2018. doi: 10. 1007/s40304-018-011
URL /ttp
Hamidreza Eivazi, Mojtaba Tahani, Philipp Schlatter, and Ricardo Vinuesa. Physics-informed neural networks for solving Reynolds- Ntvier-Stokes equations. arXiv:2107.1077 [physics], July 2021. URL http:://araat http://wwww.w..aa//////aaaaaaaa/
10711. arXiv: 2107.10711.
Yuwei Fan, Jordi Feliu-Fabà, Lin Lin, Lexing Ying, and Leonardo Zepeda-Núñezz A multiscale neural network based on hierarchical nested bases. Research in the Mathematical Sciences, 6, Mar. 2019. ISSN 2197-9847, doi: 10. 10.10077/44067--19--1655
Hugo Frezat, Julien Le Sommer, Ronan Fablet, Guillaume Balarac, and Redouane Lguensat. A posteriori learning for geastrggostt turbulence parametrization. arXiv, April 2022.
B. G. Galerkin. Series occurring in various questions concerning the clastic equilibrium of rods and plates. Vestnik Inzhenernov Tekhnikov, 19:897-908, 1915.
Han Gao, Matthew J. Zahr, and Mun- Wang. Physics-informed graph neural Galerkin networks:
A unified framework for solving PDE-governed forward and inverse problems. Computer Methods in Applied Mechanics and Engineering. 390, 2022. ISSN 0045-7825. doi: https://doi.or//1..1011 j.cma.2021.1
11452. 14502.
Rudy Geelen, Stephen Wright, and Karen Willcox. Operator inference for non-intrusive model reduction with nonlinear manifolds. arXiv:2205.00 02.304 [math.NA], May 2022. URL https: wttp://///bb////////////////// 2205. 02304. arXiv: 2205.02304.
Xu Han, Han Gao, Tobias Pfaff, Tian-Xun Wang, and Li- Li-Ping Liu. Predicting physics in mesh-reduced space with temporal attention. CORR, abs/2201.09 133911,, 2022.
09113.
Kaiming He. Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016.
Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (gelus). arXiv: 1eo::666608.11 2016.
arXiv preprint
11

---

Published as a conference paper at ICLR 2023
Lars Hörmander. The Analysis of Linear Partial Differrntial Operators. III: Pseudo-Differnttall Operators, volume 63 of Classics in Mathematics. Springer, Berlin, 2007.
Lars Hörmander. The Analysis of Linear Partial Differential Operators. IV: Fourier Integral Opera- tors, volume 63 of Classics in Mathematics. Springer, Berlin, 2009.
K. Hornik, M. Stinchcombe, and H. White. Universal approximation of unknown mapping and its derivatives using multilayer feedforward networks. Neural networks, 365)551-550, 1990.
Sergey Ioffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In International conference on machine learning, pp. 448-4544 PMLR, 2015.
Ryan Keisler. Forecasting global weather with graph neural networks. arXiv preprint arXiv:2202.07575 2022.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Yoshua Bengio and Yann LeCun (eds.), 3rd International Conference Leurning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Truck Proceedings, 2015. URL http: //aaxxi..rg//////........o./.t.t abs///////a///aaaaaa/..a/a/aaaa/ aaaaabaabb111 /1122.41..211 6980.
Dmitrii Kochkov, Alvaro Sanchez-Gonzales Jamie Alexander Smith, Tobias Joachim Pfaff, Peter Battaglia, and Michael Brenner. Learning latent field dynamics of pdes. In Third Workshop Machine Learning and the Physical Sciences (NeurIPS 2020), 2020.
Dmitrii Kochkov, Jamie A Smith, Ayya Alieva, Qing Wang, Michael P Brenner, and Stephan Hoyer. Machine learning- accelerated computational fluid dynamics. Proc. Natl. Acad. Sci. U. S. A., 118 (21), May 2021. URL https: wwww. pnas org/ content/t///////o content/118////eeee
Nathan Kutz, Steven I.. Brunton, Bingni W. Brunton, and Joshua I.. Proctor. Dynamic Mode De- composition. Society for Industrial and Applied Mathematics, Philadelphia, PA, 2016. doi: 10.113771.9786611174667 URL https: /app:://epuae /epubs.. siam. org/ otgg/oo/abs//....../oo/ 1137/1. 978161197450
Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Multipole graph neural operator for parametric partial differential equations. In Proceedings of the 34th International Conference on Neural Information Processing Systems, NIPS' 20, Red Hook, NY, USA, 2020a. Curran Associates Inc ISBN 9781 188111 13829546..
Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Fourier neural operator for parametric partial differential equations, 2020b.
Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Fourier Neural Operator for Parametric Partial Differential Equations. arXiv:2010.0888: [cs, math], May 2021. URL http: wttpaatwwwww..a////////aaaaaaaea 08895. arXiv: 2010.08895.
Bjöm List, Li- Chen, and Nils Thuerey. Learned Turbulence Modelling with Differentiable Fluid Solvers. arXiv:2202.06600 [physicx], February 2022. URL htt::t/t /ttp::://///vvv.// http://awwwi...ooaa//////////// 06988. arXiv: 2202.06985
Lu Lu, Pengzhan Jin, Guofei Pang, Zhongqiang Zhang, and George Em Karniadakis. Learning nonlinear operators via DeepONet based on the universal approximation theorem of operators. Nature Machine Intelligence, (0))29,,,2,, 2021.
Per-Gunnar Martinsson. Fast Direct Solvers for Elliptic PDEs. Society for Industrial and Applied Mathematics, Philadelphia, PA, 2019. doi: 10. s am. wwwg///ooooooooo aabs/1b/000 0 11311113111 137/1 01111/997711111977611
URL pittts::/eppbee
Ben Mildenhall, Pratul P. Srinivasan. Matthew Tancik, Jonathan T. Barron, Ravi Ramamoorthi, and Ren Ng. Nerf: Representing scenes neural radiance fields for view synthesis. In ECCV, 2020.
12

---

Published as a conference paper at ICLR 2023
Siddhartha Mishra. A machine learning framework for data driven acceleration of computations of differential equations. Mathematics in Engineering, 1(1):118-146, 2018. ISSN 2640-3501 doi: 10.3934/Mine.20158
URL https:
com/ wwom/aaee artiicltteaa..aa /doi/
10. (0.393. 30..394/MMimn ne. 201 1.18.
Jeong Joon Park, Peter Florence, Julian Straub, Richard Newcombe, and Steven Lovegrove. Deepsdf: Learning continuous signed distance functions for shape representation. Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 165-174, 2019.
Benjamin Peherstorfer. Model reduction for transport-dominated problems via online adaptive bases and adaptive sampling. SIAM Journal on Scientific Computingg, 42853 2020. doi: 10.1 (9.1133//MM25572777 URL
wwt.:0..........oo0o 1137 9MM257775.
Tobias Pfaff, Meire Fortunato, Alvaro Sanchez-Gonzale,, and Peter W Battaglia. Learning mesh- based simulation with graph networks. arXiv preprint arXiv:2010.034000 2020.
A. Pinkus. Ergebnisse der Mathematik und ihrer Grenzgebiete. 3. Folge A Series of Modern Surveys in Mathematics, volume 1. pringer Berlin Heidelberg, 2012.
M. Raissi, P. Perdikaris, and G. E. Karniadakis. Physics-informel neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations. Journal of Computational Physics, 578:686-707, February 2019. ISSN 0021-9991. doi: j0.10166j.jep.201.. 10.045. URL https:/ www.wwwww.oo
wttpscccccss/.eeeeeeeeeeee...eee science/
Prajit Ramachandran, Barret Zoph, and Quoc V. Le. Searching for activation functions. In 6th International Conference on Learning Representations, ICLR 2018, Vancouver, BC, Canada, April 30 - May 3, 2018, Workshop Track Proceedings. OpenReview.net, 2018. URL https: // //opnreeview..o net/ wtw.:r/wwwwe..et/ooom...... Herrumiidkkuuq2EPPP.
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In Nassir Navab, Joachim Hornegger, William M. Wells, and Alejandro E. Frangi (eds.), Medical Image Computing and Computer-Assisted Intervention - MICCAI 2015, pp. 234-241, Cham, 2015. Springer International Publishing. ISBN 978-3-319-24574-4.
Alvaro GonzalzZ Jonathan Godwin, Tobias Pfaff, Rex Ying. Jure Leskovee, and Peter W. Battaglia. Learning to Simulate Complex Physics with Graph Networks. arXiv:2002.094/ 09405 [physics, stat], September 2020. URL http:///wwxxxv.oo/aaa///a/aase 2002.09405.
org wwtp/////////oooooo//////////o https/////2..esooosee 09405. arXiv:
Peter Schmid. Dynamic mode decomposition of numerical and experimental data. Journal of fluid mechanics, 656:5-28, 2010.
Kimberly Stachenfeld, Drummond B. Fielding, Dmitrii Kochkov, Miles Cranmer, Tobias Pfaff, Sanchez-Gonzalez. Learned Coarse Models for Efficient Turbulence Simulation. arXiv:2112. 2175 [physics], January 2022.
Jonathan Godwin, Can Cui, Shirley Ho, Peter Battaglia, and Alvaro Sanchez- Gonzalez.
URL http: /twpvv.rvvvvv/iv///vv//... org/ abs /2112211222 15275. arXiv: 2112.15275.
Alasdair Tran, Alexander Mathews, Lexing Xie, and Cheng Soon Ong. Factorized Fourier Neural
Operators. arXiv:2111.13802 [cs], November 2021. arXiv: 2111.13802.
Lloyd N. Trefethen. Spectral Methods in MATLAB. Society for industrial and applied mathematics (SIAM), 2000.
Rui Wang, Karthik Kashinath, Mustafa Mustafa, Adrian Albert, and Rose Yu. Towards physics- informed deep learning for turbulent flow prediction. In Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, 2020. URL http:
doi Coi..rrg/10. 114.1114/3333 486 300.9444884434111111 3403198.
Sifan Wang, Hanwen Wang. and Paris Perdikaris. Learning the solution operator of parametric partial differential equations with thysics-informar deeponets. Science Advances. 7(40) 2021. doi: 10.1 126/sciadw.abi866666 URL https: /www. www.wwscciieee scwenceieccccccc......co...e org/ doi/abs////////// /10 1126/ abi 8605
13

---

Published as a conference paper at ICLR 2023
K. Willcox. Unsteady flow sensing and estimation via the gappy proper orthogonal decompo- sition. Computers & Fluids, 35(2):208-22) 226. 2006. ISSN 0045-7930. doi: https://doi.or//
10.0016f..ommftid1200.1.....0000 URL https:/ com. science/ articcle/e1/55555700000011
Yaohua Zang, Gang Bao, Xiaojing Ye, and Haomin Zhou. Weak adversarial networks for high- dimensional partial differential equations. Journal of Computational Physies, 411:109409, 2020. ISSN 0021-9991. doi: https:/dooooog///.00oooo..pgg010 URL https: /tt:://wwwww /www.
sci enced i rect com ottp 832.
14

---

Published as a conference paper at ICLR 2023
A APPENDIX
A.1 DATASETS
In this paper, we consider PDES of the form
d,u F(u) == Du N(u),
(17)
plus initial and boundary conditions. Here D is a linear partial differential operator, and N is nonlinear term. For the models considered in this paper, the differential operator is typically either diffusive, D 8e dispersive, D 0%. or hyper-diffusive, D 0 and the nonlinearity is convective term, N = 0z (u²)= ud, u. Roughly, diffusion tends to blur out the solution, rendering the solution more homogeneous; dispersion produces propagating waves whose velocity depends on the Fourier content of the solution; finally, advection also propagates waves whose velocity depends on the local value of the solution.
In practice, PDEs are solved by discretizing space and time, which converts the continuous PDE into a set of update rules for vectors of coefficients to approximate the state u in some discrete basis, e.g., on grid. For time-dependent PDEs, temporal resolution must be scaled proportionally to spatial resolution to retain an accurate and stable solution, so the runtime for PDE solvers scales at best like O(n ), where d is the number of spatial dimensions and is the number of discretization points along any dimension. In general, however, dispersion errors tend to dominate the tre-asympootic regime one would need a large amount of points in space, and in turn a large number of time steps.
To circumvent this issue, we choose a pseudo-spectral discretization, which is known to be dispersion free, due to the exact evaluation of the derivatives in Fourier space, and it possesses excellent approximation guarantees (Trefethen, 2000). Thus, relatively few discretization points are needed to represent solutions that are smooth.
We assume that, after appropriate rescaling in space, (xxx()) [0, 2m| x R+ + R is one-dimensional, 2 2y-periodic and inteerabete for all times t. Consider the Fourier coefficients ã' of u(z,)) (x. truncated to the lowest K + frequencies (for even K):
[o 2 u(x, t)e t)e--)ee---e
T' = (ū' K//'' 0t (2//2) where
ik- dxx.
(18)
2x
G is a vector representing the solution at time I, containing the coefficients of the Fourier basis eike for k E (-K/2..... K/2). The integral ut is approximated using a trapezoidal quadrature on K + 1 points, i.c., sampling u(z, t) on an equi-spaced grid, and then computing the Fourier coefficients efficiently in log-linear time using the Fast Fourier Transform (FFT) (Cooley & Tukey, 1965). Differentiation in the Fourier domain is a diagonal operator. In fact, it can be calculated by element-wise multiplication according to the identity ū,ūg ik.ū.. This makes applying and inverting linear differential operators trivial since they are simply element-wise operations (Trefethen, 2000).
The nonlinear term N (1u) can be computed via an expensive convolution, or by using Plancherel's theorem to pivot between real and Fourier space to evaluate these terms quasilinear time.
This procedure leads to a system in Fourier domain of the form
0, Dũ +N('')
(19)
where D encodes the D operator in the Fourier domain and is often diagonal matrix whose entries only depend on the wavenumber N encodes the nonlinear part. This system may be solved using a Crack-Nicolson Runge- Kutta scheme (Canuto et al., 2007). In this case,
4th order implicit-explins we treat the linear part implicitly and the nonlinear one explicitly.
We point out that this transformation still carries some of the difficulties of the original problem. In particular, as l increases the conditioning number of A will also increase, which may require refining the time-step in order to avoid instabilities and satisfy the Shannon-Nyquist sampling criterion.
The solver used to generate data is implemented using jaxcffd where we use 512 grid points.
Viscid Burgers Equation We solve the equation given by in [-1, x R
0,u + uº,uu vθ
(20)
15

---

Published as a conference paper at ICLR 2023
with periodic boundary conditions. We follow the set up in (Wang et al., 2021), in which v 0.01, and the initial condition is chosen following a Gaussian process.
We point out that the combination of the high-viscosity and the smoothness of the initial conditions results in trajectories that are fairly stable with waves travelling only small distances. No sub-grid shocks are formed.
Kuramoto-Sivashinskk Equation We solve the equation given by
d,u + ud,u u + vd - vu
in 00, x
(21)
with periodic boundary conditions, and L = 64. Here the domain is rescaled in order to balance the diffusion and anti-diflusion components so the solutions are chaotic (Dresdner et al., 2022).
The initial conditions given by
no Σ k=1;-1
un (3r)
ag sin(wy 0z),
(22)
where wy is chosen randomly from [T/ 2t/ T,3r 3m LL) aj is sampled from a uniform distribution in 0.5), and phase o follows uniform distribution in [0, 2Tr.. We use ne 30.
KdV Equation
We solve the equation
- - u + Paaz
in L-22, L/2]
(23)
with periodic boundary conditions, and 32. Here the domain is again rescaled to obtain travelling waves. We point out that this equation is highly dispersive, so an spectral methods is required to solve this equation accurately.
The initial conditions for the KdV equation also follow Eq. (22), but with = 32 and n. = 10. This choice of initial conditions allows for sol-soniton solutions. with few large wave packages being propagated together with several small dispersive waves in the background.
Note that to our knowledge this is first time that ML method able compute non-soliton solution of the KdV equation for long times.
A.2 ADDITIONAL TRAINING AND EVALUATION DETAILS
For all training stages, we use the Adam optimizer (Kingma & Ba, 2015) with 31 = 0.9, |3 = -0.99 and = 8
Ansatz evaluation Prior carrying the multi-stage learning protocol described in the main text, we first perform a preliminary evaluation of the ansatz network via batch fitting [0"} to a random selection of snapshots {w:} drawn from the dataset, i.e. for each we minimize the representation loss
N Σ I 1 ? N
1
Coop (18=))
(24)
which provides us with an assessment of the expressive power of the particular functional form of t with respect to the problem at hand. Note that this process, otherwise known as auto-decoder training (Park et al., 2019), can be done independently and we can easily compare amongst multiple potential candidates and rule out ones which are not suitable. In particular, as a rule of thumb, we exclude ansatz expressions that cannot achieve below 1% relative reconstruction RMSE. As an example, in Table 3 we show the metrie sweep of different decoder configurations for KdV. Similar procedures carried out to determine the ansatzes for VB and KS.
16

---

Published as a conference paper at ICLR 2023
Table 3: ReIRMSE of the decoder approximate snapshots of the KdV training trajectories resulting from different hyperparameters by minimizing Eq. (24). In this case we consider a dyadic scaling of the frequencies. We mark bold font reIRMSE below one percent.
Act. fun.
relu
swish
Feats/No. freqs
3
6
9
3
6
3
6
9
(4,
0.1095
0.0255
0.0151
0.0803
0.0164
0.0110
0.1198
0.0164
0.0071
(4, 4,4)
0.1010
0.0249
0.0142
0.0511
0.0129
0.0085
0.0799
0.0120
0.0059
(8, 8)
0.0273
0.0120
0.0117
0.0343
0.0095
0.0118
0.0637
0.0094
0.0054
8,88,8)
0.0195
0.0101
0.0094
0.0178
0.0082
0.0107
0.0333
0.0066
0.0050
(16, 16)
0.0113
0.0090
0.0124
0.0202
0.0100
0.0181
0.0381
0.0076 0.0073
(16, 16, 16)
0.0095
0.0079 0.0096 0.0115 0.0099
0.0178
0.0186
0.0057 0.0054
Table 4: Training specifications for encoder learning. Learning rate (LR) follows an exponentially decaying schedule, i.e. multiplying by a factor after fixed period of steps.
System
y
Batch Size
Training Steps
Base LR
LR decay
Steps Per Decay
VB
10
16
8M
10-4
0.912
160K
KS
10
32
4M
3 x 10-4
0.892
80K
Kdv
10a
16
4M
3 x 10-4
0.945
40K
Encoder training This step is performed by minimizng loss funciton Eq. (11). It acts another round of validation to help us see if a particular ansatz have sufficient accuracy when being encoded. auto-decoding evaluation since the encoder
This is stronger condition than the Above-mentionee adds extra regularization. We again set the bar to be 1% relative reconstruction error, which works well for us empirically.
Notably, ansatzes that have trainable frequencies (w,, in Eq. (10)) fail this validation despite passing the auto-decoder test. This suggests that having these extra degrees of freedom the ansatz is not as useful as hard-coding the periodic boundary condition directly, in terms of encoder learning.
Of all ansatzes that pass the 1% criterion, we choose the one with the fewest parameters, corresponding to the ones displayed in Table 6.
The encoder training specifications are summarized in Table 4.
Dynamics training As mentioned in the main text, a two-stage pretrain/fine-tue protocol is adopted. The training parameters are listed in Table 5. We additionally employ gradient clipping (scaling norm to 0.25) to help stabilize training.
Device info
All training and inference runs are performed on single Nvidia V100 GPUs.
A.3 MODEL DETAILS
The hyperparameters of the ansatz used for each system are summarized in Table 6. The columns are explained as follows:
Ansatzes
features: MLP hidden layer sizes
activation: nonlinearities applied after each hidden layer of the MIP: swish (Ramachandran et al., 2018) or sin function
frequency profile: if linear, }wx) follow an arithmetic sequence {T/ L, 2n L, 3n 3r J.. if dyadic, the sequence is geometric (x/L, 2n 4T .... in powers of
frequency dimensions: each non-zero frequency has associated sin and cos functions with distinct phases; the input and output dimensions of the ML. are both 2 x No. Freqs - |[as Zero Freq].
17

---

Published as a conference paper at ICLR 2023
Table Training specifications for dynamics learning. Pretrain settings are the same for all systems. Seq Length refers to the number of steps over which to optimize the dynamical model parameters (= for pretraining: 2 for fine-tuning)
Stage
Batch Size Training Steps Seq Length
Base LR
LR decay
Steps Per Decay
Pretrain
32
4M
2
10
0.955
40K
Fine-tune, VB
32
IM
10
3 x 10-4
0.955
10K
Fine-tune, KS
32
1M
20
5
10-5
0.955
10K
Fine-tune, KdV
32
1M
20
1 10-5
0.955
10K
Table 6: Hyperparameters for ansatz networks.
System
MLP
Frequency
No. Params
Features
Activation
No. Freqs
Profile dyadic
Has Zero Freq
VB
(4, 4)
swish
No
KS
(8, 8)
sin
3
linear
No
188
Kdv
(8, 8)
swish
dyadic
Yes
297
Encoder The encoder architecture consists multiple levels of double (one downsampling and one regular) residual blocks. Each residual block follows the sequential layer structure: Conv BatchNorm - sin activation - Conv + BatchNorm, where downsampling is applied via striding in the first Cony layer when applicable. Overall, each level downsamples the spatial dimension and increases the number of channels, both by a factor of A standard fully-connected layer is applied the end of the network obtain the desired output dimension.
Encoder hyperparameters differ very little across systems, with VB encoder using 5 downsampling levels, compared to 4 used in KS and KdV (see Table 7).
Hyper U-Net
The transforms represented by the arrows in Fig. 2 are follows:
channel-wise projection: locally connected layer (with kernel size I and no bias); projects weights from D., dimensions (first layer) and back (last layer)
channel-wise mixing: sequence transformation Dense - Swish activation - Dense Residual - LayerNorm applied to channels within the same layer (or to the global embedding channels); parameters are unshared among layers
locally connected layer compressing all weights channels within the same layer into single layer embedding of dimension Dú; similarly, at the global level, all layer embeddings are compressed into a single global embedding of dimension D3
downsanpling:
upsampling: the inverse of the downsampling transform mapping single size-D, embeddings back multiple size-D,, ones (also the size-D, global embedding into Nr x Di layer embeddings)
copy: similar to the original U-Net, channels from the downsampling stage are concatenated during the upsampling stage
Specifically for our ansatzes, we consider each MLP layer to encompass the weights and bias (W, 6} associated with the layer compute oWW + b). The phases {a=*} are considered a separate layer collectively.
Our Hyper U-Net architectures for different systems vary in the hyperparameters (Dw Dr, Dg). Their values are listed in Table 7.
A.4 BENCHMARK MODEL DETAILS
A.4.1 NEURAL GALERKIN
The Neural Galerkin approach (Bruna et al., 2022) seeks to compute evolution equations using calculus of variations in time-continuous fashion.
18

---

Published as a conference paper at ICLR 2023
Table Hyperparameters for encoder and hyper U-Net
System
Encoder
Dynamical
Model (Hyper U-Net)
Levels
No. Params
Dw
D1
Dg
No. Params
VB
186,268
4
128
512
1,846,612
KS
218,116
4
512
1024
17,963,932
Kdv
218,116
4 512
1024
17,963,932
At a high level, the Neural Galerkin method writes an evolution equation for the weights following:
M F(O))
(25)
where
M(0) = /' VoVo(r) @ VoVo(r)dr, and F(0)= { (oVo(t)F(VoV((())l
(26)
is the neural ansatz described in Eq. (10) and F is given by the underlying PDE in Eq. (17).
The matrix (0) is computed by sampling Vo(z). We employ uniform sampling on the domain to approximate the integral, namely,
Naurpies
(
Σ
VoVo(I,) & VoWo(Zy).
(27)
Naarples
We also use a trapezoidal rule to compute the integrals for the KdV equation, which works better for this particular equation.
In order to provide a fair comparison, we have attenuated two main issues with this method:
1. The ansatz, which is chosen a priori, may not expressive enough represent the trajectories fully. For fair comparison use the neural ansatz in Eq. (10) and with the same hyperparameters Table 6, instead of the ones used in the original paper (Bruna et al., 2022).
2. (()) is often ill- conditniiio,,,, so we add a small shift M (0) = ((0) +cI, for 1e-4. Even with with regularization the resulting ODE is still stiff (in our experiments the condition number is around 10*). Thus, to avoid this numerical issue, we perform the computation in double precision.
Finally, the initial condition needs to be optimized inference time, which takes considerable amount of time. In this case we use as an initial condition the condition given by our fully trained encoder
Matab (0) = F()). Aa. (un).
(28)
The method is implemented in JAX (Bradbury et al., 2018). We use the adaptive-step Dormand-Prinee integrator (Dormand & Prince, 1980) implemented in scipy. inteegrate volveivvp, and we allow the routine choose the adaptive time step. This choice allows to identify when the ODE system becomes locally stiff.
As mentioned in the main text, However, for the KdV equation the ODE integration goes unstable in all of the trajectories. Even after playing with smaller tolerances, different regularization constant, and more accurate initial conditions, we were not able to avoid the stability issues.
able compute the solutions for the VB and KS equations.
A.4.2 NEURAL DECODER VARIANT
In our main model, the ansatz network plays the role of a decoder, which is responsible for trans- forming from the latent to the ambient space. The biggest difference from typical decoder is that the ansatz essentially pins down the latent space and as a result the decoding map is not paramet- ric. In the neural decoder variant (NDV), consider the typical decoder setup and implement
19

---

Published as a conference paper at ICLR 2023
Table 8: Specifications for NDV training.
Stage
Batch Size
Training Steps
Seq Length
Base LR
LR decay
Steps Per Decay
Pretrain
8
4M
2
3 x 10-4
0.95
40K
Fine-tune
8
1M
20
5 x 10-5
0.95
10K
convolution-base decoder, whose trainable weights parametrizes the decoding map. This allows us to contrast and study the impact of having a fixed-form ansatz that is tailored to the problem.
For NDV, we use exactly the same encoder (architecture and latent dimensions) as our main model. The decoder architecture is essentially the reverse of the encoder. It consists of de-projection layer (reversing the last encoder layer), followed by a sequence upsampling blocks (alternating regular and strided convolution layers), each of which increases the resolution and decreases the number of features by a factor 2.
The latent dynamical model is MLP with (2048, 2048) features in its hidden layers. This corresponds to around 4.5M parameters, which smaller compared to our hyper U-Net model. We did experiment with larger models but did not observe any performance gain.
We follow the same multi-stage protocol to train NDV. The specifications are summarized in Table 8.
We explored the use of consistency regularization the NDV well and did not observe any clear advantage in terms of reconstruction accuracy or smoothness gain. A possible explanation is that the latent space under NDV has the freedom to scale down and let the decoder change accordingly to maintain the same level of reconstruction performance. This results in lower values in the consistency regularization without necessarily improving the conditioning of latent trajectories.
For the VB system, we explored a myriad of hyperparameter combinations including layer depths/widths, activation functions and learning rates, none of which resulted in a model that can be considered generalizing well to all snapshots.
A.4.3 DEEPONET
The DeepONet is a general framework to approximate operators using neural networks. In nutshell, for an operator G between two function spaces such that
G(w).
(29)
DeepONet seeks to approximate this map point- i.c..
G(w) = v(y),
(30)
where G the DeepONet and y are arbitrary points in the domain u. Following the example shown above, DeepONets are usually split into two networks: the trunk network, whose inputs are samples of w, and the branch network, whose inputs are the evaluation points y.
In this case we use them to approximate the trajectory given by an initial condition up at any point (z, t) for t > and x the domain. Namely we consider
u(x,1) G(un)(z 1) -N.uun NNbzanch (zr,);
(31)
where Nerunk RRN -> Rm and Merannh I + R. are neural networks.
For both the trunk and branch networks we use fully-connected MLPs. We the model on the VB, KS and KdV datasets, each of which has N Ngamples 512. For the VB data, linear unit (ReLU) activation functions. For the KS and the KdV datasets, we found that tanh and Gaussian Error Linear Units (GELU) activation functions (Hendrycks & Gimpel, 2016) worked better, respectively.
rectified
For the branch network, modified the input slightly. Instead of feeding (æ, t) directly, we compute periodic features for the spatial coordinate x using sines and cosines at different frequencies. The main rationale for this is to impose periodicity in space in the network. Thus, the input to the branch
20

---

Published as a conference paper at ICLR 2023
Table 9: FNO hyperparameters.
Layer Width
Layer Width
System
P K
Num Layers
Num Modes
(SpectralConv)
(Final Projection)
VB
20
64
8
128
KS
32
20
256
4
16
128
KdV
32
20
256
16
128
network can be rewritten
sin (2n/ Lx)
cos/ 1vT T.x))
sin 4TT Txx)
cos( 4 /Tx)
(32)
sin (6tt Lx))
cos 6T/ Lxx
sin [8T/ Lx)
cos 8T/ Lx)
and the branch network is then given by
Naranca (3r, MLP(5,();
(33)
where MLP is a fully connected network. For the VB system, both the trunk and the branch network has hidden layers with 256 neurons each. The full network has 528, 896 parameters in this case. For KS and KdV data, we use layers with 512 neurons each. This resulted in total of 2, 369, 536 parameters.
The model was trained using the Adam optimizer (Kingma & Ba, 2015). After hyper parameter sweep, the learning rate that provided the best results was 10-4 with an exponential decay every 10000 steps by a factor of 0.977. We used time horizon of 300 steps to train the networks in each case. We used 5000 epochs, with a batch size of 512, in which each element of the batch is a combination of uy and (z, t).
A.4.4 FOURIER NEURAL OPERATOR
The Fourier Neural Operator (FNO) (Li et al., 2020b) is an operator learning framework centered around the operator map/layer
v' defined as:
/ k(x, y)u()) dy + Wu(r),
o (x)
(34)
where k(x,-) and W are parametrized by trainable parameters. The integral term (a.k.a. spectral convolution, or SpectralConv operation) is conveniently computed in the Fourier space, where low- pass filtering may be additionally applied (controlled by the number of modes). To make up the overall operator map, the input function is first raised to a higher dimension, and then goes through multiple FNO layers Eq. (34) interleaved with nonlinear activation functions, before finally projected to the desired output dimension using a local two-layer MLP.
In our benchmark experiments, follow the training approach outlined in the original paper, using FNO to map from a stack of P past snapshots {u--P,- ...., u"-1] } to u", with the grid locations x additional constant feature in the input. We train the model in autoregressive manner by iteratively rolling out K prediction steps, each step appending the current predicted snapshot to the input for predicting the next step (and removing the oldest keep the input shape constant). L2 loss is averaged over all K prediction steps and minimized with respect to the model parameters.
For each system, we performed sweeps over hyperparameters. The best-performing ones are listed in Table 9.
A.4.5 DYNAMIC MODE DECOMPOSITION
The Dynamic Mode Decomposition (DMD) method (Schmid, 2010) is data-driven, equation- -fee method of modeling the spatiotemporal dynamics of system. The key idea is obtain low-rank
21

---

Published as a conference paper at ICLR 2023
linear operator which, when applied to the current state of the system, approximates the state of the system at the next timestep.
Suppose are given the snapshots u", u' ... u " where u denotes the state timestep & (i.e. at time k(t) and N, is the total number time steps. We gather the imput and output matrices, U and U' by stacking the each snapshot columnwise in U and the corresponding next timestep in U Concretely,
U
[1
u²
n²
U |n2
113
(35)
DMD learns a linear operator Á of a specified rank K such that the following quantity is minimized.
A U
(36)
where || || denotes the Frobenius norm. Without the rank constraint, the quantity ||AU - U' minimized by = U'Ut, where Ut denotes the Moore-Penrose pseudoinverse of U. However, for large dimensionality and large number of snapshots, it is intractable compute the matrix A
explicitly.
The DMD algorithm efficiently computes the approximate low rank operator À A by exploiting the singular value decomposition (SVD) of U. See Schmid (2010) Kutz et al. (2016) for more details on the DMD algorithm.
Since our datasets consist of multiple trajectories, we stack together the U and U' matrices from each trajectory to form the matrices X and X', respectively. For Nera; trajectories consisting of N₃ time steps each, we compute the matrices X, X' E RdxM where M Neras (N₂ - 1) and d the dimensionality of the data. For evaluation, an initial condition u e Rd can rolled forward to predict the future time steps i > 0, by repeated application of the operator Á.
For the VB, KS and KdV datasets, the number of trajectories N. is 1024, 800 and 2048 re- spectively; and the number of time steps per trajectory N, is 300, 1200 and 800 respectively. The dimensionality of the data is d = 512 in each case. Following (Kutz et al., 2016), we choose the desired rank K of the approximate operator À by inspecting the rate of decay of the singular values of the snapshot matrix X. For the VB and KS datasets, we chose the rank to be 64. For the KdV dataset, we chose the rank to be 80. We used only 1024 of the 2048 KdV trajectories to avoid the DMD algorithm running out of memory.
A.5 ABLATION STUDIES
Consistency regularization In Fig. 3, the relative reconstruction error is given by Eq, (15), and the relative consistency error is defined in similar fashion
||0 Ao o 1 (0)
rel PeIRMSE (0)
(37)
|
through which we measure whether we can recover the same after decoding and encoding again. The maximum component-wise 1st order total variation norm is computed as
Q N
max
Al gga+1.g - 0
(38)
where g denotes index of evaluation cases and the maximum is taken over latent dimensions i. This quantity measures how fast/slow the latent variables changes over fixed time window.
In Fig- 5, show samples of latent trajectories obtained using various y values. We observe that trajectories under high y values are visibly smoother.
Fine-tuning by multi-step training In Fig. 6, we compare rollouts generated by the models obtained after the pretraining and fine-tuning stages respectively, to highlight the importance of the latter. We observe that for all systems fine-tuning is able to significantly reduce the point-wise errors and, in the case of KS and KdV, remove patches in the predictions that look rather 'unphysical".
22

---

Published as a conference paper at ICLR 2023
0
10
10
10!
25
-50
10
20
30
40
50
10
20
50
10
20
30
40
50
time
time
time
time
Figure 5: Latent trajectory samples resulted from training with different y values for KS system. Trajectories plotted are for the same 10 latent/weight dimensions (identified by the same colors) randomly selected for a fixed ansatz.
ReIRMSE
Sample Rollouts
10
PaIRMSE
VB
time
10a
MIRMSE
20-1
KS
10-3
20
40
60
80
100
10a
rAIRMSE
10-1
Kdv
10-4
10a4
10
20
time
Pretrained:
Fine-tuned:
Figure 6: Relative RMSE (left) and sample rollouts (right) demonstrating the effects of fine-tuning with multi-step training. Rollouts are plotted in (x, () plane with three rows showing the ground truths, pretrained and fine-tuned predictions respectively.
End-to-end (e2e) training our multi-stage training protocol, we freeze our encoders when training the dynamical model. As ablation study, also attempted adding additional multi- step, e2e training stage where the parameters for both the encoder and dynamical are allowed to change, optimizing for single rollout loss in the ambient space. Despite observing steady decrease in
23

---

Published as a conference paper at ICLR 2023
10.
10a
10
reRMSS
10-4
10P
10
10-2
10
20
100
time
time
Figure 7: Relative RMSE for latent space fine-tuned (-) anddd e2e-trained (-) for KS (left) and KdV (right) systems.
Table 10: Mean reIRMSE of the decoder to approximate snapshots of the KdV training trajectories Eq. (24). We consider MLP with different number of features and activation functions. The number of layers for the first and last layer is chosen that the total number of parameters is similar to using three frequencies.
Features/Act. fun.
relu
sin
sin(
swish
tanh
66 4, 4, 6)
0.6751
0.5585
0.0583
0.6395
0.6650
(6, 4, 4, 4,
0.6496
0.4208
0.0432
0.5343
0.5945
46,4,4,4,4,00
0.6420
0.3234
0.0837
0.4534
0.5362
(6,8,800
0.4452
0.4291
0.0137
0.5653
0.5491
(6,8,8,8,6)
0.3336
0.2325
0.0098
0.3953
0.4311
(6, 16,16,6,66
0.3260
0.3330
0.0092
0.4871
0.4211
16, 16,6)
0.1873
0.1077
0.0092
0.2306
0.2268
(6,32,
0.2297
0.2634
0.0083
0.3768
0.2634
(6,32,32,3,00
0.0878
0.0494
0.0105
0.0873
0.0643
(6, 64, 64,6)
0.1229
0.2009
0.0091
0.2568
0.1086
64,4,,6,,6,,, 64
0.0424
0.0274
0.0117
0.0361
0.0214
the training loss, this process (perhaps counter-intuitiveyy was in reality found hurt the quality of rollouts significantly (errors shown in Fig. 7), especially in its stability. We suspect the reason is that without enforcing the latent evolution to remain close highly regularized the encoder output, is easy for the optimizer to move to a nearby local minima with equal or potentially better performance the ambient loss. It is highly likely that the resulting dynamics, on the other hand, not amenable to long-range rollouts in the absence of proper regularization.
Decoder ansatz We benchmarked the approximation power of classical MIPs of different depth, width, and different activation functions. The benchmark is similar to the in Table 3, in which minimized Eq. (24) for trajectories of the KdV equation. For each experiment we solved the problem in Eq. (24) using SGD with learning rate of 10 batch size of 4 with 0.5.M training steps. For each configuration, this experiment was repeated for 1000 snapshots sampled from different trajectories of the KdV equation, and we computed the mean relative RMSE at the end of the optimization. The mean relative RMSE for each configuration is depicted in Table 10. in which we can observe that the errors are much higher than the ones reported in Table 3. In addition, Table 10 shows that imposing the periodicity condition into the activation function greatly helps to reduce the approximation error.
A.6 ADDITIONAL METRICS
We provide some additional metrics including more spatial energy spectra at various prediction ranges (Fig. 8) and sample rollouts demonstrating richer range of patterns which are well captured by our models (Fig. 9).
24

---

Published as a conference paper at ICLR 2023
VB
30-
30-
-10
-20
-20
-20
-20
-20
30-1
10-1
107
10-1
11
16-1
30
10'
10-1
137
KS
20
30
10
50
50
-20
-20
-20
KdV
30
30
30
30
-20
-20
-20-
-30
-30
-30
Ours-NFA: NG-Lo:
NG-Hi: Ours-NDV:
DeepONet: FNO: DMD:
Figure 8: Spatial energy log ratio, computed using Eq. (16), for different prediction ranges.
VB
KS
KdV
Figure 9: Additional evaluation rollouts plotted in (z, () plane. Top rows are ground truths and bottom rows the corresponding predictions rolled from the same initial conditions. Prediction range: VB 3, KS 80, KdV 30.
25


Reviews:
[{"review_id": "tplcV9qnqj", "reviewer": "Reviewer_WMEw", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper introduces a latent representation of dynamic systems by proposing an autoencoder framework specialised to PDEs. The authors motivate and introduce several components for the loss functions such as reconstruction, consistency as well a single step pre-training. Good performance is demonstrated on several examples. ", "summary_of_the_review": "The paper has several positive aspects and seems to be thought out well. Empirical results are promising, so I am tending towards acceptance.  ", "strength_and_weaknesses": "__Strengths__:\n- Interesting idea on how to make use of latent variables for dynamic data.  \n- The idea is well thought out in several aspects, e.g. different loss functions, simpler pretraining, the smooth dynamics etc.\n- The performance seems good across tasks\n- Comprehensive list of relevant work\n\n__Weaknesses__:\n\n- On a high level, the idea of producing latent trajectories to encode temporal data has been present for some time and was also discussed in the original neural ODE paper, which doesn't seem to be mentioned here. But it is also only a very high level connection. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "eAoLoc3YqD5", "reviewer": "Reviewer_haAi", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper proposes a method to model dynamical systems in a latent space in a smooth way. To go back and forth from the latent space, a decoder (called \"ansatz\") and an encoder network are used, that are task dependent. One first trick proposed in the paper is to propose some cyclic consistency for this encoder decoder, so that both latent->state->latent and state->latent->state should be identity functions. This is advocated to help for the smoothness of the latent variables.\nThen, the core contribution of the paper in my view lies in the fact that the latent variables are not seen as an input to the decoder, but rather as encoding the *weights* of a network (in an hypernet way),  that inputs the coordinates (time, space) and outputs the value (nerf-style). This is advocated as yielding a much more powerful representation.\nThe rest of the paper is busy with describing important learning tricks and describing experiments.", "summary_of_the_review": "I liked the paper and the approach looks promising to me, puting together some cyclic consistency idea for latent space modeling as well as nonparametric hypernet decoders.", "strength_and_weaknesses": "As a set of strength, I can say the paper is first useful for the not so knowledgeable reader to get references in the field. Then, the proposed method looks original to me, with its two nice contributions (cyclic losses, hypernet decoder). The method seems to be applicable even when the number of trajectories is quite limited (1000 here apparently). It was a nice read and I recommend it for publication.\n\nAs weaknesses, I could mention:\n* The experiments look quite weak to me. I understand that great care was taken on this, but I basically wonder about real world usecases, which I cannot see here.\n* I wonder about how critical is the choice of the decoder. The fact is that the chosen Fourier-based decoder looks pretty appropriate for the datasets at hand (figure 4), but it may be that the system would collapse for other kinds of dynamic data. Would there be some guidelines regarding what kind of decoder are appropriate or some ablation studies on that ?\n* Likewise, even assuming a Fourier-based decoder, the choice of the frequencies look mysterious and I wonder whether their choice needs a lot of care. Why are they not trained ?\n* Table 1 is not satisfying at all. I would have preferred some objective assessment of the encoder/decoder experiments, instead of this \"working/not working\" thing. What does \"working\" mean exactly anyways ?\n* I don't understand the \"spatial energy spectra\" as a metric. Could you please explain it differently ? ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "NqSoyYm-TE", "reviewer": "Reviewer_mB8q", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper proposes latent PDE learning with novel autoencoder-type losses.\n", "summary_of_the_review": "\nThis is an excellent paper with refreshing, new ideas on this domain, good enough results and excellent writing.\n", "strength_and_weaknesses": "\ns: this paper has excellent and effortless writing that demonstrates both clarity and expertise. I enjoyed reading this manuscript a lot! I also like how the paper makes does a transparent discussion of the training difficulties/annoyances\n\ns: the idea is sensible, and turning the latent state into parameters of the decoder is elegant, which conceptually marries latent dynamics with operator learning. The construction of eqs 6-9 is very clever.\n\ns: the results are good, and sufficiently elucidated\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: notable-top-25%