What is the problem we aimed to solve?

But first, what is differentiable physics, and why should I care?

Differentiable physics refers to physics solvers written in frameworks that are automatic differentiation (autodiff) friendly, i.e., using autodiff, we can easily compute Jacobians of these solvers. For the uninitiated - autodiff is the algorithm of gradient computation underlying most modern gradient-based optimization, especially neural network training.

Then the big idea is - through the gradient, the learning pipeline "interacts" with the physics of the problem. This interaction allows machine learning models that incorporate physical laws (rather than pure data-driven learning).

Differentiable physics is but one of the several methods enabling physics-aware models. Some examples where it has been successful:

Solving inverse problems (Bendsoe & Sigmund, 2013)
Integrating physical constraints (Raissi et al., 2019; Li et al., 2024
Creating hybrid models that blend classical numerical techniques with learned components (Um et al., 2020; Kochkov et al., 2021; 2024)

Differentiable Physics sounds great! What's the problem?

Physics solvers are typically iterative. Consider a heat diffusion equation, discretized in space using the finite volumes method. This forms a sparse linear system solved using an iterative solver e.g. Jacobi, GMRES, etc.
Repeatedly querying a solver with several iterations in the forward pass - and differentiating through its iterations in the backward pass - can introduces a severe computational bottleneck during training. See figure 1.

Figure 1: A typical neural correction-learning pipeline (Um et al., 2020) that uses a differentiable physics solver $\mathcal{P}_K$ in the loop. Black arrows show the forward pass, grey dashed arrows represent the backward pass, and elements in red represent the bottleneck. As the number of solver iterations $K$ grows, the cost of passes through $\mathcal{P}_K$ becomes severe.

That's a problem! So what is PRDP?

We often train models using differentiable physics at our research group, and it's costly! So, in the summer of 2024, we set out to devise a method that makes this training cheaper (hence, faster) but doesn't need fancy (read: difficult) method implementations from the user. Enter: Progressively Refined Differentiable Physics (PRDP).

PRDP works by letting the learning algorithm use cheap low-accuracy runs of the physics solver in the beginning, and adaptively refines the physics as the epochs progress (1). The method also prevents superfluous refinement of the physics solver, based on our observation that full convergence of the physics is not necessary - neural emulators, even when trained through incompletely (but sufficiently) converged physics, perform as well as their counterparts trained through fully-converged physics (2).
These two ideas - (1) and (2) - are henceforth referred to as Progressive Refinement (PR) and Incomplete Convergence (IC).

Figure 2: PRDP reduces the training time of neural networks containing numerical solver components (c). The fidelity of iterative components is increased only if validation metrics of the network training plateau. This leads to savings by using fewer iterations in the beginning (PR savings in (b)) and by ending at a refinement level significantly below full fidelity (IC savings in (b)). The achieved validation error is identical (a).

Sweet! How does PRDP work?

PRDP's core idea revolves around balancing the compute-accuracy trade-off of iterative physics solvers. The nuancy is - we assert only the accuracy of the neural network being trained, not the physics solver. Hence, the physics solver is not converged to numerical precision, but is allowed to run only for about as many iterations as sufficient to gain significant training progress in the neural network.

How much physics refinement is sufficient?
We contribute an algorithm that determines physics refinement adaptively based on the plateauing of training progress measured on a validation metric.

Training begins with cheap coarse physics.
When the training progress stagnates, a refinement of the physics is invoked.
If the training progress has stangated over successive physics refinements, indicating that further physics refinement would be superfluous, refinement is stopped.

Figure 3: Left: the typical training progress of a neural network supported by PRDP. Right: a simplified flowchart representation of the PRDP control algorithm.

This throttling of the physics solver saves significant compute, especially in cases where physics represents a high proportion of the training cost. For instance, in our test problem of training a neural emulator on a 3D heat equation solver, PRDP saves about 78% in total training time.

I'm sold. How can I use PRDP?

Be sure to check out the code examples in our GitHub repository
Using PRDP involves integrating it into your existing machine learning pipeline that includes a physics solver. Obtain the file prdp.py from the code repo. The should_refine function implements the PRDP control algorithm and needs to be imported in your code. Here are the steps to get started:

Identify the Physics Solver $\mathcal{P}$ in your learning pipeline. Typically Physics solvers have at their core an iterative fixed-point algorithm $\Phi$. This is the one who's convergence we throttle using PRDP during the learning process.
Decide the parameter that determines solver refinement: number of iterations $K$, or tolerance (the arguments called abstol and reltol in most solver implementations). We use number of iterations in this demonstration.
Record Validation metric: Create an array that records the validation loss (or training loss) of the last $\delta$ epochs. While training loss can be used, we recommend using the validation loss. Refer appendix G.4 of the paper for the explanation.
Monitor Convergence: Implement a mechanism to monitor the convergence of the physics solver and stop refinement when sufficient accuracy is achieved for training.
Integrate with Training Loop: At the end of the first update step, using the should_refine function, determine whether the physics is to be refined. If yes, refine the physics by incrementing the number of solver iterations $K$. Refer the code below. .
Experiment and Tune: The parameters stepping_threshold and nmax_threshold control the speed of progressive refinement and the strictness on incomplete convergence, respectively. A lower value of stepping_threshold enables faster refinement. This leads to potentially lower savings, but can also give faster training progress. While the default parameter values are expected to work reasonably well in most cases, the user is encouraged to experiment with different values.

Abstract

The physics solvers employed for neural network training are primarily iterative, and hence, differentiating through them introduces a severe computational burden as iterations grow large. Inspired by works in bilevel optimization, we show that full accuracy of the network is achievable through physics significantly coarser than fully converged solvers. We propose Progressively Refined Differentiable Physics (PRDP), an approach that identifies the level of physics refinement sufficient for full training accuracy. By beginning with coarse physics, adaptively refining it during training, and stopping refinement at the level adequate for training, it enables significant compute savings without sacrificing network accuracy. Our focus is on differentiating iterative linear solvers for sparsely discretized differential operators, which are fundamental to scientific computing. PRDP is applicable to both unrolled and implicit differentiation. We validate its performance on a variety of learning scenarios involving differentiable physics solvers such as inverse problems, autoregressive neural emulators, and correction-based neural- hybrid solvers. In the challenging example of emulating the Navier-Stokes equations, we reduce training time by 62%.

BibTeX


@article{bhatia2025prdp,
    title={Progressively Refined Differentiable Physics},
    author={Kanishk Bhatia and Felix Koehler and Nils Thuerey},
    journal={International Conference on Learning Representations (ICLR)},
    volume={13},
    year={2025}
}

PRDP: Progressively Refined Differentiable Physics

ICLR 2025