[Paper Review] Solver-in-the-Loop: Learning from Differentiable Physics to Interact with Iterative PDE-Solvers

Kiwon Um, Robert Brand+4 more|arXiv (Cornell University)|Jun 30, 2020

Model Reduction and Neural Networks76 references117 citations

TL;DR

The paper trains neural corrections that interact with iterative PDE solvers in a differentiable framework, markedly reducing numerical errors across multiple PDE scenarios and outperforming non-interacting or pre-computed correction approaches.

ABSTRACT

Finding accurate solutions to partial differential equations (PDEs) is a crucial task in all scientific and engineering disciplines. It has recently been shown that machine learning methods can improve the solution accuracy by correcting for effects not captured by the discretized PDE. We target the problem of reducing numerical errors of iterative PDE solvers and compare different learning approaches for finding complex correction functions. We find that previously used learning approaches are significantly outperformed by methods that integrate the solver into the training loop and thereby allow the model to interact with the PDE during training. This provides the model with realistic input distributions that take previous corrections into account, yielding improvements in accuracy with stable rollouts of several hundred recurrent evaluation steps and surpassing even tailored supervised variants. We highlight the performance of the differentiable physics networks for a wide variety of PDEs, from non-linear advection-diffusion systems to three-dimensional Navier-Stokes flows.

Motivation & Objective

• Motivate reducing discretization and numerical errors in iterative PDE solvers.
• Propose a correction function learned by neural networks that can interact with the solver.
• Compare three training regimes (non-interacting, pre-computed, solver-in-the-loop) for accuracy and stability.
• Demonstrate gains across a suite of PDEs from 2D/3D flows to Poisson problems.
• Evaluate how look-ahead rollout during training impacts long-term accuracy and stability.

Proposed method

• Model a correction function C(s|θ) as a neural network that adds a correction to solver state s.
• Embed the neural correction inside a differentiable PDE solver to enable end-to-end training (solver-in-the-loop).
• Contrast NON (no interaction), PRE (pre-computed interaction), and SOL (solver-in-the-loop) training regimes.
• Use a fully convolutional network (10 layers, 16 features) trained with ADAM (lr=1e-4).
• Assess performance by mean absolute error between corrected trajectory and reference over n steps, with look-ahead n (SOL n).
• Apply to advection-diffusion, 2D/3D Navier–Stokes, buoyancy-driven flows, and Poisson-related CG initialization tasks.

Experimental results

Research questions

• RQ1 Can neural corrections learned inside a differentiable PDE solver outperform traditional supervised or pre-computed corrections for iterative solvers?
• RQ2 How does solver-in-the-loop training with varying look-ahead horizons affect long-term accuracy and stability across different PDEs?
• RQ3 What are the comparative benefits and limitations of NON, PRE, and SOL interaction modes in reducing numerical errors?
• RQ4 Do differentiable-physics-trained corrections generalize to out-of-distribution initial conditions and higher-dimensional flows?

Key findings

• Solver-in-the-loop corrections yield substantial accuracy gains over non-interacting and pre-computed methods.
• SOL models with longer look-ahead horizons dramatically reduce errors (e.g., up to ~60% relative improvement in some cases).
• Differentiable physics training yields improvements beyond supervised or pre-computed corrections, with stable long-horizon rollouts.
• 3D wake flows show >22% improvement in numerical accuracy using SOL corrections.
• Training cost increases but inference remains unchanged, and speedups are achieved in end-to-end simulations (e.g., 68x faster than CPU reference in one scenario).
• Corrections generalize to out-of-distribution initial conditions and improve stability for recurrent PDE solves.