Learning to Accelerate Partial Differential Equations via Latent Global Evolution
Anonymous Author(s)
AfﬁliationAddress email
Abstract
Simulating the time evolution of Partial Differential Equations (PDEs) of largescale systems is crucial in many scientiﬁc and engineering domains such as ﬂuiddynamics, weather forecasting and their inverse optimization problems.
However,both classical solvers and recent deep learning-based surrogate models are typicallyextremely computationally intensive, because of their local evolution: they need toupdate the state of each discretized cell at each time step during inference.
Here wedevelop Latent Evolution of PDEs (LE-PDE), a simple, fast and scalable methodto accelerate the simulation and inverse optimization of PDEs.
LE-PDE learns acompact, global representation of the system and efﬁciently evolves it fully in thelatent space with learned latent evolution models.
LE-PDE achieves speedup byhaving a much smaller latent dimension to update during long rollout as comparedto updating in the input space.
We introduce new learning objectives to effectivelylearn such latent dynamics to ensure long-term stability.
We further introducetechniques for speeding-up inverse optimization of boundary conditions for PDEsvia backpropagation through time in latent space, and an annealing technique toaddress the non-differentiability and sparse interaction of boundary conditions.
Wetest our method in a 1D benchmark of nonlinear PDEs, 2D Navier-Stokes ﬂowsinto turbulent phase and an inverse optimization of boundary conditions in 2D
Navier-Stokes ﬂow.
Compared to state-of-the-art deep learning-based surrogatemodels and other strong baselines, we demonstrate up to 128 ⇥ reduction in thedimensions to update, and up to 15 ⇥ improvement in speed, while achievingcompetitive accuracy.
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 a (complex)system evolves with time and numerically evolving such equations are used for forward predictionand inverse optimization across many disciplines.
Example application includes weather forecasting[1], jet engine design [2], nuclear fusion [3], laser-plasma interaction [4], astronomical simulation[5], and molecular modeling [6].
To numerically evolve such PDEs, decades of works have yielded (classical) PDE solvers that aretailored to each speciﬁc problem domain [7].
Albeit principled and accurate, classical PDE solversare 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 thedomain is discretized into [8].
For practical problems in science and engineering, the number ofcells at each time step can easily go up to millions or billions and may even require massivelyparallel supercomputing resources [9, 10].
Besides forward modeling, inverse problems, such asinverse optimization of system parameters and inverse parameter inference, also share similar scaling challenge [11].
How to effectively speed up the simulation while maintaining reasonable accuracyremains an important open problem.
Recently, deep learning-based surrogate models have emerged as attractive alternative to complement[12] or replace classical solvers [13, 14].
They directly learn the dynamics from data and alleviatemuch engineering effort.
They typically offer speedup due to explicit forward mapping [15], largertime intervals [14], or modeling on a coarser grid [12, 16].
However, their evolution scales with thediscretization, since they typically need to update the state of each discretized cell at each time step,due to the local nature of PDEs [17].
For example, if a problem is discretized into 1 million cells,deep learning-based surrogate models (e.g., CNNs, Graph Networks, Neural Operators) will needto evolve these 1 million cells per time step.
How to go beyond updating each individual cells andfurther speed up such models remains a challenge.
Here we present Latent Evolution of PDEs (LE-PDE) (Fig. 1), a simple, fast and scalable methodto accelerate the simulation and inverse optimization of PDEs.
Our key insight is that a commonfeature of the dynamics of many systems of interest is the presence of dominant, low-dimensionalcoherent structures, suggesting the possibility of efﬁciently evolving the system in a low-dimensionalglobal latent space.
Based on this observation, we develop Latent Evolution of PDEs (LE-PDE),which learns the evolution of dynamics in a global latent space.
Here by “global” we mean that thedimension of the latent state is ﬁxed, instead of scaling linearly with the number of cells as in localmodels.
LE-PDE consists of a dynamic encoder that compresses the input state into a dynamic latentvector, a static encoder that encodes boundary conditions and equation parameters into a static latentvector, and a latent evolution model that evolves the dynamic latent vector fully in the latent space,and decode via a decoder only as needed.
Although the idea of latent evolution has appeared in otherdomains, such as in computer vision [18, 19, 20] and robotics [21, 22, 23], these domains typicallyhave clear object structure in visual inputs allowing compact representation.
PDEs, on the other hand,model dynamics of continuum (e.g., ﬂuids, materials) with inﬁnite dimensions, without a clear objectstructure, and sometimes with chaotic turbulent dynamics, and it is pivotal to model their long-termevolution 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-termevolution in latent space.
Together with the multi-step loss in the input space, they encourage moreaccurate long-term prediction.
To accelerate inverse optimization of PDEs which is pivotal in engineering (e.g. optimize the boundary condition so that the evolution of the dynamics optimizes certainpredeﬁned objective), we show that LE-PDE can allow faster optimization, via backpropagationthrough time in latent space instead of in input space.
To address the challenge that the boundarycondition may be non-differentiable or too sparse to receive any gradient, we design an annealingtechnique for the boundary mask during inverse optimization.
We demonstrate our LE-PDE method in standard PDE-learning benchmarks of a 1D family ofnonlinear PDEs and a 2D Navier-Stokes ﬂow into turbulent phase, and design an inverse optimizationproblem in 2D Navier-Stokes ﬂow probe its capability.
Compared with state-of-the-art deep learningbased surrogate models and other strong baselines, we show up to 128 ⇥ reduction in the dimensionsto update and up to 15 ⇥ speedup compared to modeling in input space, and competitive accuracy.
Problem Setting and Related Work
We consider temporal Partial Differential Equations (PDEs) w.r.t.
time t 2 [0 , T ] and multiple spatialdimensions x = [ x 1 , x 2 , ...
x D ] 2 X ✓ R D .
We follow similar notation as in [7].
Speciﬁcally,
[Equation]
Here u : [0 , T ] ⇥ X !
R n is the solution, which is an inﬁnite-dimensional function.
a is timeindependent static parameters of the system, which can be deﬁned on each location x , e.g. diffusioncoefﬁcient that varies in space but static in time, or a global parameter.
F is a linear or nonlinearfunction on the arguments of ( x , a , u , @ x u , @
xx u , ... ) .
Note that in this work we consider timeindependent PDEs where F does not explicitly depend on t .
u 0 ( x ) is the initial condition, and
B [ u ]( t, x ) = 0 is the boundary condition when x is on the boundary of the domain @
X across alltime t 2 [0 , T ] .
Here @ x u = @ u @ x , @ xx u
= @ 2 u @ x 2 are ﬁrst- and second-order partial derivatives, whichare a matrix and a 3-order tensor, respectively (since x is a vector).
Solving such temporal PDEsmeans computing the state u ( t, x ) for any time t 2 [0 , T ] and location x 2 X that satisﬁes given theabove initial and boundary conditions.
Classical solvers for solving PDEs .
To numerically solve the above PDEs, classical numericalsolvers typically discretize the domain X into a ﬁnite grid or mesh X = { c i } Ni =1 with N nonoverlapping cells.
Then the inﬁnite-dimensional solution function of u ( t, x ) is discretized into
U k
= { u k i } N i = 1 2 U for each cell i and time t =
t k , k = 1 , 2 , ...K .
a is similarly discretized into{ a i } Ni =1 with values in each cell.
Mainstream numerical methods, including Finite Difference Method 95 (FDM) and Finite Volume Method (FVM), proceed to evolve such temporal PDEs by solving theequation at state { u k +1 i } at time t = t k +1 from state { u ki } at time t k .
These solvers are typically 97 slow due to small time/space intervals required for numerical stability, and needing to update eachcell at each time steps.
For more detailed information on classical solvers, see Appendix A
Deep learning-based surrogate modeling .
There are two main approaches in deep learningbased surrogate modeling.
The ﬁrst class of method is autoregressive methods, which learns themapping f ✓ with parameter ✓ of the discretized states U k between consecutive time steps t k andt k
+1 :
ˆ U k +
=
f ✓ ( ˆ U k , p ) , k = 0 , 1 , 2 , ... .
Here ˆ U k
= { ˆ u ki } Ni =1 is the model
f ✓ ’s predicted statefor
U k
= { u ki } Ni =1 at time t k , with ˆ U 0 :
= U 0 .
p = ( @ X , { a i } Ni =1 ) is the system parameter whichincludes the boundary domain @
X and discretized static parameters { a i } Ni =1 .
Repetitively apply f ✓ at 105 inference time results in autoregressive rollout
[Equation]
Here f ✓ ( · , p ) : U !
U is a partial function whose second argument is fulﬁlled by the static systemparameter p .
Typically f ✓ 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 allinvolve local computation, where the value u k +1 i at cell i at time t k +1 depend on its neighbors{ u kj } j 2 N ( i ) at time t k , where N ( i ) is the set of neighbors up to certain hops.
Such formulation 111 includes CNN-based models [24], GNN-based models [25, 7, 26] and their hierarchical counterparts[27, 13].
The surrogate modeling with local dynamics makes sense, since the underlying PDEis essentially a local equation that stipulates how the solution function u ’s value at location xdepends on the values at its inﬁnitesimal neighborhood.
The second class of method is Neural Operators [28, 29, 30, 31, 32, 33, 14, 34], which learns a neural networks that approximate a mappingbetween inﬁnite-dimensional functions.
Although having the advantage that the learned mappingis discretization invariant, given a speciﬁc discretization, Neural Operators still needs to update thestate at each cell based on neighboring cells (and potentially cells far away), which is still inefﬁcientat inference time, especially dealing with larger-scale problems.
In contrast to the above two classesof deep learning-based approaches that both requires local evolution at inference time, our LE-PDEmethod focuses on improving efﬁciency.
Using a learned global latent space, LE-PDE relieves oflocal evolution and can directly evolve the system dynamics via a global latent vectors z k 2 R d z fortime t k .
This offers great potential for speedup due to the signiﬁcant reduction in representation.
Inverse optimization .
Inverse optimization is the problem of optimizing the parameters p of the
PDE, including boundary @
X or static parameter a of the equation, so that a predeﬁned objective
L d [ a , @
X ] =
R t e t = t s ` d ( u ( t, x )) dt is minimized.
Here the state u ( t, x ) implicitly depends on a , @
Xthrough the PDE (Eq. 1) and the boundary condition (Eq. 2).
Such problems have huge importancein engineering, e.g. in designing jet engines [2] and materials [35] where the objective can beminimizing drag or maximizing durability, and inverse parameter inference (i.e. history matching)
[36, 37, 38] where the objective can be maximum a posteriori estimation.
To solve such problem,classical methods includes adjoint method [39, 40], shooting method [41], collocation method [42],etc.
One recent work [43] explores optimization via backpropagation through differential physics inthe input space, demonstrating speed-up and improved accuracy compared to classical CEM method[44].
However, for long rollout and large input size, the computation becomes intensive to the pointof needing to save gradient in ﬁles.
In comparison, LE-PDE allows backpropagation in latent space,and due to the much smaller latent dimension and evolution model, it can signiﬁcantly reduce thetime complexity in inverse optimization.
Reduced-order modeling .
A related class of work is reduced-order modeling.
Past efforts typicallyuse linear projection into certain basis functions [45, 46, 47, 48, 49, 50, 51, 52] which may not haveenough representation power.
A few recent works explore NN-based projection [53, 54, 55, 56, 57]for ﬂuid modeling.
Compared to the above works, we focus on speeding up simulation and inverseoptimization of more general PDEs using expressive NNs, with novel objectives, and demonstratecompetitive performance compared to state-of-the-art deep learning-based models for PDEs.
Our approach LE-PDE
In this section, we detail our Latent Evolution of Partial Differential Equations (LE-PDE) method.
We ﬁrst introduce the model architecture (Sec. 3.1, and then we introduce learning objective toeffectively learn faithfully long-term evolution (Sec. 3.2).
In Sec. 3.3, we introduce efﬁcient inverseoptimization 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 !
R d zthat maps the input state
U k
= { u ki } Ni =1 2 U to a latent vector z k 2 R
d
z ; (2) an (optional) staticencoder r : P !
R d zp that maps the (optional) system parameter p 2 P to a static latent embeddingz p ; (3) a decoder h : R d z !
U that maps the latent vector
z k 2 R d z back to the input state U k ;(4) a latent evolution model g : R d z ⇥ R d
zp !
R d z that maps z k 2 R d z at time t k and static latentembedding z p 2 R d zp to z k +1 2 R d z at time t k +1 .
We employ the temporal bundling trick [7] whereeach input state U k can include states over a ﬁxed length S of consecutive time steps, in which caseeach latent vector z k will encode such bundle of states, and each latent evolution will predict thelatent vector for the next bundle of S steps.
S is a hyperparameter and may be chosen dependingon the problem, and S = 1 reduces to no bundling.
A schematic of the model architecture and itsinference is illustrated in Fig.
Importantly, we require that for the dynamic encoder q , it needsto have a ﬂatten operation and MultiLayer Perception (MLP) head that maps the feature map intoa single ﬁxed-length vector z 2 R d z .
In this way, the dimension of the latent space does not scalelinearly with the dimension of the input, which has the potential to signiﬁcantly compress the input,and can make the long-term prediction much more efﬁcient.
At inference time, LE-PDE performsautoregressive rollout in latent space R d z :
[Equation]
Compared to autoregressive rollout in input space (Eq. 3), LE-PDE can signiﬁcantly improveefﬁciency with a much smaller dimension of z k 2 R d z compared to U k 2 U .
Here we do not limitthe architecture for encoder, decoder and latent evolution models.
Depending on the input U k , theencoder q and decoder h can be a CNN or GNN with a (required) MLP head.
In this work, we focuson input that is discretized as grid, so the encoder and decoder are both CNN+MLP, and leave otherarchitecture (e.g. GNN+MLP) for future work.
For static encoder r , it can be a simple MLP if thesystem parameter p is a vector (e.g. equation parameters) or CNN+MLP if p is a 2D or 3D tensor(e.g. boundary mask, spatially varying diffusion coefﬁcient).
We model the latent evolution model gas an MLP with residual connection from input to output.
The architectures used in our experiments,are detailed in Appendix C, together with discussion of its current limitations.
3.2 Learning objective
Learning surrogate models that can faithfully roll out long-term is an important challenge.
Givendiscretized inputs { U k } , k = 1 , ...
K + M , we introduce the following objective to address it:
[Equation]
where
[Equation]
Here ` is the loss function for individual predictions, which can typically be MSE or L2 loss.
ˆ U k
+ mis given in Eq.
L k recons aims to reduce reconstruction loss.
L k multi-step performs latent multi-stepevolution given in Eq.
(4) and compare with the target U k
+ m in input space, up to time horizon
M .
↵ m are weights for each time step, which we ﬁnd that ( ↵ 1 , ↵ 2 , ... ↵ M ) =
1)works well.
Besides encouraging better prediction in input space via L k multi-step , we also want astable long-term rollout in latent space.
This is because in inference time, we want to mainlyperform autoregressive rollout in latent space, and decode to input space only when needed.
Thus,we introduce a novel latent consistency loss L k consistency , which compares the m -step latent rolloutg ( · , r ( p )) ( m )  q ( U k ) with the latent target q
( U k + m ) in latent space.
The denominator || q
( U k + m )
||
22serves as normalization to prevent the trivial solution that the latent space collapses to a single point.
Taken together, the three terms encourage a more accurate and consistent long-term evolution both inlatent and input space.
3.3 Accelerating inverse optimization
In addition to improved efﬁciency for forward simulation, LE-PDE also allows more efﬁcient inverseoptimization, via backpropagation through time (BPTT) in latent space.
Given a speciﬁed objective
L d [ p ] =
P k e k
=
k s ` ( U k ) which is a discretized version of L d [ a , @
X ] in Sec. 2, we deﬁne the objective:
[Equation]
where ˆ U m =
ˆ U m ( p ) is given by Eq.
(4) setting k = 0 using our learned LE-PDE, which starts atinitial state of U 0 , encode it and p into latent space, evolves the dynamics in latent space and decodeto ˆ U m as needed.
The static latent embedding z p = r ( p ) inﬂuences the latent evolution at each timestep via g ( · , r ( p )) .
An approximately optimal parameter p can then be found by computing gradients@ L d [ p ] @ p , using optimizers such as Adam [58] (The gradient ﬂow is visualized as the red arrows in Fig.1).
When p is a boundary parameter, e.g. location of the boundary segments or obstacles, there isa challenge.
Speciﬁcally, for CNN encoder q , the boundary information is typically provided as abinary mask indicating which cells are outside the simulation domain ⌦ .
This prevents the gradient topass through to the boundary parameter p such as continuous location.
Moreover, the boundary cellsmay 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, andduring inverse optimization, anneal the temperature from high to low.
This allows the gradient topass through mask to p , and stronger gradient signal.
For more information, see Appendix B.
Experiments
In the experiments, we aim to answer the following questions: (1) Does LE-PDE able to learnaccurately the long-term evolution of challenging systems, and compare competitively with state-ofthe-art methods?
(2) How much can LE-PDE reduce representation dimension and improving speed,especially with larger systems?
(3) Can LE-PDE improve and speed up inverse optimization?
Forthe ﬁrst and second question, since in general there is a fundamental tradeoff between compression(reduction of dimensions to represent a state) and accuracy [59, 60], i.e. the larger the compressionto improve speed, the more lossy the representation is, we will need to sacriﬁce certain amount ofaccuracy.
Therefore, the goal of LE-PDE is to maintain a reasonable or competitive accuracy (maybeslightly underperform state-of-the-art), while achieving signiﬁcant compression and speed up.
Thus,to answer these two questions, we test LE-PDE in standard benchmarks of a 1D family of nonlinear
PDEs to test its generalization to new system parameters (Sec. 4.1), and a 2D Navier-Stokes ﬂow up
[Table]
4 ⇥ 10     to turbulent phase (Sec. 4.2).
The PDEs in the above scenarios have wide and important applicationin science and engineering.
In each domain, we compare LE-PDE’s long-term evolution performance,speed and representation dimension with state-of-the-art deep learning-based surrogate models inthe domain.
Then we answer question (3) in Section 4.3.
Finally, in Section 4.4, we investigate theimpact 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 equationswith different parameters in a given family.
We use the 1D benchmark in [7], whose PDEs are
[Equation]
Here the parameter p
= ( ↵ ,  ,  ) .
The term  is a forcing term [61] with J = 5 , L = 16 andcoefﬁcients A j and !
j sampled uniformly from A j ⇠ U [  0 . 5 , 0 . 5] , !
j ⇠ U [  0 . 4 , 0 . 4] , ` j 2{ 1 , 2 , 3 } ,  j ⇠ U [0 , 2 ⇡ ) .
Space is uniformly discretized to n x = 200 in [0 , 16) and time is uniformlydiscretized to n t = 200 points in [0 , 4] .
Space and time are further downsampled to resolutions of( n t , n x ) 2 { (250 , 100) , (250 , 50) , (250 , 40) } .
The @ x ( ↵ u 2 ) advection term makes the PDE nonlinear.
There are 3 scenarios with increasing difﬁculty: E1:
Burgurs’ equation without diffusion p = (1 , 0 , 0) ;
E2 :
Burgurs’ equation with variable diffusion p = (1 , ⌘ , 0) where ⌘ 2 [0 , 0 , 2] ; E3 : mixed scenariowith p = ( ↵ ,  ,  ) where ↵ 2 [0 , 3] ,  2 [0 , 0 . 4] and  2 [0 , 1] .
E1 tests the model’s ability togeneralize to new conditions with same equation.
E2 and E test the model’s ability to generalizeto novel parameters of PDE with the same family.
We compare LE-PDE with state-of-the-art deep learning-based surrogate models for this dataset, speciﬁcally MP-PDE [ (a GNN-based model )and Fourier Neural Operators ( FNO)
For FNO, we compare with two versions
: FNO-RNN isthe autoregressive version in Section 3 of their paper trained with autoregressive rollout ; FNO-PF is FNO improved with the temporal bundling and push-forward trick as implemented in [7]
Toensure 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 , and use the model with best validation performance.
In addition, we compare with downsampled ground
-truth(WEBO5), which uses a classical 5 th -order WEBO scheme [62] and explicit Runge Kutta 4 solver[63, 64] to generate the ground-truth data and downsampled to the speciﬁed resolution.
For allmodels we autoregressively roll out to predict the states starting at step 50 until step and recordthe accumulated MSE, runtime and representation dimension (the dimension of state to update ateach time step).
Details of the experiments are given in Appendix D.
Results .
The result is shown in Table 1.
We see that since LE-PDE uses 7.8 to 39-fold smallerrepresentation dimension, it achieves signiﬁcant smaller runtime compared to the GNN-based MPPDE model (which is already much faster than the classical WEBO5 scheme).
Here we record thelatent evolution time (LE-PDE evo) which is the total time for 200-step latent evolution, and the fulltime (LE-PDE 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-termfuture (e.g. [4]); the time for LE-PDE full is relevant when every intermediate prediction is alsoimportant.
LE-PDE achieves up to 15 ⇥ speedup with “LE-PDE evo” and 4 ⇥ speedup with “LE-PDEfull”.
[Figure]
With above 7.8 ⇥ compression and above 4 ⇥ speedup, LE-PDE still achieves competitive accuracy.
For E1 scenario, it signiﬁcantly outperforms both original versions of FNO-RNN and MP-PDE, andonly worse than the improved version of FNO-PF.
For E3 , 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-art models by a large margin.
Fig.
4 in Appendix D shows ourmodel’s representative rollout compared to ground-truth.
We see that during long-rollout, our modelcaptures the shock formation faithfully.
This 1D benchmark shows that LE-PDE is able to achievesigniﬁcant speedup, generalize to novel PDE parameters and achieve competitive long-term rollout.
4.2 2D Navier-Stokes ﬂow
Data and Experiments .
We test LE-PDE in a 2D benchmark [14] of Navier-Stokes equation.
Navier-Stokes equation has wide application science and engineering, 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.
Speciﬁcally, we test our model in a viscous,incompressible ﬂuid in vorticity form in a unit torus:
[Equation]
Here w ( t, x ) =
r ⇥ u ( t, x ) is the vorticity, ⌫ 2 R + is the viscosity coefﬁcient.
The domain isdiscretized into 64 ⇥ 64 grid.
We test with viscosities of ⌫ = 10  3 , 10  4 , 10  5 .
The ﬂuid is turbulentfor ⌫ = 10  4 , 10  5 ( Re  10 4 ).
We compare state-of-the-art learning-based model Fourier Neural
Operator (FNO) [14] for this problem, and strong baselines of TF-Net [24], U-Net [65] and ResNet[66].
For FNO, the FNO-2D performs autoregressive rollout, and FNO-3D directly maps the past 10steps into all future steps.
To ensure a fair comparison, here our LE-PDE uses past 10 steps to predictone future step and temporal bundling S
= 1 (no bundling), the same setting as in FNO-2D. We userelative L2 norm (normalized by ground-truth’s L2 norm) as metric, same as in [14].
Results .
, The results are shown in Table 2.
Similar to 1D case, LE-PDE is able to compress therepresentation dimension by 16-fold.
Hence, compared with FNO-2D which is also autoregressive,
LE-PDE achieves 17-fold speedup with latent evolution and 5.8-fold speedup with full decoding.
Compared with FNO-3D that directly maps all input time steps to all output times steps (which cannotgeneralize beyond the time steps given), LE-PDE’s runtime is still on par (for full decoding) or 2.9 ⇥
[Table]
[Table]
[Figure]
faster. For rollout L2 loss, LE-PDE signiﬁcantly outperforms strong baselines of ResNet and U-Net,and TF-Net which is designed to model turbulent ﬂow.
Its performance is on par with FNO-3Dwith ⌫ = 10  4 , N = 1000 and the most difﬁcult ⌫ = 10  5 , N = 1000 and slightly underperforms
FNO-2D in other scenarios.
Fig.
2 shows the visualization of LE-PDE comparing with ground-truth,under the turbulent ⌫ = 10  4 , N = 10000 scenario.
We see that LE-PDE captures the detaileddynamics accurately.
For more details, see Appendix E.
To explore how LE-PDE can model andaccelerate the simulation of systems with a larger scale, in Appendix F we explore modeling a 3D
Navier-Stokes ﬂow with millions of cells per time step, and show more signiﬁcant speedup.
4.3 Accelerating inverse optimization of boundary conditions
[Table]
and much lower error computed by ground-truth solver.
[Table]
Data and Experiments .
In this subsection, weset out to answer question (3), i.e. Can LEPDE improve and speed up inverse optimization?
We are interested in long time frame scenarios where the pre-deﬁned objective L d in Eq.(7) depends on the prediction after long-termrollout.
Such problems are challenging and haveimplications in engineering, e.g. ﬂuid control[67, 68], laser design for laser-plasma interaction [4] and nuclear fusion [69].
To evaluate, we build a 2D Navier-Stokes ﬂow in a family ofboundary conditions using PhiFlow [70] as our ground-truth solver, shown in Fig.
3a, 3b.
Speciﬁcally,we create a cubical boundary with one inlet and two outlets on a grid space of size 128 2 .
We initializethe velocity and smoke on this domain and advect the dynamics by performing rollout.
The objectiveof the inverse design here is to optimize the boundary parameter p , i.e. the y -locations of the inletand outlets, so that the amount of smoke passing through the two outlets coincides with pre-speciﬁedproportions 0 .
3 and 0 .
This setting is challenging since a slight change in boundary (up to a fewcells) can have large inﬂuence in long-term rollout and the predeﬁned 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 [14], both of which update the states in input space, while LE-PDE evolves in a 128 -dimensional latent space.
To ensure a fair comparison, allmodels predict the next step using 1 past step without temporal bundling, and trained with 4-steprollout.
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 p with the trainedmodels using Eq. 7, starting with 50 initial conﬁgurations each with random initial location of smokeand random initial conﬁguration of p .
For LE-PDE- ⇠⇠⇠ latent and FNO-2D, they need to backpropagatethrough 80 steps of rollout in input space as in [43], while our LE-PDE backpropagates through 80steps of latent rollout.
Then the optimized boundary parameter is fed into the ground-truth solver forrollout and evaluate.
The evaluation metric is average amount of the advected smoke simulated bythe solver across all 50 conﬁgurations. For more details see Appendix G.
Results .
We observe that LE-PDE improves the overall speed by 73% compared with LE-PDE- ⇠⇠⇠ latentand by 123% compared with FNO-2D (Fig.
3c, Table 3).
The result indicates a corollary of the useof low dimensional representation because Jacobian matrix of evolution operator is reduced to be ofsmaller size and suppresses the complexity associated with the chain rule to compute gradients of theobjective function.
While achieving the signiﬁcant speedup, the capability of the LE-PDE to designthe boundary is also reasonable.
Fig.
3d shows the loss of the objective function achieved the lowestvalue while the others are comparably large.
The estimated proportion of smoke hit the target fraction0 .
3 at an early stage of design iterations and coincide with the fraction simulated by the ground-truthsolver in the end (Fig. 3e).
We note that, as Table 3 shows, FNO-2D achieved the lowest score inmodel estimated error from the target fraction 0.3 while its ground-truth solver (GT-solver) erroris 30 ⇥ larger.
This shows “overﬁtting” of the boundary parameter by FNO-2D, i.e. the optimizedparameter is not sufﬁciently generalized to work for a ground-truth solver.
In this sense, LE-PDEachieved to design the most generalized boundary parameter: the difference between the two errors isthe smallest among the others.
4.4 Ablation study
[Table]
In this section, we investigate how each component ofour LE-PDE inﬂuences the performance.
Importantly, weare interested in how the multi-step loss L multi-step , latentconsistency loss L consistency contribute to the performance,and how the time horizon M and the latent dimension d zinﬂuence the result.
We focus on representative scenariosin 1D (Sec. 4.1) and 2D (Sec. 4.2), speciﬁcally the E2scenario with ( n t , n x ) =
(250 , 50) for 1D, and ( ⌫ =10  5 , T = 20 , N = 1000) scenario for 2D, which liesat mid- to difﬁcult spectrum of each dataset.
We haveobserved similar trend in other scenarios.
From Table 4,we see that both L multi-step and L consistency 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 ), it does not encourageaccurate long-term rollout.
Increasing M helps reducing error, but will be countered by less numberof examples (since having to leave room for more steps in the future).
We ﬁnd the sweep spot isat M = 4 , which achieves a good tradeoff.
In Fig.
6 in Appendix H, we show how the error andevolution runtime change with varying size of latent dimension d z .
We observe that reduction ofruntime with decreasing latent dimension d z , and that the error is lowest at d z = 64 for 1D andd z = 256 for 2D, suggesting the intrinsic dimension of each problem.
Discussion and Conclusion
In this work, we have introduced LE-PDE, a simple, fast and scalable method for acceleratingsimulation and inverse optimization of PDEs, including its simple architecture, objective and inverseoptimization techniques.
Compared with state-of-the-art deep learning-based surrogate models,we demonstrate that it achieves up to 128 ⇥ reduction in the dimensions to update and up to 15 ⇥improvement in speed, while achieving competitive accuracy.
Ablation study shows both multi-stepobjective and latent-consistency objectives are pivotal in ensuring accurate long-term rollout.
Wehope our method will make a useful step in accelerating simulation and inverse optimization of PDEs,pivotal in science and engineering.
Checklist
The checklist follows the references.
Please read the checklist guidelines carefully for information onhow to answer these questions.
For each question, change the default [TODO] to [Yes] , [No] , or[N/A] .
You are strongly encouraged to include a justiﬁcation to your answer , either by referencingthe appropriate section of your paper or providing a brief inline description.
For example:
• Did you include the license to the code and datasets?
[Yes] See Section 4.
• Did you include the license to the code and datasets?
[No] The code and the data areproprietary.
• Did you include the license to the code and datasets?
[N/A]
Please do not modify the questions and only use the provided macros for your answers.
Note that the
Checklist section does not count towards the page limit.
In your paper, please delete this instructionsblock and only keep the Checklist section heading above along with the questions/answers below.
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s
contributions and scope?
[Yes]
(b) Did you describe the limitations of your work?
[Yes] In Appendix C.
(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 to
them?
[Yes
If you are 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]
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi-
mental results (either in 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
to reproduce the experiments are included in the Appendix.
(c) Did you report error bars (e.g., with respect to the random seed after running ex-
periments multiple times)?
[Yes] Yes in Section 4.4.
For Sections 4.1 and 4.2, the
benchmarks do not include the error bars and we use the same format.
(d) Did you include the total amount of compute and the type of resources used (e g., type
of GPUs, internal cluster, or cloud provider)?
[Yes] In Appendix D,E,F,G.
If you are using existing assets (e g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[Yes] In Appendix D,E,F,G.
(b) Did you mention the license of the assets?
[Yes] In Appendix D,E,F,G.
(c) Did you include any new assets either in the supplemental material or as a URL?
[Yes]
In Appendix D,E,F,G.
(d) Did you discuss whether and how consent was obtained from people whose data you’re
using/curating?
[Yes] In Appendix D,E,F,G.
(e) Did you discuss whether the data you are using/curating contains personally identiﬁable
information or offensive content?
[Yes] In Appendix D,E,F,G.
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if
applicable?
[N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable?
[N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount
spent on participant compensation?
[N/A]