[Paper Review] Learning Controllable Adaptive Simulation for Multi-resolution Physics

Tailin Wu, Takashi Maruyama+3 more|arXiv (Cornell University)|May 1, 2023

Scientific Computing and Data Management4 citations

TL;DR

LAMP is a deep learning-based surrogate model that jointly learns physical system evolution and adaptive spatial refinement/coarsening policies using graph neural networks and reinforcement learning. It reduces long-term prediction error by 33.7% in 1D nonlinear PDEs and outperforms state-of-the-art methods in 2D mesh-based simulations by dynamically allocating computation to high-dynamic regions while balancing error and computational cost.

ABSTRACT

Simulating the time evolution of physical systems is pivotal in many scientific and engineering problems. An open challenge in simulating such systems is their multi-resolution dynamics: a small fraction of the system is extremely dynamic, and requires very fine-grained resolution, while a majority of the system is changing slowly and can be modeled by coarser spatial scales. Typical learning-based surrogate models use a uniform spatial scale, which needs to resolve to the finest required scale and can waste a huge compute to achieve required accuracy. In this work, we introduce Learning controllable Adaptive simulation for Multi-resolution Physics (LAMP) as the first full deep learning-based surrogate model that jointly learns the evolution model and optimizes appropriate spatial resolutions that devote more compute to the highly dynamic regions. LAMP consists of a Graph Neural Network (GNN) for learning the forward evolution, and a GNN-based actor-critic for learning the policy of spatial refinement and coarsening. We introduce learning techniques that optimizes LAMP with weighted sum of error and computational cost as objective, allowing LAMP to adapt to varying relative importance of error vs. computation tradeoff at inference time. We evaluate our method in a 1D benchmark of nonlinear PDEs and a challenging 2D mesh-based simulation. We demonstrate that our LAMP outperforms state-of-the-art deep learning surrogate models, and can adaptively trade-off computation to improve long-term prediction error: it achieves an average of 33.7% error reduction for 1D nonlinear PDEs, and outperforms MeshGraphNets + classical Adaptive Mesh Refinement (AMR) in 2D mesh-based simulations. Project website with data and code can be found at: http://snap.stanford.edu/lamp.

Motivation & Objective

• To address the challenge of simulating multi-resolution physical systems where only a small fraction is highly dynamic, requiring fine resolution, while most regions evolve slowly.
• To overcome the inefficiency of uniform-resolution deep learning surrogate models that waste compute on slow regions and may under-resolve dynamic ones.
• To develop a fully differentiable, end-to-end deep learning framework that jointly optimizes evolution modeling and spatial resolution adaptation.
• To enable dynamic trade-offs between prediction error and computational cost at inference time via a learnable policy.
• To demonstrate superior performance over state-of-the-art surrogate models and classical adaptive mesh refinement in long-term simulation accuracy.

Proposed method

• LAMP uses a Graph Neural Network (GNN) to learn the forward evolution of physical systems on dynamic meshes.
• It employs a GNN-based actor-critic reinforcement learning framework to learn discrete policies for local mesh refinement and coarsening.
• The policy is conditioned on local state features and a hyperparameter β that controls the trade-off between error and computational cost.
• The system is trained via alternating optimization: supervised learning for the evolution model and reinforcement learning for the policy.
• The reward function is defined as a weighted sum of error reduction and computational cost reduction over multi-step rollouts.
• Remeshing is performed by selecting edges to refine or coarsen based on the policy output, followed by interpolation to maintain consistency across time steps.

Figure 1: LAMP schematic. The forward iteration (upper box) first uses the policy $f_{\varphi}^{\text{policy}}$ to decide the number $K^{\text{re}}$ and $K^{\text{co}}$ of edges as well as which edges among the full mesh to be refined or coarsened, and then executes remeshing and interpolation. The

Experimental results

Research questions

• RQ1 Can a deep learning-based surrogate model jointly learn physical system evolution and adaptive spatial resolution selection?
• RQ2 Can such a model dynamically allocate computational resources to highly dynamic regions while reducing long-term prediction error?
• RQ3 How does the performance of a learned remeshing policy compare to heuristic or ground-truth remeshing in complex simulations?
• RQ4 To what extent can the model adapt to varying trade-offs between accuracy and computational cost at inference time?
• RQ5 Can the model generalize across different physical systems within the same family of PDEs?

Key findings

• LAMP reduces long-term prediction error by 33.7% compared to state-of-the-art deep learning surrogate models on a 1D nonlinear PDE benchmark.
• In 2D paper folding simulations, LAMP achieves a lower error (5.80e-4 MSE) than both the MeshGraphNets + heuristic remeshing baseline (6.38e-4 MSE) and the no-remeshing ablation (6.13e-4 MSE).
• LAMP outperforms the strong baseline of MeshGraphNets with ground-truth remeshing, achieving better accuracy despite using a learned remeshing policy.
• The model successfully learns to refine high-curvature regions (e.g., folding parts) and coarsen flat, slow-changing regions, as visualized in 2D simulations.
• LAMP maintains a low average number of vertices (123.1) while achieving high accuracy, indicating efficient computation use.
• The model demonstrates robust adaptability to different β values at inference time, enabling dynamic trade-offs between error and computational cost.

Figure 2: Example rollout result of our LAMP on 1D nonlinear PDEs. The rollout is performed over 200 time steps, where different color denotes the system’s state at different time. On each state, we also plot the corresponding mesh as black dots. The upper green and lower blue band shows the added a