Learning to Accelerate Partial Differential Equations via Latent Global Evolution
Authors: Tailin Wu
Takashi Maruyama
Jure Leskovec
Venue: NeurIPS 2022 Accept (2022)
Keywords: accelerate
partial differential equation
latent global evolution
inverse optimization

Abstract: Simulating the time evolution of Partial Differential Equations (PDEs) of large-scale systems is crucial in many scientific and engineering domains such as fluid dynamics, weather forecasting and their inverse optimization problems. However, both classical solvers and recent deep learning-based surrogate models are typically extremely computationally intensive, because of their local evolution: they need to update the state of each discretized cell at each time step during inference. Here we develop Latent Evolution of PDEs (LE-PDE), a simple, fast and scalable method to accelerate the simulation and inverse optimization of PDEs. LE-PDE learns a compact, global representation of the system and efficiently evolves it fully in the latent space with learned latent evolution models. LE-PDE achieves speedup by having a much smaller latent dimension to update during long rollout as compared to updating in the input space. We introduce new learning objectives to effectively learn such latent dynamics to ensure long-term stability. We further introduce techniques for speeding-up inverse optimization of boundary conditions for PDEs via backpropagation through time in latent space, and an annealing technique to address the non-differentiability and sparse interaction of boundary conditions. We test our method in a 1D benchmark of nonlinear PDEs, 2D  Navier-Stokes flows into turbulent phase and an inverse optimization of boundary conditions in 2D Navier-Stokes flow. Compared to state-of-the-art deep learning-based surrogate models and other strong baselines, we demonstrate up to 128x reduction in the dimensions to update, and up to 15x improvement in speed, while achieving competitive accuracy.

Learning to Accelerate Partial Differential Equations via Latent Global Evolution
Tailin Wu Department of Computer Science Stanford University tailin@cs. stanford... edu
Takashi Maruyama
NEC Corp. & Stanford University
49takashi@ner eneaasshieieee... com takashi279@cs. .ateess.ttee stanford... edu
Jure Leskovee
Department of Computer Science
Stanford University jure@cs. stanford edu
Abstract
Simulating the time evolution of Partial Differential Equations (PDEs) of large- scale systems crucial many scientific and engineering domains such as fluid dynamics, weather forecasting and their inverse optimization problems. However, both classical solvers and recent deep learning-based surrogate models are typically extremely computationally intensive, because of their local evolution: they need to update the state of each discretized cell at each time step during inference. Here we develop Latent Evolution of PDEs (LE-PDE), simple, fast and scalable method to accelerate the simulation and inverse optimization compact, global representation of the system and efficiently evolves fully in the latent space with learned evolution models. LE-PDE achieves speed-up by having much smaller latent dimension update during long rollout as compared to updating in the input space. We introduce new learning objectives to effectively learn such latent dynamics to ensure long-term stability. We further introduce techniques for speeding up inverse optimization of boundary conditions for PDEs via backpropagation through time in latent space, and an annealing technique address the non- diniifferreniiabiiiyy and sparse interaction of boundary conditions. We test our method ID benchmark of nonlinear PDEs, 2D Navier-Stokes flows into turbulent phase and an inverse optimization of boundary conditions in 2D Navier-Stokes flow. Compared other strong baselines, we demonstrate up 128x x reduction in the dimensions to update, and up to 15x improvement in speed, while achieving competitive accuracy.
PDEs. LE-PDE learns
Introduction
Many problems across science and engineering are described by Partial Differential Equations (PDEs). Among them, temporal PDEs are of huge importance. They describe how the state of (complex) system evolves with time, and numerically evolving such equations are used for forward prediction and inverse optimization across many disciplines. Example application includes weather forecasting
M. jet engine design [2), nuclear fusion [3], laser-plasma interaction astronomical simulation and molecular modeling (
To numerically evolve such PDEs, decades of works have yielded (classical) PDE solvers that are tailored to each specific problem domain 7. Albeit principled and accurate, classical PDE solvers
Project website and code can be found at http: htttp://snap..e wwwp:sss..aaaaadeoeomoomooo.eeee edu/le_pde/
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

static
p
static encoder
dg
forward mode
parameter
inverse oplimization mode
dzp
Zp
p
Z
p
T
dynamic encoder
9
z²
g
z
g
z
g
Z
decoder
Loss
latent space
ovolutionary model
input state
Latent Evolution of PDEs (LE-PDE)
output state
Figure 1: LE-PDE schematic. In forward mode (green), LE-PDE evolves the dynamics in global latent space. In inverse optimization mode (red), optimizes parameter p (e.g. boundary) through latent unrolling. The compressed latent vector and dynamics can significantly speed up both modes.
are typically slow due to the small time steps or implicit method required for numerical stability, and their time complexity typically scales linearly or super-linearly with the number of cells the domain is discretized into [8. For practical problems in science and engineering, the number of cells at each time step can easily go up to millions or billions and may even require massively parallel supercomputing resources p [0 Besides forward modeling, inverse problems, such as inverse optimization of system parameters and inverse parameter inference, also share similar scaling challenge [ How to effectively speed up the simulation while maintaining reasonable accuracy remains an important open problem.
Recently, deep learning- -based surrogate models have emerged as attractive alternative to complement [12 or replace classical solvers [13 [4] They directly learn the dynamics from data and alleviate much engineering effort. They typically offer speed-up due to explicit forward mapping [5] [6].
larger time intervals [14], or modeling on coarser grid 12] [7]. However, their evolution scales with the discretization, since they typically need to update the state of each discretized cell each time step, due to the local nature of PDEs [8]. For example, if problem is discretized into I million surrogate models (e.g.. CNNs, Graph Networks, Neural Operators) will
cells, deep learning-basse need to evolve these I million cells per time step. How to go beyond updating each individual cells and further speed up such models remains a challenge.
Here we present Latent Evolution of PDEs (LE-PDE) (Fig. a simple, fast and scalable method to accelerate the simulation and inverse optimization of PDES. Our key insight is that common feature of the dynamics of many systems of interest is the presence of dominant, low-dimensional coherent structures, suggesting the possibility of efficiently evolving the system in a fow-dimensional global latent space. Based on this observation, we develop LE-PDE, which learns the evolution of dynamics in a global latent space. Here by "global" we mean that the dimension the latent state is fixed, instead of scaling linearly with the number of cells as in local models. LE-PDE consists of a dynamic encoder that compresses the input state into a dynamic latent vector, a static encoder that encodes boundary conditions and equation parameters into static latent vector, and latent evolution model that evolves the dynamic latent vector fully in the latent space, and decode via decoder only needed. Although the idea of latent evolution has appeared in other domains, such as in computer vision 19 20 21 and robotics 22 23 24 25 these domains typically have clear object structure in visual inputs allowing compact representation. PDEs, the other hand, model dynamics of continuum (e.g., fluids, materials) with infinite dimensions, without a clear object structure, and sometimes with chaotic turbulent dynamics, and it is pivotal to model their long- term evolution accurately. Thus, learning the latent dynamics of PDEs presents unique challenges.
We introduce a multi-step latent consistency objective, to encourage learning more stable long-term cvolution in latent space. Together with the multi-step loss in the input space, they encourage more accurate long-ten long-term prediction. To accelerate inverse optimization of PDEs which is pivotal in engi- neering (e.g. optimize the boundary condition so that the evolution of the dynamics optimizes certain predefined objective), we show that LE-PDE can allow faster optimization, via backpropagation through time latent space instead of in input space. To address the challenge that the boundary condition may be non-differentiable or too sparse receive any gradient, design an annealing technique for the boundary mask during inverse optimization.
We demonstrate our LE-PDE in standard PDE-learning benchmarks of a 1D family of nonlinear PDEs and 2D Aavier-Stokes flow into turbulent phase, and design an inverse optimization problem in 2D Navier-Stole flow to probe its capability. Compared with ofte-of-the-art deep learning- based

---

surrogate models and other strong baselines, we show up to 128x reduction in the dimensions to update and up 15x speed-up compared to modeling input space, and competitive accuracy.
Problem Setting and Related Work
We consider temporal Partial Differential Equations (PDEs) w.r.t. time t [0. T and multiple spatial dimensions x= [21, 22, ...II
Ip] E X _ R.. We follow similar notation as in 77 Specifically,
p]
0.u F a, u. axu, axxu,
(t,x) E [0, T] X X, E x, (1,x) E [0, T) OX.
(1)
u(0, x:) = (x), Blui((, 0,
(2)
Here u [0, [] X + R" is the solution, which is infiniteddimensional function. a is time- independent static parameters of the system, which can be defined on each location x, egg. diffusion coefficient that varies in space but static in time, global parameter. F is linear nonlinear function on the arguments of (x, a, u, 0xu, Oxx Note that in this work we consider time- independent PDEs where F does not explicitly depend on t. u ux(x) is the initial condition, and B[u] x) = 0 is the boundary condition when x is on the boundary of the domain 0X across all time / E [0, 7J. Here 0хu= are a matrix and a 3-order tensor, respectively (since x is a vector). Solving such temporal PDEs means computing the state u(t, x) for any time E [0, and location x E X given the above initial and boundary conditions.
du 0xxu Su are first- and second-order partial derivatives, which dx?
8x
Classical solvers for solving PDEs. To numerically solve the above PDEs, classical numerical solvers typically discretize the domain X into a finite grid or mesh X = {c,) with N non- overlapping cells. Then the diinite-dimmesiionll solution function of uff,x) is discretized into {u, ) E U for each cell i and time i = 42 k = 1, ..... a is similarly discretized into [a 1 with values in each cell. Mainstream numerical methods, including Finite Difference Method (FDM) and Finite Volume Method (FVM), proceed to evolve such temporal PDEs by solving the equation at state [uę+*} at time fx+1 from state {u²} time f4. These solvers are typically slow due to small time/space intervals required for numerical stability, and needing to update each cell at each time steps. For more detailed information on classical solvers, see Appendix [
U
Deep Barring--aseed surrogate modeling. There are two main approaches in deep learning- based surrogate modeling. The first class of method is autoregressive methods, which learns the mapping fo with parameter of the discretized states U² between consecutive time steps t, and Here i = ( is the model fo 's predicted state
ft-1: (+ = fo( p), k = 0, for U (ut J==1 at time tg. with (/= == U p (0x, {ar) the system parameter which includes the boundary domain 0% and discretized static parameters fa; inference time results in autoregressive rollout
[u
1 Repetitively apply fo at
(3) Here fol:p): : U Us a partial function whose second argument is fulfilled by the static system parameter p. Typically Sa is modeled using CNNs (if the domain X is discretized into a grid), Graph Neural Networks (GNNs, if the domain X is discretized into a mesh). These methods all involve local computation, where the value u {ur (uf) JEN(()))) at time ls, where A(()) is the set of neighbors up certain hops. Such formulation includes CNN-based models 226.. GNN-based models 7 27 28 and their hierarchical counterparts
(+ = (fo(:))) (i*)) m 1.
cell i at time fx+1 depend on its neighbors
The surrogate modeling with local dynamics makes sense, since the underlying PDE is essentially a local equation that stipulates how the solution function u's value at location x depends on the values at its infinitesimal neighborhood. The second class of method Neural Operators 332 33 |3]] (35] 36] which learns a neural network (NN) that approximates mapping functions. Although having the advantage that the learned mapping is discretization invariant, given a specific discretization, Neural Operators still needs to update the state at each cell based on neighboring cells (and potentially cells far away), which is still inefficient at inference time, especially dealing with larger-scale problems. In contrast to the above classes of deep learning-based approaches that both requires local evolution at inference time, our LE-PDE method focuses improving efficiency. Using learned global latent space, LE-PDE removes the need for local evolution and can directly evolve the system dynamics via global latent vectors z* E R&ds for time t.. This offers great potential for speed-up due to the significant reduction in representation.
18 29
14 30 3
between andinite-dimension
3

---

Inverse optimization. Inverse optimization is the problem of optimizing the parameters p of the PDE, including boundary 0% or static parameter of the equation, so that a predefined objective Lala, 0x] = S lel (e,u(t, x]]dt is minimized. Here the state u(t, x) implicitly depends on a, 0x through the PDE (Eq. D and the boundary condition (Eq. 2) Such problems have huge importance in engineering, e.g. in designing jet engines 2 and materials
[ where the objective can be minimizing drag or maximizing durability, and inverse parameter inference (i.e. history matching) be maximum posteriori estimation. To solve such problem, classical methods include adjoint method 41 12], shooting method [ collocation method [4),, etc. One recent work 4 explores optimization via backpropagation through differential physics in the input space, demonstrating speed-up and improved accuracy compared to classical CEM method 46. However, for long rollout and large input size, the computation becomes intensive to the point of needing to save gradients in files. In comparison, LE-PDE allows backpropagation in latent space, and due to the much smaller latent dimension and evolution model, it can significantly reduce the time complexity in inverse optimization.
38
39
40
where the objective
Reduced-order modeling. A related class of work is reduced-order modeling. Past efforts typically use linear projection into certain basis functions 47
50
52
53
541 which may not have
enough representation power. A few recent works explore NN-based encoding 55
56
57
58
59 60
for fluid modeling. Compared to the above works, we focus on speeding up simulation and inverse optimization of more general PDEs using expressive NNs, with novel objectives, and demonstrate competitive performance compared to state-of-the-ar deep learning-basse models for PDEs.
3 Our approach LE-PDE
In this section, we detail our Latent Evolution of Partial Differential Equations (LE-PDE) method. We first introduce the model architecture (Sec. 3.1 and then we introduce learning objective to effectively learn faithfully long-term evolution (Sec. 3.2p. In Sec. 3.3, we introduce efficient inverse optimization in latent space endowed by our method.
3.1 Model architecture
The model architecture of LE-PDE consists of four components: (1) a dynamic encoder q U - Rd that maps the input state U {uf N E U latent vector z" R=; (2) an (optional) static encoder r P + Edsr that maps the (optional) system parameter p P to a static latent embedding Zpi (3) a decoder h Rdd + U that maps the latent vector zt E gd, back to the input state UK; (4) a latent evolution model g Rdd x Rdsa + IRd that maps z² E RRds at time l4 and static latent embedding 2p RRd ++1 E dR at time (++1- We employ the temporal bundling trick 7 where each input state U can include states over fixed length S of consecutive time steps, in which case each latent vector zą will encode such bundle of states, and each latent evolution will predict the latent vector for the next bundle of S steps. S is hyperparameter and may be chosen depending on the problem, and S = 1 reduces to no bundling. A schematic of the model architecture and its inference illustrated in Fig. [ Importantly, we require that for the dynamic encoder q, it needs to have flatten operation and MultiLayer Perception (MLP) head that maps the feature map into a single fixed-length vector z E Rdd In this way, the dimension of the latent space does not scale linearly with the dimension of the input, which has the potential to significantly compress the input, and can make the long-term prediction much more efficient. At inference time, L.E-PDE performs autoregressive rollout in latent space Rdd.:
f = h o g ((p)) q((*) = (1 ((rr(())) o... o g(. r(p)) (1(04))
(4)
composing mn times
Compared autoregressive rollout in input space (Eq. B LE-PDE can significantly improve efficiency with much smaller dimension of zt E ddd compared to 1 E U. Here we do not limit the architecture for encoder, decoder and latent evolution models. Depending on the input U* the encoder g and decoder h can be CNN or GNN with a (required) MLP head. In this work, we focus on input that discretized as grid, so the encoder and decoder are both CNN+MLP, and leave other architecture (e.g. GNN+MLP) for future work. For static encoder r. it can be a simple MIP if the system parameter p is a vector (e.g. equation parameters) or CNN+MLP p is 2D or 3D tensor (e.g. boundary mask, spatially varying diffusion coefficient). We model the latent evolution model g

---

as an MLP with residual connection from input to output. The architectures used in our experiments, together with discussion of its current limitations.
are detailed Appendix
C
3.2 Learning objective
Learning surrogate models that can faithfully roll out long-term is an important challenge. Given discretized inputs {U*},, k I...K M, introduce the following objective address it:
K
1
L
stuti-step + Liccons + consistency)
(5)
K
km1
>Em= em((Om _ L (((h(((U))) U²)
L Limultisste step
where
ressons
(6)
L Leonnistency
-m (o(-r(p))) og(l>*) 2
Here é is the loss function for individual predictions, which can typically be MSE or L2 loss. Ü++m given in Eq. $ Lieeom aims reduce reconstruction loss. L stpp performs latent multi-step
evolution given in Eq. 4 and compare with the target U in input space, up to time horizon M. Om are weights for each time step, which we find that (2,,02, 02; ....rs)) = (1, 01,
works well. Besides encouraging better prediction in input space via =uuti-sspp we also want stable long-term rollout in latent space. This is because in inference time, we want to mainly perform autoregressive rollout in latent space, and decode to input space only when needed. Thus, we introduce novel latent consistency loss Consiistecy which compares the m-step latent rollout g (s (()))) (m)) (g(U*) with the latent target ((U++m) in latent space. The denominator [[ serves as normalization to prevent the trivial solution that the latent space collapses to single point. Taken together, the three terms encourage
accurate and consistent long-tet term evolution both in latent and input space. In Sec. 4.4 we will investigate the influence of L comsistency and L stlti-step
3.3 Accelerating inverse optimization
In addition improved efficiency for forward simulation, LE-PDE also allows more efficient inverse optimization, via backpropagation through time (BPTT) in latent space. Given a specified objective Lalp) = Ex_ RR {(T/)) which is a discretized version of T.a [a, 0X) in Sec. we define the objective:
ke
La(p] = (((rT(p))
where tmm (Tm(p) is given by Eq. setting 8: = 0 using our learned LE-PDE, which starts initial state of U encode it and p into latent space, evolves the dynamics in latent space and decode to (Tm as needed. The static latent embedding zp r(p) influences the latent evolution at each time step via g(,, r())). An approximately optimal parameter p can then be found by computing gradients aLalpl
Op using optimizers such as Adam 66 (The gradient flow is visualized as the red arrows in Fig.
D. When is a boundary parameter, e.g. location of the boundary segments or obstacles, there is challenge. Specifically, for CNN encoder q. the boundary information typically provided as binary mask indicating which cells are outside the simulation domain Ω. The discreteness of the mask prevents the backpropagation of the model. Moreover, the boundary cells may interact sparsely with the bulk, which can lead to vanishing gradient during inverse optimization. To address this, we introduce a function that maps p to a soft boundary mask with temperature, and during inverse optimization, anneal the temperature from high to low. This allows the gradient to pass through mask to p, and stronger gradient signal. For more information, see Appendix
4 Experiments
In the experiments, we aim to answer the following questions: (1) Does LE-PDE able to learn accurately the long-term evolution of challenging systems, and compare competitively with state-of- the-art methods? (2) How much can LE-PDE reduce representation dimension and improving speed, (3) Can LE-PDE improve and speed up inverse optimization? For
especially with larger systems? the first and second question, since in general there is a fundamental tradeoff between compression
5

---

surrogate models in Table 1: Performance of models in ID for scenarios E1,E2, E3, Accumulated error 1 E.E,
2, MSE.
Representation dimension (= S X fg here) is the number of dimensions to update at each time step. The bold values represent the best performance for experiments and underline shows second best.
Accumulated Error 1
Runtime [ms]]
Representation dim
(R))
WENOS FNO-RNN
PNO-PF
MP-PDE
LE-PDE (0ur5)
WENOS
MP-PDE
Full (ours)
LE-PDE
LE-PDE
MP-PDE
LE-PDE
(ours)
(250, 100)
2,022
11.95
0.54
1.13
10"
20
2500
128
E1
(250,50)
6.23
29.98
0,51
1.67
1.20
18 102
80
20
8
1250
128
(250,
9.63
10.44
0.57
1.47
1.17
1.7 102
20
8
1000
128
K2
(250, 100)
1.19
17.09
2.53
1.58
0.77
1,9 x 10"
20
8
2500
128
E2
(250,50) 50;
5.35
3.57
2.27
1.63
1.13
18 x 102
90
20
8
1250
128
(250,
8.05
3.26
2.38
1445
1,03
1.7 102
20
1000
128
(250, 100)
4.71
10.16
5.69
4.26
3.39
4.8 10"
19
6
2500
64
E3
(250,50)
11.71
14.49
5.39
3.74
3.82
4.5 x 10"
90
19
6
1250
64
(250,40) 15.94
20,90
3,70
3.78
10
20
1.28
(reduction of dimensions to represent a state) and accuracy 62 63 i.e. the larger the compression to improve speed, the more lossy the representation is, will need to sacrifice certain amount of accuracy. Therefore, the goal of LE-PDE is to maintain a reasonable or competitive accuracy (maybe slightly underperform stat-of-thhearde while achieving significant compression and speed up. Thus, to answer these two questions, we test LE-PDE in standard benchmarks of 1D family of nonlinear PDEs to test its generalization to new system parameters (Sec. 4.1. and a 2D Navier-Stoker flow up to turbulent phase (Sec. 4.2). The PDEs in the above scenarios have wide and important application in science and enginecring. In each domain, we compare LE-PDE's s long-ten evolution performance. speed and representation dimension with state-of-the art deep
the domain. Then we answer question (3) in Section 4.3 Finally, in Section 4.4 we investigate the impact of different components of LE-PDE and important hyperparameters.
4.1
1D family of nonlinear PDEs
Data and Experiments. In this section, we test LE-PDE's ability to generalize to unseen equations with different parameters in given family. We use the ID benchmark in 77. whose PDEs are
0.u + Ox(22² -30,u + 7 (1, = 8(1,z)
(8)
u(0, x) 8(0, 5(t, x) Σ A,sin(wyt + 2 6)) j=1
(9)
Here the parameter p (a. 3) The term 5 is forcing term
64 with J = 5, 16 and
coefficients Aj and Wj sampled uniformly from A; ~ 0.5 Wi U1-0.400.4) E, (1,2,3), Ø3 U[, 2T). Space is uniformly discretized nz 200 in [0, 16) and time is uniformly discretized to n4y = 250 points in [0. 4]. Space and time are further downsampled to resolutions of (z), {2500, 100), (250, 50), (250, 40)3. The (a.²² advection term makes the PDE nonlinear. There are 3 scenarios with increasing difficulty: E1: Burgers' equation without diffusion
(1, 0, 0); E2: Burgers' equation with variable diffusion p (I, 77, 0) where 77 E [0, 0.2]: E3: mixed scenario with p (a, 3, 7) where a E 30, E [0, 0.4] and E [0, 1]. E1 tests the model's ability to generalize to new conditions with same equation. E2 and E3 test the model's ability to generalize to novel parameters of PDE with the same family. We compare LE-PDE with sfate-of-the-arr deep dearning-basee surrogate models for this dataset, specifically MP-PDE 7 GNN- model) and Fourier Neural Operators (FNO) [4]. For FNO, we compare with two versions: FNO-RNN is the autoregressive version in Section 5.3 of their paper, trained with autoregressive rollout; FNO-PF is FNO improved with the temporal bundling and push-forward trick as implemented in 77 To ensure a fair comparison, our LE-PDE use temporal bundling of S = 25 time steps as in MP-PDE and FNO-PF, We perform hyperparameter search over latent dimension of {64, 128} and use the model with best validation performance. In addition, we compare with downsampled ground-truth (WENOS), which uses classical 5th -orde WENO scheme 665 and explicit Runge- Kutge-Kutta> 4 solver 666 67 to generate the ground-truth data and downsampled to the specified resolution. For all models, we autoregressively roll out predict the states starting at step 50 until step 250, and record the accumulated MSE, runtime and representation dimension (the dimension of state to update at each time step). Details of the experiments are given in Appendix D
Results. The result shown Table 0 We see that since L.E-PDE uses 7.8 to 39-fold smaller representation dimension, it achieves significant smaller runtime compared to the MP-PDE model

---

Initial varticity
t=5
t=10 t=15 t=20
LE-PDE Prediction:
Figure Visualization of rollout for 2D Navier-Stokes PDE (Re 10*), for ground-trute (upper) and LE-PDE (lower, trained with 10-4, N 10²). LE-PDE captures detailed dynamics faithfully.
(which is much faster than the classical WENO5 scheme). Here we record the latent evolution time (LE-PDE evo) which the total time for 200-step latent evolution, and the full time (LEE-DDE full), which also includes decoding to the input space at each time step. The time for "LE-PDE evo" is relevant when the downstream application is only concerned with state at long-term future (e.g. 4 the time for "LE-PDE full" is relevant when every intermediate prediction is also important. LE-PDE achieves up 15x speed-up with "LE-PDE evo" and speed-up with "LE-PDE full""
With above 7.8x compression and above 4x speed-up. LE-PDE still achieves competitive accuracy, For E1 scenario, significantly outperforms both original versions of FNO-RNN and MP-PDE, and only worse than the improved version of FNO-PF. For E3, LE-PH LE-PDE outperforms both versions of FNO-RNN and FNO-PF, and the performance is on par with MP-PDE and sometimes better. For E2, LE-PDE outperforms all state-of-the-ant models by large margin. Fig. + in Appendix D shows our model's representative rollout compared to ground-truth. We see that during long-rollout, our model captures the shock formation faithfully. This ID benchmark shows that LE-PDE is able achieve significant speed-up, generalize to novel PDE parameters and achieve competitive long-term rollout.
4.2 2D Navier-Stokes flow
Data and Experiments. We test LE-PDE in 2D benchmark (14] of Navier-Stoken equation. Navier-Stokes equation has wide application science and engincering, including weather forecasting. jet engine design, etc. It becomes more challenging to simulate when entering the turbulent phase, which shows multiscale dynamics and chaotic behavior. Specifically, we test our model in a viscous, incompressible fluid in vorticity form in unit torus:
0,w(1,x) u(), Vw((,x) vAw(t, f(x),
E (0, (0,
(10)
V u-t, 0,
(0, t [0,
(11)
w(0, x)) (2),
10,
(12)
Here w(t, 2) x u(t, x) is the vorticity, v E R4 is the viscosity coefficient. The domain discretized into 64 64 grid. We test with viscosities of v 10 1 10 10 The fluid is turbulent for 10 10 (Re 2 104). We compare of-te-t learning-bast based model Fourier Neural Operator (FNO) 14! for this problem, and strong baselines of Te-Net (26). U-Net (68 and ResNet [69] For FNO, the FNO-2D performs autoregressive rollout, and FNO-3D directly maps the past 10 steps into all future steps. To ensure fair comparison, here our LE-PDE uses past 10 steps to predict
Table 2: Performance of different models in 2D Navier-Stokes flow. Runtime is using the
10 3 N 1000 for predicting 40 steps in the future.
10
10
10
10-0
Representation dimensions
Runtime
Runtime
Method
T 50
T 30
T 3
T 20
full [ms]
evo [ms]]
N 1000
N 1000
N 10000
N 1000
FNO-3D
4096
0.0088
0.1918
0.0820
0.1893
FNO-2D
4096
0.0128
0.1559
0.0834
0.1556
U-Net 68
4096
813
813
0.0245
0.2051
0.1190
0.1982
TF-Net
26
4096
428
428
0.0225
0.2253
0.1168
0.2268
ResNet
69
4096
317
317
0.0701
0.2871
0.2311
0.2753
LE-PDE (ours)
256
48
15
0.0146
0.1936
0.1115
0.1862

---

T=0 T=60 T=70
T=0
T=60
T=70
C
(a) Trajectory generated by ground-trude truth solver with an initial randomly generated boundary parameter.
(b) Trajectory generated by ground solver with boundary parameter optimized by LE-PDE.
FMO-20
LE-PDE-latemi
GT solver
100
E.PDD--been
timels)
LE-PDE
LE-PDE
fraction
0.6
MSE
0.2
0.2
0.0
0.0
20
40
60
80
100
20
40
60
80
100
20
40
60
80
100
Iteration
Iterationn
(c) Runtime.
(d) Objective function.
(e) Smoke amount at lower outlet.
Figure 3: Numerical results associated with inverse optimization of boundary (inlet and outlet
designing) in Sec. 1.3 (a) shows a trajectory generated by ground-truth (GT) solver with an initial randomly generated boundary parameters (y-position of inlet and two outlets), with lower outlet passing 55.18% smoke; (b) with optimized boundary parameters, with lower outlet passing 31.79% of smoke, very near the objective percentage of 30%. (c) Runtime and (d) learning curve (Eq. 7 for inverse optimization at different iteration steps: (e) For LE-PDE, fraction smoke passing through the lower outlet computed by GT solver (green) and estimated by LE-PDE (orange). Error bar denotes 95% confidence interval over 50 runs with random initial conditions.
one future step and temporal bundling S = 1 (no bundling), the same setting as in FNO-2D. We use relative L.2 norm (normalized by ground-truth" L.2 norm) metric, same in 14]
Results. The results are shown in Table 2 Similar to ID case, LE-PDE is able to compress the representation dimension by 16-fold. Hence, compared with FNO-2D which is also autoregressive, LE-PDE achieves 9.3-fold speed-up with latent evolution and 2.9-fold speed-up with full decoding. Compared with FNO-3D that directly maps all input time steps all output times steps (which cannot generalize beyond the time range given), LE-PDE's runtime still 1.6x faster for latent evolution. For rollout L2 loss, LE-PDE significantly outperforms strong baselines of ResNet and U-Net, and TF-Net which is designed to model turbulent flow. Its performance is on par with FNO-3D with 1000 and the most difficult
= 10-4, FNO-2D in other scenarios. Fig. 2 shows the visualization of LE-PDE comparing with ground-truth, under the turbulent dynamics accurately. For more details, see Appendix E To explore how LE-PDE can model and accelerate the simulation of systems with a larger scale, in Appendix F
10-5 N 1000 and slightly underperforms
10
10000 scenario. We see that LE-PDE captures the detailed
we explore modeling a 3D Navier-Stokes flow with millions of cells per time step, and show more significant speed-up.
4.3 Accelerating inverse optimization of boundary conditions
Data and Experiments. In this subsection, set out to answer question (3), iee. Can LE- PDE improve and speed up inverse optimiza- tion? We are interested in long time frame see- narios where the pre-defined objective La Eq.
Table 3: Comparison of LE-PDE with baselines.
GT-solver Error
Runtime
(Model estimated Error)
[d]
LF-
0,305 (0. 123) 0.124 (0.004)
86.42
7 depends on the prediction after long-term rollout. Such problems are challenging and have implications in engineering, e.g. fluid control 170. [71. laser design for laser aaser-plasma interaction
ENO-2D
111.14
LE-PDE (ours)
0.035 (0.036)
49.81
and nuclear fusion [72]. evaluate, we build a Navier-Stokes flow in a family of boundary conditions using PhiFlow
[ as our ground-truth solver, shown in Fig. 3a
3b Specifically, we
create a cubical boundary with one inlet and two outlets on grid space of size 1282. We initialize
the velocity and smoke on this domain and advect the dynamics by performing rollout. The objective

---

of the inverse design here is to optimize the boundary parameter p. i.e. the y-locations of the inlet and outlets, so that the amount of smoke passing through the two outlets coincides with pre-specified proportions 0.3 and 0.7. This setting is challenging since slight change in boundary (up to a few cells) can have large influence in long-term rollout and the predefined objective.
As baseline methods, we use our LE-PDE's ablated version without latent evolution (essentially a CNN, which we call LE-PDE-Latent) and the FNO-2D [. both of which update the states in input space, while LE-PDE evolves in a 128-dimensional latent space (128x compression). To ensure a fair comparison, all models predict the next step using past step without temporal bundling, and trained with 4-step rollout. We train all models with 400 generated trajectories of length 100 and test with 40 trajectories. After training, we perform inverse optimization w.r.t. the boundary parameter with the trained models using Eq. [ starting with 50 initial configurations each with random initial location of smoke and random initial configuration of p. For L.F-PDF-Jatent and FNO-2D, they need to backpropagate through 80 steps of rollout in input space as in 45 74), while LE-PDE backpropagates through 80 steps of latent rollout. Then the optimized boundary parameter is fed to the ground-truth truth solver for rollout and evaluate. For the optimized parameter, we measure the total amount of smoke simulated by the solver passing through two respective outlets and take their ratio. The evaluation metric is the average ratio across all 50 configurations: see also Appendix G
Results. We observe that LE-PDE improves the overall speed by 73% compared with LE-PDE-latent and by 123% compared with FNO-2D (Fig. 3c Table Bb. The result indicates a corollary of the use of low dimensional representation because Jacobian matrix of evolution operator is reduced to be of smaller size and suppresses the complexity associated with the chain rule to compute gradients of the objective function. While achieving the significant speed-up, the capability of the LE-PDE to design the boundary is also reasonable. Fig. Ed] shows the loss of the objective function achieved the lowest value while the others are comparably large. The estimated proportion of smoke hit the target fraction 0.3 at an early stage of design iterations and coincide with the fraction simulated by the ground-truth solver in the end (Fig. 3e As Table 3 shows, FNO-2D achieves the lowest score in model estimated error from the target fraction 0.3 while its ground-truth solver (GT-solver) error is 30x larger. This shows "overfitting" of the boundary parameter by FNO-2D, i.e. the optimized parameter is not sufficiently generalized to work for ground-truth solver. In this sense, LE-PDE achieved to design the most generalized boundary parameter: the difference between the two errors is the smallest among the others.
4.4 Ablation study
In this section, we investigate how each component of our LE-PDE influences the performance. Importantly, we are interested in how each of the three components: multi-step loss Ltmulti-stpp latent consistency loss Leunsistency and recon- struction loss Lrccoms contribute to the performance, and how the time horizon M and the latent dimension d, in- fluence the result. For dataset, we focus on representative scenarios in 1D (Sec. 4. and 2D (Sec.
Table 4: Error for ablated versions of LE-PDE in ID and 2D.
ID
2D
LE-PDE (ours)
1.127
0.1861
no T-multi-step
3.337
0.2156
no Lecasistency
6.386
0.2316
no Lereooms
1.506
0.2025
specifically the F22 scenario with (n,, flg) (250, 50) for ID, and = 10=10-5 T = 20, N = 1000) scenario for 2D, which lies at mid- to difficult spectrum of each dataset. We have observed similar trends in other scenarios. From Table we see that all three components Liullti-tepp Lecnssstency and Liecons are necessary and pivotal in ensuring a good performance. The time horizon M in the loss is also important. If too short (e.g. M = 1). does not encourage accurate long-term rollout. Increasing M helps reducing error, but will be countered by less number of examples (since having to leave room for more steps in the future). We find the sweet spot at M- = 4, which achieves a good tradeoff. In Fig. 3 in Appendix H. we show how the error and evolution runtime change with varying size of latent dimension dz. We observe that reduction of runtime with decreasing latent dimension d,, and that the error is lowest at d, = 64 for ID and d2 256 for 2D, suggesting the intrinsic dimension of each problem.
Time horizon M-
5.710
0.2860
Time horizon = 1.234
0.2010
Time horizon M = 4 1.127
0.1861
Time horizon M = 6 1.924
0.1923

---

5 Discussion and Conclusion
In this work, we have introduced LE-PDE, simple, fast and scalable method for accelerating simulation and inverse optimization of PDEs, including its simple architecture, objective and inverse optimization techniques. Compared with state-of- ofte-oof-the-rr deep learning-basse based surrogate models, we demonstrate that it achieves up to 128 x reduction in the dimensions to update and up to 15x improvement speed, while achieving competitive aceuracy. Ablation study shows both multi-step objective and latent-consistecyy objectives are pivotal in ensuring accurate long-term rollout. We hope our method will make useful step accelerating simulation and inverse optimization of PDEs, pivotal in science and engineering.
Acknowledgments and Disclosure of Funding
We thank Sophia Kivelson, Jacqueline Yau, Rex Ying, Paulo Alves, Frederico Fiuza, Jason Chou, Qingqing Zhao for discussions and for providing feedback on manuscript. We also grate- fully Meknowlldge the support of DARPA under Nos. HR001 12190039 (TAMI), N66001 1924033 (MCS); ARO under Nos. W911NF-16 16-N--6--- (MURI), W911NNFI6-1-01 Nos. OAC-1835598 (CINES), OAC-1934578 (HDR), CCF-1918940 (Expeditions), NIH under No. 3U54HG010426-044 (HuBMAP), Stanford Data Science Initiative, Wu Tsai Neurosciences Institute, Amazon, Docomo, GSK, Hitachi, Intel, JPMorgan Chase, Juniper Networks, KDDI, NEC, and Toshiba.
(DURIP); NSF under
The content is solely the responsibility of the authors and does not necessarily represent the official views of the funding entities.
References
[1] P. Lynch, "The origins of computer weather prediction and climate modeling," Journal of computational physics, vol. 227, no. 7, pp. 3431-3444, 2008
[2] M. Athanasopoulos, H. Ugail, and G. G. Castro, "Parametric design of aircraft geometry using
partial differential equations," Advances in Engineering Software vol. 40, no. 7, pp. 479-486, 2009.
[3] F. Carpanese, "Development of free-bbundarr equilibrium and transport solvers for simulation and real-time interpretation of tokamak experiments," EPFL, Tech. Rep., 2021.
[4] N. Sircombe, Arber, and R. Dendy, "Kinetic effects laser-plasma coupling: Vlasov theory and computations," in Journal de Physique IV (Proceedings), vol. 133. EDP sciences, 2006, pp. 277-281.
[5] R. Courant, K. Friedrichs, and H. Lewy, "On the partial difference equations of mathematical physics," " IBM journal of Research and Development, vol. 11, no. 2, pp. 215-234, 1967.
[6] T. Lelievre and G. Stoltz, "Partial differential equations and stochastic methods in molecular dynamics," Acta Numerica, vol. 25, pp. 681-880) 2016.
[7] J. Brandstetter, D. E. Worrall, and M. Welling, "Message passing neural PDE solvers," International Conference on Leurning Representations, 2022. [Online]. Available: http://eeweeeeeeeee///o/mdddddda
[8] D. E. Keyes, D. R. Reynolds, and C. S. Woodward, "Implicit solvers for large-scale nonlinear problems," in Journal of Physics: Conference Series, vol. 46, no. 1. IOP Publishing, 2006, p. 060.
[9] Y. Dubois and R. Teyssier, "Cosmological MHD simulation of a cooling flow cluster," Astronomy & Astrophysics, vol. 482, no. pp. L13-L16, 2008
[10] P. Chatelain, A. Curioni, M. Bergdorf, D. Rossinelli, W. Andreoni, and P. Koumoutsakos, "Billion vortex particle direct numerical simulations of aircraft wakes." Computer Methods in Applied Mechanics and Engineering, vol. 197, no. 13-16, pp. 1296-1304, 2008
[11] L. T. Biegler, O. Ghattas, M. Heinkenschloss, and B. Bloemen Waanders, "Large-scale pde-constrained optimization: an introduction," in Large-Scale PDE-Constrained Optimization. Springer, 2003, pp. 3-13.
10

---

[12] K. Um, R. Brand, Y. R. Fei, P. Holl, and N. Thuerey. "Solver-in-the Learning from differentiable physies to interact with iterative pde-solvers," Advances in Neural Information Processing Systems, vol. 33, pp. 6111-6122, 2020.
[13] A. Sanehez-Gonzalez, J. Godwin, T. Pfaff, R. Ying, J. Leskovec, and P. Battaglia, "Learning to simulate complex physics with graph networks," in International Conference Machine Learning. PMLR, 2020, pp. 8459-8468.
[14] Z. Li, N. B. Kovachki, K. Azizzadenesheli, B. liu, K. Bhattacharya, A. Stuart, and A. Anandkumar, "Fourier neural operator for parametric partial differential equations," " in International Conference on Learning Representations, 2021. [Online]. Available:
https//opaaccvie..ee///oommcmmcp http://orrumnddddPPo/VVoii VimnO
[15] M. Tang. Y. Liu, and L. J. Durlofsky, "A dcep-learningg-asee based surrogate model for data assimi- lation in dynamic subsurface flow problems," Journal of Computational Physics, vol. 413, p. 109456, 2020.
[16] T. Wu, Q. Wang, Y. Zhang, R. Ying, K. Cao, R. Sosic, R. Jalali, H. Hamam, M. Maucec, and J. Leskovec, "Learning large-scale subsurface simulations with a hybrid graph network simulator," in Proceedings of the 28th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, 2022, pp. 4 4184-4194.
[17] D. Kochkov, J. A. Smith, A. Alieva, Q. Wang, M. P. Brenner, and S. Hoyer, "Machine learning- accelerated computational fluid dynamics," Proceedings of the National Academy of Sciences, vol. 118, no. 21, 2021.
[18] A. Sanchez, D. Kochkov, J. A. Smith, M. Brenner, P. Battaglia, and T. J. Pfaff, "Learning latent field dynamics of PDEs," Advances in Neural Information Processing Systems, 2020.
[19] N. Watters, D. Zoran, T. Weber, P. Battaglia, R. Pascanu, and A. Tacchetti, "Visual interaction networks: Learning physics simulator from video," Advances in neural information processing systems, vol. 30, 2017.
[20] S. van Steenkiste, M. Chang, K. Greff, and J. Schmidhuber, "Relational neural expectation maximization: Unsupervised discovery of objects and their interactions," in International Conference on Learning Representations, 2018. [Online]. Available:
[21] S.-M. Udrescu and M. Tegmark, "Symbolic pregression: discovering physical laws from distorted vidco," Physical Review E, vol. 103, no. 4, p. 043307, 2021.
[22] C. Gelada, S. Kumar, Buckman, O. Nachum, and M. G. Bellemare, "Deepmdp: Learning
continuous latent space models for representation learning." in International Conference
Machine Learning. PMLR, 2019, pp. 2170-2179.
[23] D. Hafner, T. Lillicrap, I. Fischer, R. Villegas, D. Ha, H. Lee, and J. Davidson, "Learning latent dynamics for planning from pixels," in International conference on machine learning. PMI.R,
2019, pp. 2555-2565.
[24] R. Julian, F. Heiden, 7. He, H. Zhang. S. Schaal, J. Lim, G. S. Sukhatme, and K. Hausman, "Scaling simulation-to-real transfer by learning latent space of robot skills," The International Journal of Robotics Research, vol. 39, no. 10-11, 1, pp. 1259-1278, 2020.
[25] A. X. Lee, A. Nagabandi, P. Abbeel, and S. Levine, "Stochastic latent actor-critic: Deep reinforcement learning with latent variable model," Advances in Neural Information Processing Systems, vol. 33, pp. 741-752, 2020.
[26] R. Wang, K. Kashinath, M. Mustafa, A. Albert, and R. Yu, "Towards physics-inforinn deep learning for turbulent flow prediction," in Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, 2020, pp. 1457-1466.
[27] T. Pfaff, M. Fortunato, A. Sanchez-Gonzater and P. W. Battaglia, "Learning mesh-base based simulation with graph networks," in International Conference on Learning Representations. 2021.
[28] Z. Li and A. B. Farimani, "Graph neural network-accelerated lagrangian fluid simulation," Computers Graphics, vol. 103, pp. 201-211, 2022.
[29] Ò. Çięek, A. Abdulkadir, S. S. Lienkamp, T. Brox, and O. Ronneberger, "3d u-net: learning dense volumetric segmentation from sparse annotation," International conference on medical image computing and computer-assisted intervention. Springer, 2016, pp. 424 4322.
11

---

[30] M. Raissi, "Deep hidden physics models: Deep learning of nonlinear partial differential equa- tions." " The Journal of Machine Learning Research, vol. 19, no. 1, pp. 932-955, 2018.
[31] Y. Zhu and N. Zabaras, "Bayesian deep convolutional encoder- networks for surrogate modeling and uncertainty quantification." Journal of Computational Physics, vol. 366, pp. 415-447, 2018.
[32] S. Bhatnagar, Y. Afshar, S. Pan, K. Duraisamy, and S. Kaushik, "Prediction of aerodynamic flow fields using convolutional neural networks," Computational Mechanics, vol. 64, no. 2. pp. 525-545, 2019.
[33] Y. Khoo, Lu, and L. Ying, "Solving parametric pde problems with artificial neural networks," European Journal of Applied Mathemalics, vol. 32, no. 3, pp. 421-435, 2021.
[34] Z. Li, N. Kovachki, K. Azizzadenesheli, B. Liu, K. Bhattacharya, A. Stuart, and A. Anandkumar, "Neural operator: Graph kernel network for partial differential equations," arXiv preprini
arXiv:2003. 03345 2020.
[35] Z. N. Kovachki, K. Azizzadenesheli, B. Liu, A. Stuart, K. Bhattacharya, and Anandkumar, "Multipole graph neural operator for parametric partial differential equations," Advances in Neural Information Prucessing Systems, vol. 33. pp. 6755-6766, 2020.
[36] L. Lu, P. Jin, G. Pang, Z. Zhang, and G. Karniadakis, "Learning nonlinear operators via deeponet based on the universal approximation theorem of operators. nature mach. intell. 3 (3), 218-229 (2021)."
[37] K. T. Butler, J. M. Frost, J. M. Skelton, K. L. Svane, and A. Walsh, "Computationall materials design of crystalline solids," Chemical Society Reviews, vol. 45, no. 22, pp. 6138-6146 2016.
[38] I. Vernon, M. Goldstein, and R. Bower, "Galaxy formation: Bayesian history matching for the observable universe," Statistical science, pp. 81-90, 2014.
[39] D. Williamson, M. Goldstein, L. Allison, A. Blaker, P. Challenor, L. Jackson, and K. Yamazaki, "History matching for exploring and reducing climate model parameter space using observations and a large perturbed physics ensemble," Climate dynamics, vol. 41, no. 7, pp. 1703-1729, 2013.
[40] D. Oliver and Y. Chen, "Recent progress on reservoir history matching: a review," Computa- tional Geoxciences, vol. 15, no. 1, pp. 185-221, 2011.
[41] O. Talagrand and P. Courtier, "Variational assimilation of meteorological observations with the adjoint vorticity equation. i: Theory." Quarterly Journal of the Royal Meteorological Society, vol. 113, no. 478, pp. 1311-113 1-1328, 1987.
[42] Tromp, C. Tape, and Q. Liu, "Seismic tomography, adjoint methods, time reversal and banana- doouhnu kernels," Geophysical Journal International, vol. 160, no. 1, pp. 195-216, 2005.
[43] H. B. Keller, Numerical solution of two point boundary value problems. SIAM, 1976.
[44] J. T. Betts, "Survey of numerical methods for trajectory optimization," Journal of guidance, control, and dynamics, vol. 21, no. 2. pp. 193-207, 1998.
[45] K. R. Allen, T. Lopez-Guevvraa, K. Stachenfeld, A. Sanchez-Gonzaler P. Battaglia, J. Ham- rick, and T. Pfaff, "Physical design using differentiable learned simulators," arXiv preprint arXiv:2202.00728. 2022.
[46] R. Y. Rubinstein and D. P. Kroese, "The cross-entropy method: A unified approach to monte carlo simulation, randomized optimization and machine learning." Information Science & Statistics, Springer Verlag, NY, 2004.
[47] A. Treuille, A. Lewis, and Z. Popović, "Model reduction for real-time fluids ACM Transactions on Graphics (TOG), vol. 25, no. 3, pp. 826-834- 2006.
[48] G. Berkooz, Holmes, and J. L. Lumley. "The proper orthogonal decomposition in the analysis of turbulent flows," Annual review fluid mechanics, vol. 25, no. 1, pp. 539-575, 1993.
[49] M. Gupta and S. G. Narasimhan, "Leggndre fluids: a unified framework for analytic reduced space modeling and rendering of participating media," in Symposium on Computer Animation. 2007, pp. 17-25.
[50] M. Wicke, M. Stanton, and A. Treuille, "Modular bases for fluid dynamies." ACM Transactions on Graphics (TOG), vol. 28, no. 3, pp. 1-8, 2009.
12

---

[51] B. Long and E. Reinhard, "Real-time fluid simulation using discrete sine/cosine transforms." in Proceedings of the 2009 symposium on Interactive 3D graphics and games, 2009, pp. 99-106.
[52] T. De Witt, C. Lessig, and Fiume. "Fluid simulation using laplacian cigenfunctions." ACM Transactions on Graphics (TOG), vol. 31, no. I, pp. 1-11, 2012.
[53] T. Kim and J. Delaney, "Subspace fluid re-simulation," ACM Transactions on Graphics (TOG),
vol. 32, no. 4, pp. 1-9, 2013.
[54] B. Liu, G. Mason, J. Hodgson, Y. Tong, and M. Desbrun, "Model-reduced variational fluid simulation," ACM Transactions on Graphics (TOG), vol. 34, no. 6, pp- 1-12, 2015.
[55] A. Radford, L. Metz, and S. Chintala, "Unsupervised representation learning with deep convolutional generative adversarial networks," in 4th International Conference on Learning Representations, 2016. [Online]. Available:
[56] S. Wiewel, M. Becher, and N. Thuerey, "Latent space physics: Towards learning the temporal evolution of fluid flow," in Computer graphics forum, vol. 38, no. 2. Wiley Online Library, 2019, pp. 71-82.
[57] B. Kim, V. C. Azevedo, N. Thuerey, T. Kim, M. Gross, and B. Solenthaler, "Deep fluids: A generative network for parameterized fluid simulations," in Computer Graphics Forum, vol. 38, no. Wiley Online Library, 2019, pp. 59-70.
[58] K. Lee and K. Carlberg, "Model reduction of dynamical systems on nonlinear manifolds using deep convolutional autoencoders," Journal of Computational Physics, vol. 404, 108973, 2020.
[59] Wiewel, B. Kim, V. C. Azevedo, B. Solenthaler, and N. Thuerey, "Latent space subdivision: stable and controllable time predictions for fluid flow," in Computer Graphics Forum, vol. 39, no. Wiley Online Library, 2020, pp. 15-25.
[60] R. Vlachas, G. Arampatzis, C. Uhler, and P. Koumoutsakos, "Multiscale simulations of complex systems by learning their effective dynamics," Nuture Machine Intelligence, vol. 4, no. 4, pp. 359-366, 2022.
[61] D. P. Kingma and Ba, "Adam: A method for stochastic optimization," in International Conference Learning Representations (Poster), 2015. [Online]. Available:
[62] N. Tishby, F. C. Pereira, and W. Bialek, "The information bottleneck method." arXiv preprint physics/0004057, 2000.
[63] T. Wu and I. Fischer, "Phase transitions for the information bottleneck in representation learning," in International Conference on Learning Representations, 2020. [Online]. Available:
YvB
[64] Y. Bar-Sinai, S. Hoyer, J. Hickey, and M. P. Brenner, "Learning data-driven discretizations for partial differential equations." Proceedings of the National Academy of Sciences, vol. 116, no. 31, pp. 15 344-15 349, 2019.
[65] C.-W. Shu, "High-order finite difference and finite volume weno schemes and discontinuous galerkin methods for efd," International Journal of Computational Fluid Dynamics, vol. 17, no. 2, pp. 107-118, 2003.
[66] C. Runge, "Über die numerische auflösung von differeentialllichhnnnn, Mathematische An- nalen, vol. 46, no. 2. pp. 167-178 1895.
[67] W. Kutta, "Beitrag
naherungsweisen integration totaler differentialgleichungen Z. Math. Phys., vol. 46, pp. 435 435-453, 1901.
[68] O. Ronneberger, P. Fischer, and Brox, "U-net: Convolutional networks for biomedical image segmentation," in International Conference Medical image computing and computer-assisted intervention. Springer, 2015, pp. 234-241.
[69] K. He, X. Zhang, S. Ren, and J. Sun, "Deep residual learning for image recognition," " in Proceedings of the IEEE conference computer vision and pattern recognition, 2016, pp. 770-778.
[70] A. McNamara, A. Treuille, Z. Popović, and Stam, "Fluid control using the adjoint method," ACM Transactions On Graphics (TOG), vol. 23, no. 3, pp. 449-456, 2004.
13

---

[71] P. Holl, N. Thuerey, and V. Koltun, "Learning to control pdes with differentiable physics," " in International Conference on Learning Representations, 2020. [Online]. Available: https//opeeneviiewnee/foooo http:/orruunndd--HyeSSerrPPa
[72] A. Zylstra, O. Hurricane, D. Callahan, A. Kritcher, Ralph, H. Robey, J. Ross, C. Young, K. Baker, D. Casey et al., "Burning plasma achieved in inertial fusion," Nature, vol. 601, no.
7894, pp- 542-548, 2022.
[73]
"Phiflow," https://gtthu..com/umu wow/tum-pps///illann
[74] Q. Zhao, D. B. Lindell, and G. Wetzstein, "Learning to solve pde-constrained inverse problems with graph networks," International Conference Machine Learning, 2022.
[75] K. Roberts and H. L. Berk, "Nonlinear evolution of a two-stream instability," Physical Review Letters, vol. 19, no. 6, p. 297, 1967.
Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's
contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes] In Appendix
(c) Did you discuss any potential negative societal impacts of your work? [Yes] In
Appendix[I [
(d) Have you read the ethics review guidelines and ensured that your paper conforms them? [Yes]
2. If you including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [N/A]
(b) Did you include complete proofs of all theoretical results? [N/A]
3. If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either the supplemental material or as a URL)? [Yes] The Appendix includes full details on model architecture, training and evaluation to reproduce the experimental results. Code and data will be released upon publication of the paper. (b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? (Yes] Important training details are included in main text, and full details reproduce the experiments (c) Did you report error bars (e.g., with respect the random seed after running ex- periments multiple times)? [Yes] Yes Section 4.4. For Sections 4.1 and 4.2, the benchmarks do not include the error bars and use the same format.
included in the Appendix.
(d) Did you include the total amount of compute and the type of resources used (c.g., type of GPUs, internal cluster, or cloud provider)? [Yes] In Appendix D
4. If you using existing assets (e.g., code, data, models) or curating/releusing
assets...
(a) If your work uses existing assets, did you cite the creators? [Yes] In Appendix D
(b) Did you mention the license of the assets? [Yes] In Appendix DEFF
(c) Did you include any new assets either the supplemental material or as URL? Yes]
In Appendix
DEFG
(d) Did you discuss whether and how consent was obtained from people whose data you're using/curating? [Yes] In Appendix G
(e) Did you discuss whether the data you using/curating contains personally identifiable information or offensive content? [Yes] In Appendix D
5. If you used crowdsourcing conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, applicable? [N/A]
(b) Did you describe any potential participant risks, with links Institutional Review
Board (IRB) approvals, applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent participant compensation? [N/]]
14


Reviews:
[{"review_id": "o_mqVYBZFUG", "reviewer": "Reviewer_bK7j", "summary": "The paper proposes a neural PDE solver based on learning how to represent the system with an evolving latent global state.\nThe advantage is reducing the cost of approximating PDE solutions without significantly reducing accuracy.\nThe method can also be used to solve some inverse problems.\n", "questions": "## Major\n\n* In the objective defined in line 127, is the right hand side missing integration over $x$? If not, then there is a different objective for every $x$ location. How are these combined? (Also, less importantly, what does the subscript $d$ refer to?)\n\n* In Table 1, why is the WEBO5 accumulated error so large if this is the ground truth method?\n\n## Minor\n\n* As someone not familiar with this field,\nI found the subsection on \"Deep learning-based surrogate modeling\" reviewed autoregressive methods very clearly.\nHowever the 1-sentence summary of \"neural operators\" was too brief for me to understand.\nCan you expand on the details of \"neural networks that approximate a mapping between infinite-dimensional functions\",\nas this seems non-trivial to do.\n* Line 107. In what sense is $f$ a partial function?\n* Table 1. Why is runtime for FNO not included?\n* Line 233-234. The text says you start with $n_t=200$ then downsample. But downsampling is to $n_t=250$, which is a larger value!\nIs there a typo here? If not, why not keep the original value of $n_t$?\n* Line 314: \"our LE-PDE's ablated version without latent evolution\". Can you describe this in more detail?\nI couldn't find a description in the main paper.\n(Maybe it's somewhere in the supplement - in which case adding a reference would be helpful.)\n* Line 354 \"Increasing M... will be countered by less number of examples (since having to leave room for more steps in the future)\"\nI don't understand this, can you explain in more detail?", "limitations": "The paper covers technical and ethical limitations clearly in the supplementary material.\nI agree with the authors' assessment that there are no obvious negative social impacts.\nI also appreciate that the authors included experimental results where existing methods (slightly) outperform the proposed approach.\nThis helps readers get a clear picture of the relative strengths of the methods in different scenarios.", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "## Summary (and strengths)\n\n* Originality. The paper proposes an approach which is simple to understand and very general (allowing any architecture to be used for the various components e.g. encoders, decoders). Furthermore, it appears novel to me, although I'm not familiar with the neural PDEs literature.\n* Quality. The proposed method is always at least competitive with existing methods, and outperforms them in some important examples. The experiments are thorough, covering all the important ablations I could think of while reading the paper.\n* Clarity. Overall the paper is clearly written and easy to understand, but there are some minor problems discussed below.\n* Significance.  The authors make the case for the importance of neural PDE solver well. \n\nIn my opinion this is a good paper and clearly above the acceptance threshold.\n\n## Weaknesses\n\nI found few important weaknesses in the paper. Some minor points are listed below and in the \"Questions\" section.\n\nWhile the paper is generally well written, there are a few typos and confusingly worded sentences.\nMost of these are unimportant, but below I list a few which made the paper harder to understand:\n\n* Line 75 \"flow probe\" -> \"flow to probe\"?\n* Line 89 \"location... that satisfies given\" Remove \"that satisfies\"?\n* Line 122 \"LE-PDE relieves of local evolution\" I don't understand \"relieves\" in this context - use another word?\n* Line 205 \"The prevents the gradient to pass through to the boundary parameter $p$ such as continuous location.\" I'm struggling to parse this sentence, though I think the point is that it's not possible to backprop through discrete variables.\n* Line 236 \"$\\eta \\in [0,0,2]$\" Should this say \"$[0,2]$\"?\n* Line 324 \"average amount of the advected smoke simulated by the solver\" I don't follow what this means and I can't find the exact details in the supplement.", "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "Yvsk7vvSCsb", "reviewer": "Reviewer_MdSE", "summary": "The authors propose a method LE-PDE to efficiently perform forward simulation and inverse optimization of Partial differential equation based models by performing the evolution in latent space. A loss term penalizing 3 terms including the consistency between latent space and the original space is defined, and optimized via backpropagation for the inverse task. Experimental evaluations including an ablation study is provided.\n", "questions": "Please have a look at the weakness mentioned in the strength and weakness section and address these. Overall the idea seems interesting, the authors need to substantiate their claims in light of existing literature and possibly a few more experiments.", "limitations": "Yes, the limitations are discussed in the appendix.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strength:\nThe authors put together a very clear introduction and motivation for their work, the method is described clearly in an easily reproducible manner.\n\nWeakness:\nThe authors state that low dimensional representation exists for high dimensional data (line 49-64). Please have a look at Johnson-Lindenstrauss lemma and cite appropriate literature applicable here.\n\nThe authors may want to try their experiments on a larger scale. Grid size of 64 is too small, as the authors mention dimensions in the millions/billions in the introduction, at least one experiment should demonstrate the efficacy of the proposed method in such a larger scale where other baseline methods are computationally very very expensive and resource consuming.\n\nAnother concern about the proposed method is the applicability of the proposed system in real-world scenarios. Experimental data are often noisy. The authors may want to look at how noise affects the latent space evolution and encoder-decoder performance.\n", "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "VqGOaHv9wJp", "reviewer": "Reviewer_jbxg", "summary": "This paper builds on recent advances in deep learning for physical simulation to reduce the computational complexity of learned surrogate models by avoiding local evolution. Instead of updating the values at each discretized point in space, the authors propose to evolve the model in a global latent space. By encoding both states and boundaries into lower-dimensional latent vectors, the model can evolve dynamics in the latent space and recover the inferred states only when needed. The paper also introduces techniques for training, such as a new consistency loss directly in latent space, a fast way of dealing with inverse optimization problems by backpropagating through time in the latent space, and an annealing technique for boundary optimization. Experiments in various PDE settings demonstrate that the proposed model, LE-PDE (Latent-Evolution of PDEs) compares competitively against state-of-the-art methods while requiring a fraction of the other methods’ computational resources.", "questions": "1. GNN+MLPs as a future direction, but it is unclear to me whether this could be possible. How would message passing be realized in a global latent space, though? Could we still use the same latent space even if we added particles or mesh points?\n\n2. About the boundary annealing, what is meant by $\\beta$ being a “temperature hyperparameter”? This notation is also not introduced in the main text.", "limitations": "See above. As stated in the main text, no major negative societal impact is to be expected from this work.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "### Strenghts\n\nThe paper is appropriately placed in the current growing literature on scientific machine learning to create a novel fast model for tackling complex problems in the realm of PDEs. The exposition for the motivation is written crisply and is generally easy to follow. The experiments are performed in challenging settings, such as boundary control with the 2D Navier Stokes equations. The results show that very large neural networks are not always necessary to represent complex physics and that “simple” latent MLPs with few parameters can learn such physics evolutions thus considerably lowering computational requirements, which is significant in deep learning for simulation.\n\n### Weaknesses\n\nAlthough the speedups of LE-PDE are considerable, it is unclear how the proposed approach would compare against other models with a similar number of parameters. What if the FNO had fewer parameters, comparable to LE-PDE? Would it still be better than the proposed approach? In other words, it is unclear whether the method is *Pareto efficient* or not since there is no Pareto plot considering an extensive hyperparameter search for both LE-PDE and close competitors such as the FNO.\n\nAbout the novelty: there is a possibly missing reference [1] (although not published in conference proceedings), in which the authors evolve a model in the latent space (see Figure 1 and Section 3.1 of the reference), and the main idea of the paper is very similar to the presented one. How does this work compare against it?\n\nThe experimental section in the Appendix about the 3D extension of the method (3D Navier Stokes) lacks some experimental details and evaluation. The learned model is not compared with other deep learning approaches, so it is difficult to assess how the proposed method would perform against other models (except for the same version of the model without the global latent space evolution). While the comparison in the main text is against other deep learning approaches, here it is against the ground truth solver which is known to be generally much slower than deep learning approaches. Moreover, there is no report on statistical errors but only a qualitative plot (that also shows noticeable artifacts).\n\n[1] Sanchez A, Kochkov D, Smith JA, Brenner M, Battaglia P, Pfaff TJ. Learning latent field dynamics of PDEs,  Third Workshop on Machine Learning and the Physical Sciences (NeurIPS 2020)\n\n### Minor comments\n\nThese are primarily typos and do not influence my score:\n\n- Line 235-236: “Burgur’s Equation”\n- Line 355: “sweep spot”\n- Line 707 (Appendix): \"The details of the dataset has already given in […]”\n- Line 748: “To explore how LE-PDE to larger scale turbulent dynamics a […]”", "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "KIsnb7cTFc", "reviewer": "Reviewer_Mytm", "summary": "This paper looks at the problem of expensive time cost when simulating the time evolution of a PDE. Existing methods employ values at different spatial positions at each time step, causing a long simulating time. LE-PDE accelerates by learning a low dimensional latent representation, and evolve the low-dimensional state, rather than the high-dimensional original variable. Optimization is performed upon reconstruction loss of the latent state and the evolution accuracy. Empirical advantages on time cost achieved over previous methods on different PDEs.", "questions": "It would be interesting to see a study of when can the proposed LE-PDE method have a reasonable result. Does the reconstruction take extra data to train? \nAlso, there has been numerous dimension reduction techniques, PCA, VAEs, just to name a few. Why did the author pursue the proposed method rather than making use of the existing ones? Are there any theoretical relationships between the proposed method and the existing ones?", "limitations": "The author proposed to accelerate evolution of PDEs by first reducing its dimension. \nThe idea is intuitive and straight forward, and is proven empirically with different PDEs. \nHowever, a theoretical study on what PDEs this reduction can be efficiently applied to is not conducted, and it is not argued why the proposed reduction method is better than existing ones like VAEs. \nThe method is flexible, and should be able to be accompanied with existing surrogate models and other method.", "rating": 5, "confidence": 2, "soundness": 2, "presentation": 3, "contribution": 1, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths: \nPaper demonstrated empirical improvements on both 1D nonlinear PDEs and 2D Navier-Stokes PDE over multiple previous methods. \nModel has a relatively flexible structure.\n\nWeaknesses: \nDid not study when can PDE states be encoded into a lower dimensional state, and.or how many dimension can be get rid of. \nDid not study the number of parameters used against other methods, using an extra encoder and decoder might take extra number of parameters, resulting in unfair advantages over other methods.\nFigure 1 is unclear. Specifically, is the box on the left the initial condition and boundary condition of the PDE? The schematic in the middle is also unclear.", "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept