Learning Controllable Adaptive Simulation for Multi-resolution Physics

1

Tailin Wu*（吴泰霖）, Takashi Maruyama*, Qingqing Zhao*, Gordon Wetzstein, Jure Leskovec

Stanford University

ICLR 2023, notable-top-25%

*equal contribution

Outline

2

1. Background: multi-resolution physics simulation and prior methods
2. Method: Learning Controllable Adaptive Simulation for Multi-resolution Physics (LAMP)
3. Experiments

Simulation is a core task in science

3

Problem definition and significance

4

Problem definition: Simulate a multi-resolution system in a accurate and efficient way.

Significance: Many physical systems in science and engineering are multi-resolution: parts of the system is highly dynamic and need to resolve to fine-grained resolution, while other parts are more static.

Weather prediction

Galaxy formation

Laser-plasma particle acceleration

Preliminaries: classical solvers and limitations

5

Classical solvers:

Based on Partial Differential Equations (PDEs)

• Pros: (1) Based on first principles and interpretable, (2) accurate, (3) have error guarantee.
• Challenges: Slow and computational expensive, due to

(1) Small time interval to ensure numerical stability, or use implicit method.

(2) For multi-resolution systems, typically need to resolve to the lowest resolution

Pros and challenges:

Preliminaries: Deep learning-based surrogate models

6

Recently, deep learning based surrogate modeling has emerged as attractive alternative to replace or complement classical solvers. They:

• Offer speedup via:
• Larger spatial resolution
• Larger time intervals
• Use explicit forward

Limitation of today’s method for multi-resolution physics:

1. Deep learning-based surrogate models: prior works mainly focus on learning the evolution with low prediction error, without directly optimizing the computational cost

• Almost all works are learning the evolution model on a fixed grid or mesh
• A few exceptions, e.g. MeshGraphNet [1] learn the remeshing using supervision from the solver, where the solver’s label is heuristics based on geometric feature (second derivative) of mesh coordinates.

7

[1] Pfaff, Tobias, et al. "Learning mesh-based simulation with graph networks." arXiv preprint arXiv:2010.03409 (2020).

Architectural backbone: MeshGraphNets [1]

8

[1] Pfaff et al. ICLR 2021

Limitation of today’s method for multi-resolution physics:

9

2. Classical solvers:

• Many solvers are based on fixed grid/mesh
• Adaptive Mesh Refinement (AMR) can adaptively update the local spatial resolution. Since it is based on classical solvers, it shares similar drawbacks (e.g. slow), and it is based on heuristics without directly optimizing the computational cost

Our contribution

10

We introduced the first deep learning-based surrogate model, which jointly learns the evolution and optimize computational cost.

(Wu et al., ICLR 2023, spotlight)

Key component is a GNN-based reinforcement learning (RL) agent, which learns to coarsen or refine the mesh, to achieve a controllable tradeoff between prediction error and computational cost.

Architecture

11

Architecture

12

Architecture

13

Architecture

14

Action space: refinement and coarsening

15

(1) Refining an edge

There are also constraints that need to be satisfied, e.g. if two edges are on the same face, they cannot be both refined or coarsened.

Learning the evolution model

The loss is based on the multi-step prediction error of the evolution model, compared with the ground-truth.

16

Learning the policy

17

for S steps

Learning the policy : reward

18

Reward:

Reward is based on the improvement of both error and computational cost.

• Error is the multi-step prediction error
• Computational cost is measured by number of vertices in the mesh

[1] Sutton, et al. NIPS 1999

is also an input to the policy

for S steps

for S steps

Learning the policy : actor objective

19

advantage

log prob. for taking

the action

entropy regularizer

sg: stop gradient

REINFORCE

Learning the policy : critic objective

20

Experiment 1

(1) Burgers’ equation (from the benchmark in [1])

21

[1] Brandstetter, Johannes, Daniel Worrall, and Max Welling. "Message passing neural PDE solvers." ICLR 2022

Experiment 1

Example rollout:

22

• Refines more near the shock front, and coarsens more in static regions
• With larger that focuses more on reducing computation, it refines less and coarsens more

Added cells

removed cells

Experiment 1

Example rollout:

23

• Refines more near the shock front, and coarsens more in static regions
• With larger that focuses more on reducing computation, it refines less and coarsens more

24

25

Result table:

Compared to state-of-the-art (FNO, MP-PDE) models and strong baselines, our model achieves large error reduction (average of 33.4%)

26

Result table:

Compared to ablation LAMP (no remeshing), our full model can reduce error (49.1%), by only modest increase computational cost.

Experiment 1

27

Experiment 1

28

Experiment 2: results

29

Experiment 2: example visualization

30

ground-truth (fine-grained)

MeshGraphNets + GT remeshing

MSE: 5.91e-4

Experiment 2: example visualization

31

LAMP + heuristic remeshing

ground-truth (fine-grained)

MSE: 6.38e-4

Experiment 2: example visualization

32

LAMP + no remeshing

ground-truth (fine-grained)

MSE: 6.13e-4

Experiment 2: example visualization

33

ground-truth (fine-grained)

34

Other examples:

Summary

1. Multi-resolution is a key characteristics in physical simulations.
1. Challenge: how simulate accurately and efficiently
2. We introduced LAMP, the first deep learning-based surrogate model (LAMP) which jointly learns the evolution and optimizes:

35

Summary

3. Experiments in 1D PDE and mesh-based simulation demonstrate LAMP’s capability,

outperform previous state-of-the-art.

​

36

Paper:

Code:

Future opportunities

37

1. Larger simulations

2. Inverse design

Up to millions to billions of nodes

Welcome collaborations (email Tailin Wu, tailin@cs.stanford.edu)!