PRDP: Progressively Refined Differentiable Physics
Authors: Kanishk Bhatia
Felix Koehler
Nils Thuerey
Venue: ICLR 2025 Poster (2025)
Keywords: differentiable physics
iterative PDE solvers
neural surrogate
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)

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%.

Published as a conference paper at ICLR 2025
PRDP: PROGRESSIVELY REFINED DIFFERENTIABLE PHYSICS
Kanishk Bhatia,, Felix Kochler', Nils Thuerey
Technical University of Munich [kanishk. wwa.baaiiaa ia, f. koehler, htt.:rrinnssstt.e...eee..e tturrey)(tum..... de
"Equal contribution
ABSTRACT
The physics solvers employed for neural network training are primarily itera- tive, and hence, differentiating through them introduces 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 differen- tiable 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 train- ing, it enables significant compute savings without sacrificing network accuracy. Our focus is on differentiating iterative linear solvers for sparsely discretized dif- ferential operators, which are fundamental scientific computing, PRDP ap- plicable 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- neural- hybrid solvers. In the challenging example of emulating the Ntvier-Stokes equa- tions, we reduce training time by 62%.
INTRODUCTION
Differentiable Physics paradigm which allows learning algorithms interact with gradients of classical physics solvers. This has proven effective across many domains, e.g., solving inverse problems (Bendsoe Sigmund, 2013), integrating physical constraints (Raissi et al., 2019; Li al., 2024), and especially, creating hybrid models that blend classical numerical techniques with learned components (Um et al., 2020; Kochkov et al., 2021; 2024). Despite their promise, cuurdlhhybrdd models for differential equations face limited adoption due to the computational cost of executing and differentiating through classical solvers during training. At the core of most classical solvers for differential equations are iterative processes that can be tuned for accuracy, typically by adjusting parameters such as step size or iteration count. Traditionally, these methods prioritize achieving the highest possible accuracy in the physics solver. In contrast, our work takes a novel approach, drawing inspiration from bi-level optimization (Pedregosa, 2016). Rather than focusing solely on maximum physics accuracy, we explore how numerical solvers can be strategically adjusted to sub- stantially accelerate the training process without losing network accuracy.
Our work applies to training pipelines involv- ing differentiable physics, as exemplified in fig- ure 1. Repeatedly querying solver with sev- eral iterations in the forward pass and differ- entiating through its iterations in the backward pass introduces severe computational bottle- neck during training. Since deep learning is in- herently based noisy gradient estimates, show that neither the physics nor the physics' Jacobian must be fully converged training
x
N
Pk
12
Figure 1: A neural network training pipeline using differentiable physics solver Px- Black and red arrows show the forward and backward passes, re- spectively. As solver iterations K grows, the cost of passes through Px becomes severe.
Munich Center for Machine Learning

---

Published as conference paper at ICLR 2025
10a
Converged
Physicx salver itterttonn
40
125
PRDP
Training Time [s]
IC Savings
Vallddation
20-1
20
75
50
Savings
10-2
25
50
75
100
25
50
75
100
Converged PRDP
Epoch (a)
Epoch (b)
(c)
Figure 2: Progressively Refined Differentiable Physics (PRDP) reduces the training time of neural networks containing numerical solver components (c). The fidelity of iterative components in- creased only if validation metrics 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).
time to attain good generalization. Using differentiable numerical solvers at level significantly coarser than needed for tight tolerances and progressively refining it starting from an even coarser level sufficient achieve full accuracy of the network.
type of bilevel optimization. In the context of hyper- parameter optimization (a typical bilevel optimization problem), Pedregosa (2016) explored the idea of successively refining an inner solver make increasingly accurate updates in outer optimizer. Under strict assumptions of convexity, it was shown that with a summable sequence of refinement levels (in terms of increasingly tighter tolerances for both forward and backward solve), conver- gence in the low-dimensional hyperparameter space can be achieved. We build upon this work but instead, treat the nonconvex learning of neural network parameters as an outer problem. Our in- ner problem is the most elementary operation in any scientific computing, the solution of linear system of equations from the discretization of physical model. In this more general setting, we observe that we cumulatively reduce inner iterations not only by progressive refinement but also by ending the refinement a level significantly coarser than needed for full physics convergence. This is possible because of the approximative nature of neural network training for which highly accurate gradients are not required. Similar approaches have been investigated for Deep Equilibrium Models (Bai et al., 2019), which have nonlinear root- not-finding with dense Jacobians as the inner problem, e.g., by (Shaban et al., 2019; Fung et al., 2021; Geng et al., 2021). In contrast, we focus on the sparse lin- systems arising from the discretization of partial differential equations (PDEs) admitting special solution characteristics that have been unexplored in the context of deep learning. For realistically large and sparse linear systems of equations, the prevalent class of solution methods are iterative linear solvers (Saad, 2003). By controlling the number of iterations in these solvers (and during their backward passes), we can directly balance physics refinement and computational cost.
Differentiable Physics can be understood
Reducing the cumulative number of solver iterations performed throughout the network training re- sults in considerable compute savings since linear solves typically dominate times. However, the optimal refinement schedule and the sufficient level of refinement are problem-dependeee automatically determine them during the network training, we present novel algorithm, Progres- sively Refined Differentiable Physics, in which the refinement of the physics is adaptively increased a plateau in terms of validation metrics is encountered. This idea visualized in figure 2.
To
Our experiments focus on efficiently training neural networks with differentiable linear solvers in the loop. We address unrolled as well as implicit differentiation methods, showing that PRDP applies ef- fectively to both. The approach is tested on training tasks across range of PDE problems, including the Poisson, heat diffusion, Burgers, and Ravier-Stoker equations. Empirical insights into PRDP's behavior presented for ID problems, and its performance is validated through more complex 2D and 3D time-stepping problems. We demonstrate its effectiveness variety settings such as inverse problems, differentiable physics losses, and correction-bast approaches.
In summary, our main contributions are the following.
We empirically demonstrate that full network performance be achieved with coarse level of physics refinement, well below the typical refinement required for full convergence, leading significant computational savings.
2

---

Published as a conference paper at ICLR 2025
We introduce the Progressively Refined Differentiable Physics (PRDP) algorithm, which adaptively identifies the optimal level of physics refinement during training.
We validate the effectiveness of PRDP across various differentiable physies learning sce- narios, demonstrating its broad applicability.
2 DIFFERENTIATING ITERATIVE LINEAR SOLVERS
A theoretical understanding of diffcrentiating iterations can be gained through the example of an inverse problem involving a parameterized linear system Aux, by Such systems arise in the numerical solution of discretized PDEs, where refers the spatial discretization width (Ames et al., 2014). For this, both the system matrix A and the right-hand side b, arise from the assembly routines A() and ((). respectively. An iterative solver creates a sequence of [0] [1]
guesses {ur 2/ ... n K. { {**(,, A. b=) that should converge to the analytical solu- tion u/ [1]] x u = A by up to a given tolerance e. We denote the number of solver iterations required to achieve this tolerance by K.. Clearly, the set of parameters
Rº affect the solution to this linear system of equations. Let Pk denote function mapping from the parameters 0 to
the iterative solution of the linear system of equations [/]]] via first assembling the system matrix A A(0) and the right-hand side by. 3(0) and then employing the iterator & for K steps. An example of the corresponding compute graph is shown in figure 21 in the appendix. The inverse problem is solved by performing an optimization discrepancy against a reference solution ni, via
the parameter space, aiming minimize the
1
min 8
L(0) = I( (0))
un (0) 2 S.L. A(0)u: 3(0)
2
This forms a bilevel problem where the outer optimization concerns minimizing the loss L(0), while the inner optimization pertains solving the linear system Px. Then, the loss function's evaluation can be explicitly written as L(0) = ((Px(0)).
Solving the outer optimization using a first-order method requires the gradient VoL. We can differ- entiate this chained function using reverse-mon modee automatic differentiation (AD, Griewank & Walther (2008)) to find the (transposed) gradient as
(VoL) = 0T
π. (1)
JP lo
(2)
primal inace.
adjoint
Clearly, the quality of the gradient depends the number of iterator steps K. Inaccuracies in the physics operator Pk propagate to the loss gradient V through two sources: primal inaccuracy, i... the Jacobian of the loss function Ji evaluated at the approximate solution Pk (x))). and adjoint inaccuracy, i.e., the Jacobian Jp la of the (approximate) iterative solver itself.
2.1 LINEAR SOLVER VJPS
AD frameworks do not assemble the full Jacobian; rather, they employ vector-Jacobian products (VJP) to reverse-propagate the gradient information (Murphy, 2023; Blondel & Roulet, 2024). The programmatic implementation for the VJP of the loss function ni, = 1 Ji Px (9) is straightforward to perform via AD. On the other hand, the VJP over the approximate physics Pk requires reverse propagation over the solver and the assembly routines, which we detail in appendix C. Conceptually, there are two approaches.
Implicit Differentiation The backpropagation over the iterative linear solye can be framed as the solution of another linear system in terms of an auxiliary variable with A ün This requires the transpose system matrix AT, Implicit differentiation over any kind of implicit relation gives rise to linear solve with the system's linearized form being transposed. It can be derived automatically with AD tools (Blondel et al., 2022). Having a linear solve in the primal execution well is special case, as the same system matrix appears in the forward as well as the backward pass. Hence, it is reasonable to employ the same iterator in the VJP as in the primal. but with the transposed system matrix and the different right- hand side AT ū.). The number of iterates required to
converge can be different from the primal and provide way control the adjoint inaccuracy.
3

---

Published as a conference paper at ICLR 2025
Unrolled Differentiation Unrolled differentiation applies AD directly to iterative program, treating each iteration as an individual computational step. It thereby accesses the VJP through the iterator & and accumulates each iterate's s contribution. This inherently requires access to all primal iterates. Since the AD engine unrolls as many iterations as in the primal pass, the adjoint accuracy naturally coupled with the primal accuracy.
2.2 SCHEDULING INNER ITERATIONS
Since the outer problem of equation is solved iteratively, the gradient Val. can be coarse at the beginning of the outer iterations and only requires the highest fidelity towards the end (Pedregosa, 2016). This fundamental idea was developed in the domain of hyperparameter optimization, and demonstrate its application to physical inverse problems. Consider the Poisson equation, which is prototypical elliptic partial differential equation found in many areas of science and engineering. Most discretization techniques lead to a linear system of equations with N degrees of freedom (for an example discretization, refer to section D.I).
Au(x) p((, 0). x E EEQC C CR
Au b/s uh: by E RRN AE MERM N
(3)
We consider a setting the one-dimensional unit interval Ω (0, 1) with homogeneous Dirichlet boundary conditions, u()) = 0 = u(1). Assuming the parameter space is one-dimensional, we design the right-hand side as p(x, 0) 0 sin(2 and discretize it on the domain. The outer optimization problem over 0 E IR can be exactly solved in this case.
We employ a Jacobi scheme solve the linear system which converges to = 10m5 within K, = 600 iterations; so does its unrolled Ja- cobian. Under gradient descent optimizer in the parameter space of 0, it requires a total of 125 outer iterations bring the relative pa- rameter suboptimality against 0" also down to 10-5. By using simple scheduling that starts the outer optimization at a coarse inner resolu- tion of K 25, and progressively increases the inner iterations by AK = in every outer iteration, all the way Ke, we can reduce the overall computational cost of the optimization while achieving the same parameter subopti- mality. These results are visualized in figure 3. Despite requiring slightly more outer iterations, fewer inner iterations are necessary. This re- sults an overall reduction of 28% the total number of inner iterations. These savings rep- resent one component of the final cost savings. As they are obtained due to scheduling fidelity of the physics solve, we denote them as progressive refinement (PR) savings.
physics solver iterations
Fully converged Scheduled
300
600
PR savings
physics solver Iterations
600
Parameter relative
10-1
400
10-3
10-5
0
100
0
100
Outer iteration
Outer iteration
Figure 3: Progressively refining the differentiable physics during outer optimization of Poisson in- verse problem achieves full convergence of the parameter with fewer cumulative iterations of the physics solver, leading to progressive refinement (PR) savings.
2.3 NETWORK TRAINING UNDER INCOMPLETELY CONVERGED DIFFERENTIABLE PHYSICS
Going beyond simple convex inverse problems, more sophisticated compute graphs arise. For ex- ample, assume the linear system is assembled from a prior variable that is given as the output of neural network gh := f(:00) (ignoring the input to the network for now). In this case, we have A A(gr) and bn ((gh). Hence, the optimization over 0 turns into the nonconvex learning problem in the weight space of the neural network. The extended reverse-mode AD operation reads
(VoL)* = 8" = -(1T" JilPx(t( Jpp Jpuls(:s0))
J
(4)
The approximate solution to the linear system of equations using K steps and its differentiation remain the sources of gradient inaccuracy. Yet, the effect of the neural network alters the solution characteristics of the iterative linear solver through the assembly of the system matrix (influencing its spectrum) and right-hand side.
4

---

Published as a conference paper at ICLR 2025
Conversely, we hypothesize that neural network training can also work under approximate gradients. To illustrate this, consider the one mne-dimensional heat equation on a periodic unit interval with a time- implicit discretization
du 8u = Oz.
a E (0, = Df Lu
[++1] = [d]
(5)
01
where C1 represents the matrix from the spatial discretization of the second derivative (see sec- tion D.2 for more details). The physies operator now advances from step the next ui ((u))) ( Consider the scenario in which a neural network performs a first prediction in time from an initial state, i... (1 = f(,) [0), 0). followed by the physics operator involving iterator steps which provides the second time step solution [?[[K] Px ).. The loss in equation is computed against reference given by applying the converged physics operator twice unh... Pk (u Hence, the network trained emulate one application of P. We solve the linear system of equa- tions using the Jacobi method (B.1), which converges within Kg 25 iterations. The implicit differentiation's 's linear solve requires equally many iterations.
i++1]
[0]
Figure 4 shows the results of multiple network training runs performed at different levels of physics refinement, i.e., different values of K that are kept constant for each training run. We observe that beyond a refinement threshold Kmax. no further improvements of the neural emulator's performance on validation metric are noticeable. Hence, we conclude that net- works can attain high level of accuracy, even when trained through incompletely converged (IC) physics. As explained in equation 2, this incomplete physics convergence subsumes coarseness of its Jacobian. In the example of the heat emulator, we obtain Kmax 10, which reduces the number of physics solver iterations for training by 60%. We denote this second type of reductions as IC savings.
# physics solver iterations
10
25
Kmz K.
Validation Error during training
10
Validation Eiror at final apoch
10a
10a
IC savings
10-1
10
Epoch
# physics salver iterations
Figure 4: Network accuracy does not improve be- yond refinement level of differentiable physics (Kmx) significantly lower than full convergence (K.) constituting incomplete convergence (IC) savings.
While we could not isolate single cause for these IC savings, we hypothesize that is likely due to combination of factors: (1) the more
critical components of the gradient may converge more rapidly than less important ones, (2) the inherent noisiness of neural network training due stochastic mini-batching, (3) the usage of momentum-based optimizers, and (4) the inherently approximative nature of machine learning mod- els. Taken together, these factors presumably reduce the need for full convergence in the differen- tiable physics solver, resulting in very substantial cost savings.
3 PROGRESSIVELY REFINED DIFFERENTIABLE PHYSICS
The inner iterations saved by progressive refinement (PR) and incomplete convergence (IC) can greatly reduce training time without impeding accuracy. However, suitable refinement schedule and Kmax are unknown a priori. They depend on the PDE, its discretization (given by and 3), the iterative linear solver, and the learning dynamics.
To arrive at a practical method, we automate the progressive refinement and the detection of Kmax based the observation that training progress in terms of validation metric stagnates when the physics inaccuracy too large. This stagnation corresponding to different physics refinement levels exemplified in figures 3 and Conversely, we track training progress over time to automatically increase physics refinement by AK linear solver iterations when training plateaus. Typically PRDP starts at Ko
TThoouhhut this work, superscripts in square brackets to denote sequence entries, for example for temporal snapshots iterates of linear solver.
5

---

Published as a conference paper at ICLR 2025
def update state_in,, state_out, inner_ iimer_iteraticnn)
$ nepra) network grvee rirst time step, prysice gives second time step stale_l model (stale_in) state_2 " physics (state_l, inner_ iteretiner)) lass, grad sellu_and__rrdd_oee modal update s_modal imodel, grod)
(state_2, state_out!
model, loss
inner_ inner_iterationa vel_loss_histon
(eveluate vel_áetaaooderr))
for epoch in reage (epochs);
for stata_in, state_out in dataloader:
model, loas update state_in, state_out, inner_ inmer_tterationg bllos__hsttrry.aggedd
(model, val_dataloadec) if should_refinee (vall aaraalloss_hhttrrne inmer_iterations
A typical mixed-chain learning pipeline (figure 1) as used in our neural emulator learning experiments. Training begins with coarse physics, using, e.g., solver iteration. The should.re re f ne function applies PRDP, determining when to refine and progressively increasing the inner solver iterations during training.
Listing 1:
Start wi coarse physics
Controlling Refinement Given an exponentially smoothed his- tory of validation metrics (ī/a)) after a representative training in- terval, for which we can either use epochs or fixed number update steps, we distinguish three trends:
Validation Metric
Refine
Reline
1. Validation metric plateaus: If the ratio of the latest vali- dation error and the value of several grace intervals 8 ear- lier, r = i.!-1) /i!-a] is above a threshold r > Tste, it indicates that the network has achieved the highest possi- ble accuracy at the current refinement level. Then, sub- ject the following checks, the refinement is increased by AK. We record the validation metric before refinement as checkpoint + _. Identifying this trend leads full utitilization of each refinement level and contributes the PR savings.
Epochs
Tatop
No
Refine
Continue
Continua referement
euperiunue)
2. Adequate refinement is achieved: We compute the improvement of the current plateaued validation metric against the checkpoint, Fe i'="'// /c. If the stagnation is greater than its threshold. i.e., ra > Tstop- the refinement level is retained. This identifies the adequate level of re- finement necessary, Kmax enabling the IC savings.
Figure 5: Top: the typ- ical training progress of neural network supported by PRDP, showing the ratios and тe. Bottom: a simpli- fied flowchart representation of the PRDP control algo- rithm.
3. Initially divergent regime:
For some scenarios, there a minimum number of iterations Kmin which if K < Kmin lead to the convergence of the network to worse value than its initialization, e.g., for K = and K 2 in figure 4. To overcome this, we also refine if re > 1.0.
Figure illustrates the basic concept of the PRDP control algorithm. The exact algorithm is detailed in pseudocode in Algorithm 4. Its implementation in training pipeline that uses differentiable physics is represented by the should-refin function in listing
Applying Refinement Intrinsically, PRDP pertains to refinement of the primal physics. As we showed in section 2.1, the adjoint refinement is controlled based on the kind of differentiation used. In implicit differentiation, tion rule. This approach inherently decouples the convergence of the primal and its Jacobian. We maintain their coupling by using the same solver and number of inner iterations for the primal solve and the VJP solve. This setup ensures that adjustments to the number of inner iterations consistently affect the primal and the VJP. When using unrolled differentiation, progressive refinement of the primal trivially extends to progressive refinement the Jacobian. This because the number of iterations unrolled by the automatic differentiation engine is the same as that of the primal solve.
additional linear system solve is performed via a custom differentia-

---

Published as a conference paper at ICLR 2025
104
2-step NMSE x-vel ooiit
coarse
Physics Refinement
2-step nMSE
5-step
20
10"2
fine
0
50
100
0
25 Epochs
50
70
0
50
100
50
100
50
100
PRDP
Epachs
Epochs
Epochs
Epochs
(a) Heat 1D
(b) Heat 2D
(c) Heat 3D
(d) Burgers
(e) Navier Stokes
Figure 6: PRDP achieves the same network accuracy fully converged physics (baseline in dark purple). It does so by adjusting physics accuracy adaptively over the training if validation perfor- mance plateaus, which particularly noticeable in Heat 2D.
4 EXPERIMENTS
To validate our approach across problems of varying complexities, combination of different learn- ing and physics scenarios are chosen. This includes the convex optimization an exactly solvable inverse problem with the Poisson equation described in section 2.2, the learning of autoregressive neural emulators in the spirit of Bar-Sinai et al. (2019) and Brandstetter et al. (2022) for the diffusion and Burgers equation but with loss setup involving differentiable physics, and third, the learning of a neural-hybrid emulator in a setting similar to Um et al. (2020). When training neural networks, our results show the aggregation over ten different initialization seeds. The validation metric used the normalized mean squared error (nMSE) over the validation dataset. Further specifics are provided in appendix
4.1 LINEAR INVERSE PROBLEM
In section 2.2, we demonstrated the potential progressive refinement savings on a doubly- convex inverse Poisson problem when employing pre-defined linear scheduling of in- ner iterations across outer iterations. This schedule was hand-tuned based on exhaustive runs. PRDP is designed be problem-independenl and adjust the inner iterations adap- tively. To confirm this, we apply PRDP in the same setting and achieve a saving 33% due to progressive refinement. This is slightly higher than the manually achieved 28%.
Moreover, we extended this inverse prob- lem to space, with each entry scaling the first three eigenmodes, and covered combi- nation of different setups, including the steepest descent solver (SD), and implicit differentiation (Impl) next to unrolled dif- ferentiation (Unr). The results presented in the table to the right show that PRDP universally applies to each combination and is, hence, agnostic to the linear solver and the differentiation method. The qualitative behavior of all combinations displayed in appendix figure 15. For brevity, only list the PR savings achieved using PRDP by displaying the cost of optimization with K, and the corresponding reduction (Red.). Importantly, under the larger param- eter space within each of the four setups, PRDP works equally well, achieving PR savings of 48% and 43% for the Jacobi method and the steepest descent solver, respectively.
firee-dimensionll parameter
P Type
Diff
Cost
Red.
PR. Sav.
Jac
Unr
75K
24.87K
33%
Jac
Impl
75K
24.87K
33%
3
Jac
Unr
337.5K
163K
48%
3
Jac
Impl
337.5K
163K
48%
3
SD
Unr
192K
82.7K
43%
3
SD
Impl
192K
82.7K
43%
4.2 LINEAR NEURAL EMULATOR LEARNING
We train multilayer perceptron (outer problem) to function as an autoregressive emulator, i.e., a replacement for the numerical time stepper of a heat diffusion equation. Following the differential physics paradigm, the numerical heat equation solver (inner problem) is in- cluded in the gradient loop during training. This training pipeline is depicted in figure

---

Published as a conference paper at ICLR 2025
In section 2.3, using ID case
Cumuuaavve salver taradoon
of the same setup. we demon- strated that savings due in- complete convergence are actu-
2500
2000
1500
2000
IC
2000
ally achievable, and that they
1000
1000
1090
PR
are likely caused by the nature of the training in deep learn- ing. We identified Kmsx based on an exhaustive search of dif- ferent K values. This is infea- sible in practice. Hence, PRDP designed to automatically find the upper refinement threshold and prevent further superfluous refinement. In figure 6(a), (b) & (c), this can be seen as its con- vergence follows certain refinement levels. We can confirm that the same Kmax = 10 found by PRDP constituting the aforementioned 60% IC savings. Moreover, as figure 7(a) reveals, the total PRDP savings are 79% against the fully converged run with K, = 25 because PRDP additionally contributes 19% savings due progressive refinement. We repeated similar experiment in two and three dimensional settings. In the 2D case, the difference between K, = 43 and Kma 12 constitutes 72% IC savings. Together with 14% PR savings, this totals 86% savings due to applying PRDP as visualized in figure 7. Similarly training ResNet (He et al., 2016) in the 3D case, the total savings due to PRDP were 81%. This corresponded to a reduction in total training time by 24% and 78% in the 2D and 3D cases, respectively, which underscores PRDP's potential for large computational savings in more difficult, higher-dimensional settings.
Converged PRDP
Converged PRDP
Conwerged PRDP
Figure 7: Total Inner iters in Heat ID, 2D 3D.
4.3 NONLINEAR NEURAL EMULATOR LEARNING
The preceding tests constituted inner problems where only the linear sys- tem's right-hand side was dependent on the trainable parameters. We ex- pect our PRDP algorithm to work equally well when training through in- ner problems with non-constant matrix assembly function A. To illus- trate this, we train a ResNet to function for the one-dimensional Burgers equation on a periodic unit interval with time-implicit upwinding discretization. Additionally, this yields non- symmetric matrix requiring the more sophisticated GMRES solver (Saad & Schultz, 1986). We use linearization around the previous point in time, re- sulting in an Oseen problem (Turek, 1999). When training a ResNet under a similar mixed scenario as in section 4.2, we again confirm the working of PRDP in this case. It amounts to 30% savings by incomplete conver- gence reducing K, 23 to Kmax 16. An additional 29% PR savings yield total of 59% PRDP savings. The Burgers emulation problem was particularly challenging, requiring us set Ko 4 to overcome strongly divergent regime if the physics was too coarse at the beginning of training. This is noticeable in the initial epochs in figure 6(c). This figure also shows that since the fanning between the different (convergent) refinement levels is not as strong as before, the IC savings are lower. However, we still see significant savings due to progressive refinement.
Cumulative Itarations
2000
1500
autoregressive emulator
PR
Converged
Figure 8: Total inner iters in Burgers.
4.4 NEURAL-HYBRID EMULATOR FOR THE NAVIER-STOKES EQUATION
Tightly combining classical numerical solvers and neural networks is a promising research di- rection (Kochkov et al., 2021; Um et al., 2020). differs from our previous experiments in that the solver not just part of the training com- pute graph but will also be executed during the inference of the model.
Pk
NN
Px
NN
-
L
Px.
NN
My
I++
Figure 9: The training (top) and evaluation (bot- tom) pipelines for neural-hybrid emulator.
We mimic the setup of Um et al. (2020) to have neural component correct the discretiza- tion error of an incompressible Navier-Stokes solver by training it against reference pro-
8

---

Published as a conference paper at ICLR 2025
duced at higher resolution of discretization. The training and evaluation setups are shown in figure 9. The outer and inner problems correspond to the network training and the Savier-Stoker solver in the training loop, respectively. Similarly to the Burgers example, based discretization with the linearization around the previous step in time. The additional incom- pressibility constraint leads saddle point structure (Turek, 1999). We choose solve it in coupled form with the GMRES algorithm (Saad & Schultz, 1986).
choose an upwind-
In this complex scenario, PRDP likewise gives substantial benefits: It is able to save 65% inner iterations due to incomplete conver- gence and 16% iterations in progressive refinement. Ultimately. reducing the number of iterations by a total of 81% results in a reduction in wall clock training time by 62%, as shown in fig- ure 2(c). Since neural-hybrid approaches execute the fully con- verged physics solver during inference and consequently during validation, PRDP's reliance
Cumulative solver Icoratonn
4000
2000
2000
1000
PR
validation metric necessitates this additional cost. In our experiments, we found that compared to training with fully converged physics without computing the vali- dation loss, PRDP was still faster by 57%.
Converged PRDP
Figure 10: Total inner iters in Navier Stokes.
5 RELATED WORK
Differentiation through implicit relations Fischer (1991) established implementation frame- work for unrolled differentiation and investigated the convergence of the derivative, focusing on linear solvers. Charles (1992) extended this work to a broader class of iterative processes. As an alternative to unrolled differentiation. Christianson (1998) derived the implicit differentiation rules over linear solves. Beyond linear systems, implicit differentiation has since gained prominence in the machine learning community, particularly for hyperparameter optimization (Bengio, 2000) and other bilevel optimization (BLO) applications that require differentiating through inner optimiza- tions, such as deep equilibrium models (Bai et al., 2019) and meta-learning (Andrychowicz et al., 2016). We view differentiable physics (Thuerey et al., 2021) as another type of BLO with the spe- cialty of sparse linear systems. The practicality of implicit differentiation is largely due to its reduced memory footprint and lower reverse-men mode computational cost, especially as modern automatic dif- ferentiation (AD) tools can automatically handle the necessary propagation rules (Blondel et al., 2022). Additionally, implicit differentiation allows for the black-box use of solver implementations, enabling the integration of third-party components into differentiable computational graphs (Giles, 2008). However, unrolled differentiation remains an active research area, with work focusing on non-asymptotic analyses (Scieur al., 2022) and other aspects (Maclaurin al., 2015; Franceschi et al., 2017; Grazzi et al., 2020; Ji et al., 2021).
Analysis and cost-reduction of bilevel optimization To address the computational cost of bilevel optimization, Fung et al. (2021) presented Jacobian-free backpropagation, where the Jacobian of the implicit solver is approximated as the identity, climinating the need for adjoint linear solves. Geng et al. (2021) introduced phantom gradients, where the matrix inversion is replaced with an approxi- mate inverse, and provided theoretical guarantees on the convergence of stochastic gradient descent as the outer problem. Lorraine al. (2020) approximates the implicit linear solves with reduced number of conjugate gradient steps. Moreover, Shaban et al. (2019) and recently Bolte et al. (2023) discuss unrolled differentiation through only reduced number of iterations. All the aforemen- tioned works consider (what we call) incomplete convergence (IC) savings, albeit in the setting of deep equilibrium models (Bai al., 2019) or hyperparameter optimization (Feurer & Hutter. 2019). On the other hand, Pedregosa (2016) studied approximate hypergradients by adjusting the tolerances of inner primal and implicit linear solves following a pre-defined schedule. This work proves (what we call) progressive refinement (PR) savings. Prior to this, Domke (2012) investigated outer train- ing with unrolled AD through incompletely converged iterations, specifically for gradient descent, heavy ball, and L-BFGS methods as inner optimizers. They note the advantage of implementing incomplete convergence using the number of inner iterations rather than inner tolerance. Our ap- proach uniquely combines PR and IC savings through both unrolled and implicit differentiation targeting iterative solvers for sparse linear systems embedded within neural network training. To our knowledge, no prior work has applied these techniques in this context.

---

Published as a conference paper at ICLR 2025
6 LIMITATIONS AND OUTLOOK
Limitations While PRDP is designed to work across a range of differentiable physics settings, there are, of course, way more potential linear systems that can arise, all with their specific charac- teristics. Albeit we believe that the approach using scheduling over iterations rather than scheduling via tolerances might be more generally applicable, remains be tested how PRDP applies to un- structured discretizations in higher dimensions potentially also involving multiple physics. PRDP is limited to settings that involve iterative linear solvers. As such, it can not be used for purely explicit numerical solvers (c.g., found in strongly hyperbolic problems) when the linear systems solved spectrally (Kochkov et al., 2024). However, most other simulations in science and engineering en- tail linear solvers either due more efficient implicit time stepping when solving steady-state problems. Moreover, whenever dynamics are constrained, e.g., the incompressible Navier-Stoff Stokes equations, even if purely explicit time stepping is used, iterative processes are required to resolve the constraints.
Outlook and Impact Progressively Refined Differentiable Physics (PRDP) provides a means to exploit both savings due to progressive refinement and incomplete convergence, thereby greatly ducing the cost of neural network training with differentiable physics. This could enable settings that so far have been infeasible due to prohibitive expenses like long temporal unrollings or differentiable physics on high resolutions or in three dimensions.
Our work opens up many interesting directions for future investigations, such smoother relations between the achieved plateau ratio r and the conducted iterations/prescribed tolerance in the physics solver. Potentially, those could be faster than the linear increments we used in this work. For neural-hybrid emulators, IC savings via Kmax could also extend to the inference stage. So far, PRDP couples primal and adjoint (in- Jaccuracy. One can use the unique properties of either unrolled or implicit differentiation for more sophisticated approaches. This can include an imbalanced number of iterations in the primal solve and the implicit linear solve. For unrolled differentiation, one can unroll a different number of iterations either reversely following Bolte et al. (2023) or Shaban et al. (2019) or from the beginning- Other levels of refinement than by the number of iterations of a linear solver are likewise interesting directions for future work, e.g., using differences in spatial or temporal resolution together with cesolutiio-agnostii neural emulators.
7 CONCLUSION
This work investigated neural network training through incompletely converged differentiable physics. Our objective was to reduce the cost of gradient computation without sacrificing training accuracy. Prevalent research on training with incompletely converged gradients focuses on bilevel optimization problems, primarily in meta-learning or hyperparameter optimization contexts. This work extends the research in the differentiable physics space, focusing on iterative linear system solves associated with discretized differential operators.
We have demonstrated that neural networks can be trained through differentiable physics solvers significantly coarser than full convergence. Our approach of Progressively Refined Differentiable Physics combines compute savings from both progressive refinement and incomplete convergence. This yielded favorable outcomes across all our test scenarios. It makes initial training progress cheap using coarse physics and carefully improves training accuracy using adaptive physics refine- ment over time, ending the training at a refinement significantly below primal convergence. In total, achieve up 86% fewer cumulative number of physics solver iterations than training with fully converged physics, and 78% reduction in training time. Our approach has the potential for numer- ous practical improvements in learning pipelines that involve differentiable numerical solvers and could facilitate integrating simulators into training that were previously considered computationally infeasible.
8 REPRODUCIBILITY STATEMENT
To ensure reproducibility, we detail all physics parameters, discretization schemes, boundary con- ditions, iterative solvers, network architectures, optimizers, learning rate schedules, data generation
10

---

Published as a conference paper at ICLR 2025
methods, train-test splits, and batch sizes in the appendix. Additionally, the full source code for experiments is available at https: /tttpps:/ggttee /ttpthutt.tttt/// Saauu--pb/PRRR /PRDP.
REFERENCES
W.F. Ames, W. Rheinboldt, and A. Jeffrey. Numerical Methods for Partial Differential Equa- fions. Applications Mathematics Series. Flsevier Science, 2014. ISBN 9781483262420.. URL https: /books google de
Marein Andrychowicz, Misha Denil, Sergio Gomez Colmenarejo, Matthew W. Hoffman, David Pfau, Tom Schaul, and Nando de Freitas. Learning to learn by gradient descent by gradient descent. Daniel D. Lee, Masashi Sugiyama, Ulrike von Luxburg, Isabelle Guyon, and Ro- man Garnett (eds.). Advances in Neural Information Processing Systems 29; Annual Conference on Neural Information Processing Systems 2016, December 5-10, 2016, Barcelona, Spain, pp. 3981 1989, 2016. URL https: http///pcceeeniggceeee newrips...p.... ccc hppppppppc11. / html.
Shaojie Bai, J. Zico Kolter, and Vladlen Koltun. Deep equilibrium models. In H. Wal- lach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett (eds.), Ad- vances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019. URL https: file/01386 1tte:/0..86666660011222440777171
Yohai Bar-Sinai, Stephan Hoyer. Jason Hickey, and Michael P Brenner. Learning data-driven dis- cretizations for partial differential equations. Proceedings of the National Academy Sciences, 116(31):15344-1539) 2019.
Martin Philip Bendsoe and Ole Sigmund. Topology optimization: theory, methods, and applications. Springer Science Business Media, 2013.
Yoshua Bengio. Gradient-based optimization of hyperparameters. Neural Comput.. 12(8):1889- 1900, 2000. doi: 10.1 10.1162089976603300115888 URL attps://dols oorgi.org/ii wt..ooor//.......o.. 1162/
089976503300051195
Mathicu Blondel and Vincent Roulet. The Elements of Differentiable Programming. arXiv preprint
2rXhv:2403.14606 2024.
Mathieu Blondel, Quentin Berthet, Marco Cuturi, Roy Frostig, Stephan Hoyer, Felipe Llinares- Lopez, Fabian Pedregosa, and Jean-Philippe Vert. Efficient and modular implicit differentiation. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural Information Processing Systems, volume 35, pp. 5230-5242. Curran Associates, Inc., 2022. URL httpss://pocceeeeeeeee http://rooeedd.........eeei...ss www.rps......... cappprr saappprrrflls_eeeeeeeeeeeeeeeses files/ patas//pppe//2222 pdf.
Jéröme Bolte, Edouard Pauwels, and Samuel Vaiter. One-step differentiation iterative algo- rithms. 2023. URL http://pppprrr.r...cccp/pppeeeee
James Bradbury. Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Miine and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL http:
// //github. com gooogle// jax.
Johannes Brandstetter, Daniel E. Worrall, and Max Welling. Message passing neural PDE solvers. In The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022. OpenReview.net, 2022. URL https: http:::s//pppeevveeee /oppnrevieeeeeeeeeeeeeeeeeee net fformm?id vSi YHSU
Gilbert Jean Charles. Automatic differentiation and iterative processes. Optimization Methods & Software, 1:13-21, 1992. URL https: :// /ppi www.ntisschooar.ccoooooo httpp::/apl...mmmiiccooo.aaaa... cttpriisssssaaaa.ooooooopppsaa 120894038.
11

---

Published as a conference paper at ICLR 2025
Bruce Christianson. Reverse accumulation and implicit functions. Optimization Methods and Soft- ware, 94):307-322, 1998. doi: 10.100/105568999880667777
DeepMind, Igor Babuschkin, Kate Baumli, Alison Bell, Surya Bhupatiraju, Jake Bruce, Peter Buchlovsky, David Budden, Trevor Cai, Aidan Clark, Ivo Danihelka, Antoine Dedieu, Clau- dio Fantacci, Jonathan Godwin, Chris Jones, Ross Hemsley, Tom Hennigan, Matteo Hessel, Shaobo Hou, Steven Kapturowski, Thomas Keck, Iurii Kemaev, Michael King. Markus Kunesch, Lena Martens, Hamza Merzic, Vladimir Mikulik, Tamara Norman, George Papamakarios, John Quan, Roman Ring, Francisco Ruiz, Alvaro Sanchez, Laurent Sartran, Rosalia Schneider, Eren Sezener, Stephen Spencer, Srivatsan Srinivasan, Miloš Stanojević, Wojciech Stokowiee, Luyu Wang, Guangyao Zhou, and Fabio Viola. The DeepMind JAX Ecosystem, 2020. URL http: ///gttub.... com/ google wenggl--deepiin deepmind.
Justin Domke. Generic methods for optimization- based modeling. Neil D. Lawrence and Mark Girolami (eds.), Proceedings of the Fifteenth International Conference Artificial Intelligence and Statistics, volume 22 of Proceedings of Machine Learning Research, pp. 318-326, La Palma, Canary Islands, 21-23 Apr 2012. PMLR. URL https: attps://proceeddee /ptpceeeeoceessssoeoee/e/// mlr. press/ v22 2022/dokkell 12 html.
Matthias Feurer and Frank Hutter. Hyperparameter optimization. Automated machine learning: Methods, systems, challenges, pp. 3 3-33, 2019.
Herbert Fischer. Automatic differentiation of the vector that solves a parametric linear system. Jour- nal of Computational and Applied Mathematics, 35(1):169-18)) 1991. ISSN 0377-0427. doi:
URL https: www. www.scieeed.rer.t 90205x.
Luca Franceschi, Michele Donini, Paolo Frasconi, and Massimiliano Pontil. Forward and reverse gradient-based hyperparameter optimization. In Doina Precup and Yee Whye Teh (eds.), Pro- ceedings of the 34th International Conference on Machine Learning, ICML 2017, Sydney, NSW; Australia, 6-11 August 2017, volume 70 of Proceedings of Machine Learning Research, pp. 1165- PMLR, 2017. URL http: http::///pccceddeeeeooooo mlr. ww....press/eemmmmm ppress/vii v70 Franceschi17a html.
Samy Wu Fung, Howard Heaton, Qiuwei Li, Daniel Mekenzie, Stanley J. Osher, and Wotao Yin. Jib: Jacobian-free backpropagation for implicit networks. In AAAI Conference on Artificial Intel- ligence, 2021. URL. https: /appi.app...////// semant wcchelccclll.....ccocccc ar. www.oo////oroso.o
Corpus 2338 98721.
Zhengyang Geng, Xin-Yu Zhang, Shaojic Bai. Yisen Wang, and Zhouchen Lin. On training im- plicit models. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 24247-24260. Cur- ran Associates, Inc., 2021. URL https: httpsttt/////cteeeee httpp//oeeeeiin...rrioo.......oo www.wnw..ii...... paper pcpppp...pr.................. paper/2021/Eii file www.acc6dc66/0. 7ct.://0000011111112222220000020
files/
Mike B Giles. Collected matrix derivative results for forward and reverse mode algorithmic differ- entiation. In Advances in Automatic Differentiation, pp. 35-44. Springer, 2008
Riccardo Grazzi, Luca Franceschi, Massimiliano Pontil, and Saverio Salzo. On the iteration com- plexity of hypergradient computation. In International Conference on Machine Learning, 2020. URL https: /api. semanti iischola.......... org/
22025503222
Andreas Griewank. Achieving logarithmic growth of temporal and spatial complexity in reverse automatic differentiation. Optimization Methods and software, 1(1):35-54 1992.
Andreas Griewank and Andrea Walther. Evaluating Derivatives. Society for Industrial and Applied Mathematics, second edition, 2008. doi: 10. 01..1177/1.788888771777777 URL https: /epubs.
siam. wwoo/oooooooooooooo /oi/ /ts/ 1137/1111111 11118///..78800 07808987100870
Francis H Harlow and Eddie Welch. Numerical calculation of time-dependent viscous incompress- ible flow of fluid with free surface. The physies of fluids, (12):2182-2189) 1965.
12

---

Published as a conference paper at ICLR 2025
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In 2016 IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2016, Las Vegas, NV, USA, June 27-30, 2016, pp. 770-778. IEEE Computer Society, 2016. doi: 10.1109/CVPR.2016..0 URL https:: //oi/.... org/ O. 11 09, Ct////VPR.22 6. 90.
Raad Issa, AD Gosman, and AP Watkins. The computation of compressible and incompressible recirculating flows by a non-iterative implicit scheme. Journal of Computational Physics, 62(1): 66-82, 1986.
Kaiyi Ji, Junjie Yang, and Yingbin Liang. Bilevel optimization: Convergence analysis and enhanced design. In International conference machine learning, pp. 188224892) PMLR, 2021.
Patrick Kidger and Cristian Garcia. Equinox: neural networks JAX via callable PyTrees and filtered transformations. Difserentiable Programming workshop at Neural Information Processing Systems 2021, 2021.
Dmitrii Kochkov, Jamie Alexander Smith, Ayya Alieva, Qing Wang, Michael Brenner, and Stephan Hoyer. Machine learning accelerated computational fluid dynamics. Proceedings of the National Academy of Sciences USA, 2021.
Dmitrii Kochkov, Janni Yuval. Ian Langmore, Peter Norgaard, Jamie Smith, Griffin Mooers. Milan Klöwer, James Lottes, Stephan Rasp, Peter Düben, et al. Neural general circulation models for weather and climate. Nature, pp. 1-7, 2024.
Zongyi Li. Hongkai Zheng, Nikola Kovachki, David Jin, Haoxuan Chen, Burigede Liu, Kamyar Azizzadenesheli, and Anima Anandkumar. thysics-informed neural operator for learning partial differential equations. ACM / IMS J. Data Sei., 1(3), may 2024. doi: 10.1 10.1145/3648506. URL
https: //oii./////i././// www.orrggoooooo..ooooooooo 1145/113000 3544560
Jonathan Lorraine, Paul Vicol, and David Duvenaud. Optimizing millions of hyperparameters by implicit differentiation. International Conference on Artificial Intelligence and Statistics, pp- 1540-1552. PMLR, 2020.
Dougal Maclaurin, David Duvenaud, and Ryan Adams. Gradient- hyperparameter optimiza- tion through reversible learning. In International conference on machine learning, pp. 2113-2122. PMLR, 2015.
Kevin P. Murphy, Probabilistic Machine Learning: Advanced Topics. MIT Press, 2023, URL http:://robtt
Fabian Pedregosa. Hyperparameter optimization with approximate gradient. In Maria Florina Balcan and Kilian Q Weinberger (eds.), Proceedings of The 33rd International Conference on Machine Learning, volume 48 of Proceedings of Machine Learning Research, pp. 737-746, New York, New York, USA, 20-22 Jun 2016. PMLR. URL wtwpttts//ppoeeeeeeeeeee press/vee /tte///8///eeee pewerr/eegggsaalll 6 html.
mlr.
J.Blair Perot. An analysis of the fractional step method. Journal of Computational Physics, 108 (1):51-58, 1993. ISSN 0021-9991. doi: http://eeeeeee/r.r..ee/////////2 htww:///wwwtttt..ctt/c www.://cceneeiii..oo rect. com/s http//sciinne/rrreee
URL https:
Maziar Raissi, Paris Perdikaris, and George E. Karniadakis. Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations. J. Compul. Phys., 378:686-707, 2019. doi: 10.1016/JJCP22018 10.004. 101 6 jep. . 201 8 10. 045.
URL www.thtps////...mm
Yousef Saad. Iterative Methods for Sparse Linear Systems. Society for Industrial and Applied Mathematics, second edition, 2003. doi: 10.1137711978888887778888 do /oo
URL https: /epubs.
s aiimm.oo wwtoorrg/ooooo abs/ 0 1137 111111111708111180111
Yousef Saad and Martin H Schultz. Gmres: A generalized minimal residual algorithm for solving nonsymmetric linear systems. SIAM Journal on Scientific and Statistical Computing, 7(3):856- 869, 1986. doi: 10.1 137/0907058.
13

---

Published as a conference paper at ICLR 2025
Damien Scieur, Gauthier Gidel, Quentin Bertrand, and Fabian Pedregosa. The curse of unrolling: Rate of differentiating through optimization.
S. Koyejo, S. Mo- hamed, A. Agarwal, D. Belgrave. K. Cho, and A. Oh (eds.), Advances in Neural In- formation Processing Systems, volume 35, pp. 17133-17145. Curran Associates, Inc., 2022. URL httpss://rroceeedd.g..... neurips
sttpapprrfllleeppperreeeee htl1://8539233002222340aaaeeeeee Paapeebdecc---perrrr-eee Conference pdf.
Amirreza Shaban, Ching-An Cheng, Nathan Hatch, and Byron Boots. Truncated propppropagttion for bilevel optimization. The 22nd International Conference on Artificial Intelligence and Statistics, pp. 1723-1732. PMLR, 2019.
Nils Thuerey, Philipp Holl, Maximilian Mueller, Patrick Schnell, Felix Trost, and Ki- won Um. Physics-based Deep Learning. physicsbaseddeep pattteeepperninnnne eeaeeeng..nn... org.
WWW, 2021. URL https://
Stefan Turek. Efficient Solvers for Incompressible Flow Problems An Algorithmic and Com- putational Approach, volume 6 of Lecture Notes in Computational Science and Engineering. Springer, 1999. ISBN 642-5--14--3 doi: 10.10077998--3-44--888889 //doi. oooi..org/100. 10000070777 101.000000797---4440400 --11---3--22-----22---3---
URL https:
Kiwon Um. Robert Brand, Yun (Raymond) Fei, Philipp Holl, and Nils Thuerey. Solver- in-the-loop: Learning from differentiable physics to interact with iterative pde-solvers. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Neu- ral Information Processing Systems, volume 33, pp. 6111-6122. Curran Associates, Inc., aeeaper_fllespppprr//aae
2020. URL httpss//pproceeddngss neurips wwwpurriss.cpppppeeeee fileAABee6aff331e0001112212111 oagp00000.ee22211eeeeeeeaaeaaeer
14

---

Published as a conference paper at ICLR 2025
A NOTATION
This section lists the primary notations used throughout the paper for clarity and ease of reference.
Ω RP: Spatial domain of the physics problem.
Spatial discretization width.
t: Time step index in time-stepping problems.
VARIABLES AND OPERATORS
Un E RR: The discretized solution vector.
: The direct solution to the linear system Au,, by-
IR-XX; The system matrix of the discretized PDE.
by E I The right-hand side vector.
A: The Laplacian operator, e.g., Au(z).
V: The gradient operator.
A: Auxiliary variable used in implicit differentiation.
3()): Function mapping parameters 0 to bj.-
A(O) Funcion mapping parameters to A.
(()): Function representing each iteration of an iterative solver.
|| ||-|: ||: Matrix or vector norm.
f.; 0): Neural network parameterized by 0.
T: Transpose of a matrix or vector.
T: Identity matrix.
0: Trainable parameters of the neural network or physical model.
K: Number of iterations performed by the iterative solver.
K₂: Number of solver iterations required to achieve tolerance e.
Pk: Differentiable physics operator, approximating the solution of linear system after K iterations.
u Approximate solution of the linear system after K iterations of the solver.
: Reference solution to the linear system, computed using either direct solver fully converged iterative solver.
L(0): Outer loss function, which measures the discrepancy between the predicted solution and a reference solution.
Ji, Jr Jacobians of the loss function and the physics operator, respectively.
PARAMETERS AND TOLERANCES
Convergence tolerance for iterative solvers.
Tstep: Threshold for the stepping criterion in PRDP.
Tstop: Threshold for the stopping criterion in PRDP.
6: Grace window for epoch intervals.
e: Exponential averaging window
A K: Increment in solver iterations used in progressive refinement.
15

---

Published as a conference paper at ICLR 2025
MISCELLANEOUS
Kg: Number of iterations required to reach a specified tolerance e the iterative solver.
K Maximum number of inner iterations sufficient for training accuracy.
min Minimum number of inner iterations required to prevent divergence during training.
re Ratios used evaluate the stepping and stopping criteria for progressive refinement, based on the validation metric's 's behavior.
B ITERATIVE LINEAR SOLVERS
The solution to linear systems of equations is fundamental seientific computing. Especially for partial differential equations discretized using fine resolutions or in higher dimensions, the discrete linear systems become large and sparse. Oftentimes, iterative solvers are the only practical way of solving them (Saad, 2003).
For our efforts to reduce the cost of differentiable physics as part neural network training, we consider three different iterative solvers. The Jacobi method belongs to the class of smooth- ing/relaxation methods. When reformulating a linear system solve as convex quadratic optimiza- tion problem, the algorithm of steepest descent naturally arises. To solve asymmetric and compli- cated systems, we also use the more sophisticated GMRES method. We implemented the Jacobi method and steepest ourselves. For the GMRES, we used the version of JAX'
tests, the linear solvers were zero-initialized. Convergence is achieved the relative residuum error using the 2-norm
<(
Auu
[ h
bnlla
(6)
by 1|2
below the threshold c, which number of iterations K as a way to control the refinement of the physics simulation. The iterative solvers return if either the maximum number of iterations K is reached the tolerance threshold
10-5 due single precision. We the maximum
((k) c is met.
B. JACOBI RELAXATION
The Jacobi method is relaxation-type method based on the decomposition of the system matrix A into a strictly lower diagonal L, diagonal D, and a strictly upper diagonal U part such that A = L + D + U. We present its algorithm in
Algorithm 1: Jacobi Method for Solving Au,, by.
Input: Matrix A NxN, vector by RR tolerance maximum iterations K Output: Approximation of the solution un E R2
Decompose L + D + U A Initialize u E R to zeros u [0]
0;
for k 0 1, K do
R++ D-1 [b (L + D)u (4 4
Compute relative residuum norm: g(k+1] || [k+1] balle/ b2ll
c[++1] then
break; end
end
[[+1]]
return u.
https::/jax.rreattheeee Hinalg.gmee gmres. html
16

---

Published as a conference paper at ICLR 2025
As smoothing method, the convergence of the Jacobi method depends the spectral radius of the iterator matrix (Saad, 2003) via
S pL+ R))l
(7)
Loosely speaking, the condition number of the system matrix A affects the spectrum of the Ja- cobi iterator matrix that high condition numbers lead to spectral radii close 1.0, causing slow convergence.
B.2 STEEPEST DESCENT
The steepest descent method follows the gradient of the convex quadratic optimization problem associated with solving linear system of equations. which is also given by the negative residuum. However, in contrast to the gradient descent methods typically found in more general optimization problems (like for training neural networks), the optimal step size for maximum decrease a [k] can be determined in each iteration (Saad, 2003). We present the steepest descent in algorithm 2.
Algorithm 2: Steepest Descent Method for Solving Aus = by.
Input: Matrix A E NNN, vector bj E RR,, tolerance maximum iterations K Output: Approximation of the solution UA RA
Initialize u E R to zeros U) 0:
Set initial residual rh bn for & = 0, 1,......... K do
[0]
Au [0]
Compute step size a
Th A 1
rh rh
[4]
Update solution u [++1] 1]
ut + [1 [K].
[6]
Update residual r []+1] r a Ar k].
Compute relative residuum norm: g( = ||r/ [+++1]
112 /
l12
if [[++1] then
break;
end end
[k+1]
return
u
It can be shown that the residuum norm converges exponentially linear in the asymptotic regime based on the condition number of the system matrix x( (A) via (Saad, 2003)
||
[K++1]
S
(x(A) Ki (AA)
|| [K] 1l.
(8)
+
B.3 GMRES
The General Method of RESiduals (GMRES) builds subset of the Krylov basis associated with system matrix A and finds an approximation to the solution of the linear system via least-squares. Typically, it is restarted rebuild construction of the entire m-diimensionl Krylov subspace and subsequent least-squares solve. We use the "bat ched" mode of the JAX implementation, which can only terminate after a restart but not within a restart.
Krylov basis each iteration. We call iteration the
C MORE DETAILS ON DIFFERENTIATING OVER ITERATIVE LINEAR SOLVERS
For convenience, we restate the simple computational graph involving a linear solver with steps subject to reverse-mode automatic differentiation (AD)
(VoL)* = -0" =0""=(11 ((1" JIlPk(S(-0)) Jrx JPx (xllsc Jsl0.
17

---

Published as a conference paper at ICLR 2025
Algorithm 3: GMRES Method for Solving Aun b₁.
Input: Matrix E NNN, vector bn RN,, tolerance e, maximum iterations K, restart parameter m
Output: Approximation of the solution uA R
Initialize un E R to zeros U/ 0;
[0]
Set initial residual rh bn Au [0]
[0]
Set 8 || [0] ||:
Set e = r/ 00
/8 (first Krylov basis vector);
for k 1.... K do
for j= 1,2, ................................ m do
1.2,.... Compute w Au;; for i 1, ................................ do
Compute has = v; wi
Update w, w hig U4; end
Set hit1, = 11:0 Normalize 4,+1 wji h/z11j
end
Solve the least-squares problem: minimize [|Ber H H 12: [s] + o
Update solution: u [ u
Compute relative residuum norm: g[++1] = [1+[+++1] /12/ /le//lbnll l1: if c[++1] < c then
break;
end
Restart GMRES after every m iterations if convergence not met;
end
return U. [++1]
18

---

Published as a conference paper at ICLR 2025
In this setting, one has to compute three vector-Jacobian products (VJPs). The VJP over the loss function /h 1 Ji Pf :0)) and into the network's parameter space 0" gh Jyle can be straightforwardly evaluated by the AD engine as they consist of explicit operations.
On the other hand, the VJP over the approximative solver 9h T J. If :)) is nontrivial as it requires handling the iterative solver and the assembly routines. We detail the two prominent approaches below. They visually depicted in figure 11.
C.I IMPLICIT DIFFERENTIATION
We first solve an auxiliary linear system for an adjoint variable X (Christianson, 1998)
ATA, An Uh.
(10)
Since the convergence behavior of most iterative linear solvers is dependent on the spectrum of the system matrix and since transposition does not effectively change this, it is reasonable to em- play the same iterator as in the primal solution but with a transposed system matrix and differ-
ent right-hand side +D- AT, ūn). The convergence of the iterates {A/o) x! ,..,,, ,!! ( [0) AAT ūn)) for tolerance 6 can be different from primal the K,-
potentially be different than in the primal solution. In other words, K.
Once X is determined, it directly equals the intermediate gradient the right-hand side b, An- The intermediate gradient on the system matrix arises as the negative outer product with the primal solution A = = Hence, it is sufficient to only save the solution ui no further iterates are required.
[X] from the primal pass;
Then, the reversely propagated intermediate gradient on the input of the physics operator is given by
i JPxls.
g% = b% J2la
A Jalgs
(11)
AAJJJla Ra
under the abuse of notation to consider the gradient matrix A as a vector. Alternatively, this can also be expressed the left double contraction. It must be noted that typically A is only evaluated on the sparsity pattern of A, if it is materialized at all.
In case only the right-hand side by is parameterized, there will be reverse propagation through the matrix assembly A. For example, this was the case for the Poisson and heat emulator examples in section 2.2 and section 2.3, respectively.
Since we use matrix-free implementations for most linear solvers, the primal system is not solved with a materialized matrix but function that is linear in its first argument
(g _ (h gr)-
(12)
In this case, we compute the system matrix' contribution to the previous intermediate gradient via the negative VJP over its second argument
T Ug Jrx lg. oh = =b -X7, J..2lg.
(13)
To employ the iterator & matrix-free version of the transposed matrix, we need to program- matically transpose the function v) which can be done with JAX (Bradbury et al., 2018). Moreover, there is function in JAX to automatically register linear solvers with correct propagation rules?.
More generally speaking, implicit differentiation is powerful because it allows differentiation over various implicit relations simply by solving linear system of equations. Modern AD engines allow the effortless linearization of optimality conditions. This allows for easily registering custom propagation rules by employing matrix-fred Krylov solvers (Blondel et al., 2022). Note, however, that the focus of this paper is the differentiation over linear system solutions in which the primal operation the same as in the implicit propagation rule.
Pttpps:/ / jax. readt nenerr_sol sol
linear
html
19

---

Published as a conference paper at ICLR 2025
C.2 UNROLLED DIFFERENTIAION
Unrolled differentiation of an iterative program is the direct application of standard automatic dif- ferentiation tools its algorithmic implementation. AD unrolls the program's iterations and writes them as individual computational steps.
We find the intermediate gradients on both the system matrix A and right-hand side oa via first computing the intermediate gradient on all iterates using the VJP of the iterator & with respect to its first argument
II I=K-1
J
(o. b.,)))'
(14)
Then we aggregate each contribution using the VJP over the iterator & with respect its condi- tioned arguments
K-1
A T unroll
u
[k]
J0.2 with =0 (A,,,,)
(15)
K-1
unroll
[R] T Up
Jp.3 with =0,, (A,)),
(16)
Since the VJP of the iterator has to be evaluated at primal inputs, typically, all iterates must be stored on the tape recomputed. There are approaches that balance compute and memory (Griewank, 1992). However, for simplicity, we only implement unrolled differentiation by retaining the entire sequence of iterates.
After the intermediate gradients have been obtained, the gradients are further backpropagated via equation 11. While the mathematical description of unrolled differentiation more elaborate than for implicit differentiation, its implementation in AD engines like JAX is easier given the algorithm fully uses differentiable operations.
The non-asymptotic study of Scieur et al. (2022) revealed that the convergence of unrolled differen- tiation can exhibit burn-in phenomenon. We also investigated this for our iterative lincar solvers but found it to be practically irrelevant. While it theoretically can occur for problems with parame- terized system matrices or for Krylov methods in general, we believe that for large enough parameter spaces, the potential burn-in of some gradient components is compensated by the entirety. Moreover, since we applied PRDP using both unrolled and implicit differentiation and observed almost identi- cal savings in both cases, we further conclude that even if there was a burn-in the PRDP approach would be unaffected by it.
D DIFFERENTIABLE SOLVERS
In this section, we describe the discretization choices behind the differentiable solvers. They are based on finite difference approximations on uniform cartesian grids with Dirichlet or periodic boundary conditions. Linear PDEs naturally lead to linear system of equations for which only the right-hand assembly remains a dependency on the prior compute graph, e.g., is built upon infor- mation previous in time. In the case of the nonlinear PDEs, choose a Ficard-based approximation leading to Oseen- like problems (Turek, 1999). Those have both the right-hand side as well as the system matrix depending on the prior parts of the compute graph.
The discretizations are implemented matrix-free in JAX (Bradbury et al., 2018). Matrices are only materialized when needed for smoothing methods or for direct decomposition- based solvers.
Once fully discretized and foroouulaeed, each problem cast in the standard form a linear system of equations
(17) using the assembly routines A- Aggn) and b, (gg)). In the simplest case have directly 9t = 0 but other setups like information from a previous time step 9t. = w [[[]]] are common. Then, define the physics operator P as the function mapping from the prior point in the compute graph
Aun. bR:
20

---

Published as a conference paper at ICLR 2025
f
A²
A
A
Figure 11: The iterative solution of a linear system of equations spawns a long sequential compute graph that constitutes the primal pass (in black). The step from prior variable gh. to approximate solution ! can be capsuled as the physics operator Pk (in blue). Under reverse-mode AD, differentiation over PK (blue curvy line) be done in ways. Implicit Differentiation spawns an auxiliary iterative linear solve using the transposed matrix AT (in purple). On the other hand, unrolled differentation opens the black box and reversely propagates through each iterator step & (in yellow). Both approaches yield intermediate gradients on system matrix A and b, which are subsequently backpropagated into gn via the VJP over their assembly routines.
gh. to the solution of the linear system of equations
P(gn) == un A b = (A(g*)) ((gh)
(18)
P is given without subscript this refers to an exact solver of the linear system. This can be a fully converged iterative solver. Either way, the
direct method based on matrix decomposition residuum norm of this solver's result guaranteed to be below the relative tolerance threshold e via
|(gg) ) B(gh)ll <e.
(19)
([3(g)]]2
If not mentioned otherwise, we floating format. An approximate solver with K iterations is written Px(-). The iterative linear solvers we employed are introduced in section B.
10-5 because we exclusively work in single precision
D.1 POISSON EQUATION IN 1D
We solve the Poisson equation the unit interval with homogeneous Dirichlet boundary conditions and parameterized right-hand side
d² ( u())= = u(1)
(20)
dx22
For tinite difference discretization, the domain Ω (0, divided equidistantly into N grid points, out of which two are related to the prescribed value on the boundary and can, hence, freedom that make up the discrete solution vector un E N Equally, the right-hand side is discretized at the same points and negated, yielding b;. At index i interior to the domain, the second derivative is approximated with the three-point stencil
be ignored. As such, there are degrees
Ax 2 (u-1 2uq + (4+1) = -b;,
(21)
21

---

Published as a conference paper at ICLR 2025
with the spacing Ar
This leads to the system of linear equations
N41-
2
1
0
0
0
M1
b1
2
0
0
U2
b2
0
1
0
0
W3
b3
(22)
(Ar}
:
0
1
2
bN
=:
in which the first matrix is due to the discretized Laplace operator in one dimension. We denote &1 (the tilde is to distinguish it from the Laplacian matrices on periodic boundaries of the following sections). The subscript is to indicate the mie-dimensional setting. It is tridiagonal and solely defines the system to be solved with A L1
Au 3bbb).
(23)
with the right-hand side assembly function just element wise negating the input. The system matrix is not parameter-dependemt, i.c., A(g) = L.
D.2 HEAT DIFFUSION IN 1D
We consider the time-dependent diffusion equation in one dimension on the unit interval under homogeneous Dirichlet boundary conditions
Du 01
8u O.2
ult, 0) (0.00)=0 - - 00((t,))
(24)
Similarly section D.I. we equidistantly divide the domain into N grid points. The matrix associated with the discretized second derivative in one dimension is again
2
1
0
0
0
E =
0
1 -2
0 0
(25)
(Ar)²
0
2
With it equation 24 can be discretized in space via the method of lines
dur vL un.
(26)
dt
Applying an implicit Euler discretization to the time derivative yields
[++1]
(c]
(27)
A
with two subsequent time levels the state vector. This can be rearranged into the standard form of linear system equations as
VALE
[++]] u u1 [
(28)
This is also referred to the backward-in-time central-in-space (BTCS) discretization of the dif- fusion equation. Advancing the state by one time level requires solving linear system with the constant system matrix A and with the assembly function 3(-) simply being the identity. Hence, the physics operator maps to the next state in time with
h
[++1]
P(u 11)
(29)
22

---

Published as a conference paper at ICLR 2025
D.3 HEAT DIFFUSION IN 2D
The equation for heat diffusion in two dimensions on the unit square with homogeneous Dirichlet boundary conditions reads
Ou 01
80uu 2r
8ºu oy²
u(1,0,v) - =((, 1,2), u(t, x, x, - =((, x.
(30)
The Laplacian matrix in two dimensions can be written in block structure using the one- dimensional Laplacian L1 and appropriately sized identity matrices I
(Ax)2C L1 21
I
I
(Ar) L1
21
I
La
1 (Ar)²
0
I
(Ax)
L1 21
0
(Ax)² L1 21
(31)
With this two-dimensional Laplacian matrix, the state vector E RN is advanced similarly via solving
- (-Ii
ui u [++1]
(32)
=AERNi
Similar to the one-dimensional BTCS scheme, the system matrix is constant. The right-hand side assembly is again the identity.
D.4 HEAT DIFFUSION IN 3D
The equation for heat diffusion three dimensions reads
Ou
(82u
yz
82u
+
(33)
Ōt
2x
Qy²
0z
Au
We again use homogeneous Dirichlet boundary conditions on all six sides of the unit cube
utt, 0, y, =) u(t, 1, u(t, x, 0, 0 u(I, 1, u(I, v. y, 0) 0 u(t, or, y, 1). We consider the unit-cube Ω (0, 1)3 with N degrees of freedom per dimension, in total N3.
(34)
The Laplacian matrix in three dimensions can again be written a block structure using the two- dimensional Laplacian L2 and appropriately sized identity matrices I
( 21
I
0
0
0
I
(Az) Le 21
I
0
0
0
0
E3
1
0
I
(Az)² C2 - 21
I
0
0
0
(
0
0
0
0
0
I (Ax)² 21
(35)
With this three-dimensionnal Laplacian matrix, the state vector up E NR is advanced similarly via solving
(-vAt(s)
[++1]
u.!!
(36)
AERN x
Similar to the one-dimensional BTCS scheme, the system matrix constant. The hight-hand side assembly is again the identity.
23

---

Published as a conference paper at ICLR 2025
D.5 BURGERS ID
The Burgers equation on the one-dimensional unit interval in non-conservative form with periodic boundary conditions reads
Ou du 01
8u
+ u
u(t, u(t,))=u(t,11 1).
(37)
Ox
d=
We will discretize the domain into N + grid points. Under periodic boundaries, only of the boundary points can be eliminated. By convention, we choose the right-most point, leaving N degrees of freedom (including the left-most point). The matrix associated with the discretized second derivative in one dimension now reads
2 1
0
1
-2
1
0
0
L1 = (Ar)²
0
1
-2
0
(38)
0
-2
It differs from the former Laplacian matrix in that it is not exclusively tri-diagonal but also has entries in the top right and bottom left corners. Additionally, we now have Ar 1/N.
The convection term requires special treatment because of its nonlinearity and the advection char- acteristic of the first derivative. Let Fi and B₁ represent the forward or backward approximation of the first derivative in one dimension on periodic boundaries, respectively, via
0
0
0
1
0
0
1
1
0
0
0
F =
0
-1
0 0
B1 =
0
-1
0
0
Ar
AT
0
1
(39)
Again, note the element in the corner entries of the matrices. Then, can build upwind differ- entiation matrix based on the winds wp.
8 (w)) + WT
81 (w)) + we o
(ww, diag
max
B, diag
max
Fr. (40)
2
positive winds
negative winds
Deducing the positive and negative winds from neighboring averages (using the periodic forward shift 8-1 and backward shift 81 operators) is necessary to have correct movement the winds change sign the domain. If we the discrete state vector ut, winds wh, can discretize the continuous equation via the method of lines as
du: di
L'(uu)u = vLluu.
(41)
Naturally, the spatial discretization of nonlinear PDE leads system of nonlinear ODEs. To fully resolve the nonlinearity, one could resort to a Newton-Raphson or a quasi-Newton method. However, for simplicity, will apply the trick linearize the upwind matrix using the state previ- ous in time (Turek, 1999), which gives
[++]]
h
u.
[t+1]
2L [++1]
(42)
At
This can be rearranged into the standard form as
I + A AtvCi
|0+1]
U 18]
(43)
=A=A(u!")
As such, the system matrix is dependent on the previous state in time, i.c., it is dependent on the previous variables in the compute graph. The right-hand side assembly routine is again the identity.
24

---

Published as a conference paper at ICLR 2025
However, different from before is that the system matrix A is now asymmetric, which necessitates special linear solvers.
Iteration re-assembly and solution is possible but omit this for simplicity, accepting that the nonlinear residuum is not fully converged. This introduces an error order ((t)), which is acceptable since temporal discretization is first order. It must be noted that despite the nonlin- ear residuum not fully converged, the linear residuum associated with the linearization will be, assuming use the converged solver P.
D.6 BURGERS 2D
While we did not investigate any experiments with two-dimensional Burgers solver, will still present here as it naturally helps with understanding the two-dimensional Aavier-Stoker solver.
In two dimensions, the continuous solution function the Burgers PDE becomes vector-valued with two channels. Using symbolic notation, we can write
du + (u Vju = vdu.
(44)
alternative way to write the two-dimensional Burgers equation its two components
Ou
Oul
Oul
8uu
82uu
01
01
0r
+ u2
Oy
v
dr
+
Qy²
(45)
Ou2
Ouz
Ou
@2
0u
Ot
+ 41
dr
+ u2 Oy
v
Ou²
0y²
Our domain Ω is again the unit square.
Under two-dimensional biperiodic domain, both the right-most and the top-most boundary nodes
are eliminated, see figure 12a. We consider the unit-square ? = (0, 1)² with equally many degrees of freedom per dimension, total 2
The Laplacian matrix in two dimensions can be written in block structure using the one- dimensional Laplacian ć and appropriately sized identity matrices I
( L1 21
I L1 (
0
0
0 0
I
I
21
I
La
1
0
I
(Az)
C1
21
I
0
0
0
(Ar)²
0
0
0
0
(Ar)² L1 21
(46)
Note the identity blocks in the top right and bottom left corner.
Let F2,1 and F2.2 be the forward derivative operator in two dimensions in the direction of dimen- sions one and two, respectively. Moreover, B2, and B2 are the same for the backward derivative operator. Hence, we can build an upwinding operator for the first direction
S (w1,) + wh.1 o)
81 (1 + Wh.1
22,1 (w)) diag
max
+diag
max
F2
positive winds
negative winds
(47)
which only needs the winds in direction one. Similarly, we get for the other direction
8 (-1( + whi2 0)
81 (wh,)) + wh.
(w)) diag max
B
max
F2.
2
positive winds
negative winds
(48)
25

---

Published as a conference paper at ICLR 2025
Let us combine these into one joint upwind discretization
(2(w)) = 22,1 (2rh (2,2(2n))
(49)
This allows for discretizing the component-wan equation 45 first in space via the method of lines, afterward similarly in time as before get the Oseen problem
[++1]
h,1
uni1
) (++1)] = 2C2u [++1] (2(
+
At
(50)
[++1]
u h.2
u
+ I2(u, juk [e+1]] v u
At
Note that for the discretization of each component, need both components from the previous time level to assemble the upwinding matrix. We can re-arrange and write the system jointly as
I+
Atl 2t
[0) u1
Atv
[++1]
(51)
0
Ar 2 AIr
11,2
[++1]
AA
The assembly routine for the system matrix yields block-diagonal structure while the right-hand side assembly is just the identity.
D.7 COUPLED NAVIER-ST -TIEESSSOKES IN 2D
We solve the incompressible Navier-Stokes equations which can be written in symbolic notation
du + (u - -
(52)
Øt V 0.
(53)
The domain is again unit-square with bi-periodic boundary conditions. This is a system of partial differential equations for three unknowns, which are the two components of the velocity u and the scalar pressure field The pressure acts constraint to enforce incompressibility given by the second equation. The coupling velocity and pressure can be challenging, but a simple way for finite difference discretization is the staggered grid (Harlow & Welch, 1965). As depicted in figure 12b, it also accounts for bi-periodicity by ignoring the top-most and right-most grid points but uses different locations
the three unknowns. In this setting, there equally many degrees of freedom per variable and direction. Each variable contributes 2 entries. Hence, in total, there 3.V²
Evaluating the convection term for both derivative directions requires mapping between the two staggered representations of the velocity grid. Let denote the mapping operator Mj that moves all variables to the grid representation of the first velocity component. The operator M2 does the same for the second velocity component. These operators can be easily realized via bi-linear interpolation, which in the uniform cartesian grid simply amounts to the average of the four neighbors. Then the convection operator for velocity components one and two are (2 (ux)) and (M)M)(u))) respectively.
The discretization of the pressure gradient and the velocity divergence requires a mapping between the velocity and the pressure representations. It turns out that forward derivative F2,1 or F2 also maps from the velocity representation to the pressure representation. Hence, we can define the divergence operator
velocity components D F2,1 and D2 F2,2, respectively. Vice versa, the gradient operators mapping from pressure to velocity representations are the backward differences, i.e., G B2,1 and G2 B2,2-
26

---

Published as a conference paper at ICLR 2025
O
(b) Backward Staggering for bi-periodic domain. Hor- izontal arrows denote the location degrees of freedom for the x-velocity are saved. Vertical arrows corre spond to y-velocity. Red circles are for pressure.
(a) Bi- Bi-periodic (collocated) grid. The top-most and right-most vertices have no degrees of freedom due to periodicity.
Figure 12: Arrangement of degrees of freedom on two-dimensional bi di-periodic domains.
The method of line discretization with the linearization of convection matrices around the previous state in time then yields
[++|] 1
u
h.1
2
+ (w)
[++1]
[[+1] G G
[++1]
(54)
A
[++|]
[t]
1.2
2h
+
+22MM221
(4) )20
[0+1] = Gru + v UL
[1+1]
(55)
At
Duu + Dun
[++1] = 0
(56)
We can write this in matrix form
[++1]
I+ (M! AtvC2
Gi
u
I+ ATT (Mg(()
AtvL2
G2
U1.2
[++1]
u
A.2
D
D2
2b h.3
[++1]]
0
A
sx 11+31
(57)
which reveals the saddle-point nature of the coupled Navier-Stokes system given by the zero block in the bottom right (Turek. 1999). Many popular solution techniques to the Navier-Stoker equations, like PISO (Issa al., 1986), can be interpreted as efficient preconditioners as argued in Perot (1993). For simplicity and to not get nested iterative linear solvers, we solve the coupled system without further modifications, which reasonable for the low employed resolution despite the considerably high condition number.
this coupled system
Similar to the Burgers simulators, the system matrix A is asymmetric and needs to be re-assembled from information given by the previous state in time. However, only the previous velocity compo- ( nents in time | and u are needed. The previous pressure state u does not affect the matrix assembly. It is also not relevant for the right-hand side assembly function 3(.).
E MORE DETAILS ON PRDP PARAMETERS
Stepping Threshold step The stepping criterion defines the threshold at which the validation met- ric is considered sufficiently flat before the physics refined. This value should not exceed as
27

---

Published as a conference paper at ICLR 2025
values greater than I would trigger refinement only when the enror increases. Higher values (less than 1) allow the metric to flatter before the physics is refined. In other words, a lower value enables faster progressive refinement and trades off PR savings. Refinement can also be made faster by making larger steps in refinement AK. In practice, we found that Tstep values between 0.9 and with AK sufficed in all our network training runs.
Stopping Threshold Tstop The stopping criterion the checkpoint ratio governs how strictly the stagnation in network accuracy should be over different refinement levels. A lower value implies stricter check, and progressive refinement would stop in lower max region, and vice versa. In other words, a lower value of Ssup enables a more aggressive strategy for IC savings.
The grace window of intervals controls the length of lookback calculate the stepping ratio r. This is helpful in the case of strong epoch-to-epose oscillations. On the other hand, the effect of oscillations on the stepping ratio mitigated through exponential smoothing. We found window of 3 worked for most scenarios. The grace window goes hand-in-hare with the stepping criterion a longer window requires stricter (smaller) Tstep values.
F EXPERIMENTAL SETUPS
An algorithmic description of Progressively Refine Differentiable Physics is given in algorithm 4. Our choice PRDP parameters based on section are listed in each problem setup as follows. Our scheduling via changing the number of inner iterations via AK each time PRDP triggers refinement. In figure 13, we demonstrate that this is identical to scheduling with relative tolerance thresholds.
Validation Metric For test cases based on the Heat diffusion, Burgers equation, and Navier Stokes, the validation metric we use for PRDP is the solution's 's error against reference solution specific time step t, normed over space, and normalized against the reference solution. This error is mean-squared over the validation set. Mathematically, this can be written as
||ue
2 1 112
Validation Error
Σ SEV
(58)
M
||
where:
u | the predicted solution time step for validation example i.
|t,.r. the reference solution at time step for validation example u,
|| ||: the norm over the spatial domain,
V: the validation set.
The corresponding specifics for the Poisson inverse problem are explained in F.I.
Seed statistics For the neural network training setups, we conducted trials with 10 different initial- ization seeds. Each uses different network initialization and different stochastic minibatching. but the same data. The results shown in the work visualize the mean over these runs (shown as solid lines), along with its variability (represented by the shaded area). The shaded areas indicate the range within one standard deviation of the mean.
F.I POISSON EQUATION INVERSE PROBLEM
The Poisson inverse problem is our simplest test example that incorporates differentiable physics into a learning pipeline for a doubly convex problem optimizing in a low-dimensional parameter space. We set up the discretized Poisson problem as described in section F.1 with degrees of freedom N 30. The resulting linear system is parameterized, similar to section 2. The iterative linear solver applied to this system represents the physies Px in this problem.
28

---

Published as a conference paper at ICLR 2025
Linest sabeer stevation
10-4
primal cusvergetoo
FRDP
10%0
prandd cosinnggtooo
madter
FROP
20
60
100
Sprch
Figure 13: PRDP refines the physics by progressively increasing the solver iterations throughout training until training accuracy in terms of a validation metric stagnates. Progressively increasing the number of iterations is equivalent to progressively decreasing the tolerance threshold, as shown here for the example of the training the Navier-Stokes emulator of section 4.4,
Algorithm 4: Determine Whether Refine Physics
Input: Tstep 0.95, Tstop 0.9, 5 = 3, Ae = 6
Data: {L
Output: Boolean indicating whether to refine [[e]
(i(e)) = =ema Jalee Ae)
Smooth with Exponential Moving Average
Lall /
change to epochs prior
r > Taep then
rc i
change previous plateau
re < Taop
V re > 1.0 then
c 1_1!
current plateau Retine!
return True
else
L return False
b Do not further refine becouse we reached Kaax.
yields IC savings
else
return False
Do not refine to keep coarse physics for economic reasoms, = yields PR savings
(a) 1P
Right-hand side
Solution
(b) 3P
0.10
7.5
5.0
0.050
0.05
Right-hand handdhhnddd side
Solution (u]
2.5
0.025
Right-hand
Solution
0.00
0.0
0.000
-2.5
-0.025
0.05
5.0
-0.050
-010
-7.5
0.00
0.25
0.50
0.75
1.00
0.00
0.25
0.50
0.75
1.00
Figure 14: The Poisson PDE can be interpreted as the deformation of a thin string subject to load. We consider two scenarios. In (a), the load (=right-hand side) only consists of one sine mode whose amplitude is scaled by single (IP) parameter. The second case (b) uses the first three sine modes with parameterized amplitude each.
29

---

Published as a conference paper at ICLR 2025
For simplicity, keep the system matrix A constant and only parameterize the right hand side by = (()). The map a sum of the first sine modes in the unit interval whose amplitudes are given by the parameter vector 0 E IRP.
P ((0) Σ (,sin(2imr))
We design two inverse problems, one with single sine mode (P = 1), and one with three sine modes (P 3). A qualitative example is given in figure 14.
[K] The outer optimization's objective the MSE emean-square error) in the physics solution of u against a reference solution WA,r- The reference is generated by direct solution of the linear system at a reference parameter value 0,. In other words, unrr = A b- where b ((,)) Optimization performed using gradient descent algorithm.
One-dimensional parameter space We use 0, 2.0, and an initial guess for gradient descent Inn 5.0. 170 update steps are performed with a constant learning rate 275. For PRDP, we set the control parameter values Ssep 0.92, Tstup 0.98, 5 2. Training linear solver iterations. At every refinement, it was incremented by AK = 10. relatively high value for AK and aggressive refinement strategy with relatively smaller value of Tatep are suitable for cases where Kmax is significantly higher than the number of outer iterations. In this case, the physics converges K, 600 (which is also K max due the double-convex problem). PRDP was successful at enabling 33% solver iteration savings
started with K, 25
dhree-nimennional parameter space We use 0, = [0.62, 1.86, 5.1]. and an initial guess for gra- dient descent Oinic [3.0,3. 3.0, 3.0). 750 update steps performed when using Jacobi the inner problem and 800 when using steepest descent, with a constant learning rate of 3500. The physics converged with 450 240 inner iterations when using Jacobi steepest descent, respectively. In this case, a less aggressive PRDP strategy was suitable since the update steps were much higher than K.. Hence, we found control parameter values Seep 0.99, Tstop 0.99, å 2. with AK and Kú 1 was suitable. PRDP was successful at enabling 40-50% savings for both Jacobi and steepest descent with both unrolled and implicit differentiation.
The results of these experiments are presented in figure 15.
F.2 HEAT DIFFUSION AUTOREGRESSIVE NEURAL EMULATOR TRAINING
In this setup, a neural emulator learns the diffusion equation solvers presented in sections D.2 and D.3. For the physics, we use a diffusivity of 0.001 and time step size of At 1, see equation 5. We discretize the space, into 32 grid points (including boundaries) in each dimension, which, after applying homogeneous Dirichlet boundary conditions, results of freedom in the 1D) and 2D cases, respectively. In the 3D case, we have 22 grid points including boundaries, ic. 20 interior grid points, corresponding to 8000 degrees freedom. The resulting linear system is solved using the Jacobi method. This constitutes the physics operator Px. We employ our own implementations of the Jacobi solver. The initial conditions are generated as a truncated Fourier series. For ID, we sum the first 5 sine and cosine modes defined the unit interval. For 2D, we use the products the first 5 sine and cosine modes defined as:
and 900 degrees
uo( v) Σ (an sinnsiinnnaaa sin((nmy) + bn cos(2nmzz) cos (2nmy)
On sin(2n cos((2n))) + d, cos cos(2nmz)) sin (2nmyy)
We similarly extend this procedure to 3D. All amplitudes are randomly sampled from uniform distribution 1)) A qualitative example of the dynamics is given in figure 16 for ID, figure 17 for 2D, and figure 18 for 3D.
Training Procedure We intentionally chose simple architectures the focus of this work is not architecture but training methodology. For 1D and 2D cases, a multilayer perceptron (MLP) fo is trained to emulate Px. For the 1D case, the ML. has 3 hidden layers, each of width 64, and for the
30

---

Published as a conference paper at ICLR 2025
physics solver iee alionio
Fully converged
# physics solver iterations
Fully converged
300
600
PROP
300
600
PRDP
Parameter relative arror
teraatons
600
PR savings
physics salver Roratiens
600
savings
Parameter relative
10
400
10
10-3
200
200
10
125 170
125 170
170
125 170
Outer iteration
Outer iteration
Outer iteration
Outer iteration
(a) P. Jacobi, Unrolled diff.
(b) P. Jacobi, Implicit diff.
physics solver iterations
Fully converged
physics solver Iterations
Fully converged
225
450
PRDP
225
450
PR savings
PR savings
Parametor ertor
Parameter relative error
# phycics zolver Iteraaiions
400
400
10-²
10-2
10-4
10-
750
750
750
750
Outer Iteration
Outer iteration
Outer Iteration
Outer Iteration
(c) P. Jacobi, Unrolled diff.
(d) P. Jacobi, Implicit diff.
physics solver iterations
Fully converged
* physics solver iterations
Fully converged
120
240
PRDP
120
240
PR savings
PR savings
Pamantterralativn enor
# physics solver iterations
Paramelee ler aive error
100
10-2
10-2
100
100
10-1
physic
800
800
800
800
Outer Iteration
Duter iteration
Outer Iteration
Outer Iteration
(e) P. Steepest Descent, Unrolled diff.
(f) P. Steepest Descent, Implicit diff.
Figure 15: Results of the Poisson equation inverse problem visualizing the parameter error during training and the savings enabled by PRDP for both the parameter and 3 parameter setup using dif- ferent combinations of physics solvers (Jacobi, steepest descent) and solver differentiation methods (unrolled, implicit).
31

---

Published as a conference paper at ICLR 2025
IC
step
2 steps
IC
IC
IC
2
1
ufx
0.0
0.5 1.0
0.0
0.5
1.0
0.0
0.5
1.0
x
x
x
Figure 16: Example of three different trajectories over three temporal snapshots for the one- dimensional heat equation.
2D case, 3 hidden layers each of width of 3000. Corresponding the physical problem's degrees of freedom, it has 30 input and 30 output nodes (900 for 2D). For the 3D case. we train classic ResNet that supports dimensional input and output tensors with channel. It contains 6 blocks, each with 32 hidden channels. Its convolutions implement homogeneous Dirichlet boundary conditions. All networks are initialized with random weights. Non-linearization is enabled by applying ReLU activation after each hidden layer.
The compute graph consists of the neural emulator fo stepping from an initial condition the first time step solution, and the physics operator Px stepping from the neural network's solution to the solution at the second time-step. The nMSE between this solution the second time-step against reference solution forms the loss function for training. This mixed-chain approach (depicted in figure 22 allows gradients to flow through Pи and lets the network training to be informed with the physics's dynamics. We implicit differentiation for the 1D case, and unrolled differentiation for the 2D and 3D cases. The custom VJP rules of the physics use the same iterative linear solvers as the primal solve and resolved the same number of iterations the primal, explained in section 2.
For the outer (training) problem, et al., 2020). The learning rate is scheduled as exponentially decaying with an initialization of 10-3 decay rate of 0.94 for ID (0.9 for 2D), and 100 transition steps, while for 3D, the initialization is 10 with a decay rate of 0.92 and 100 transition steps. We train in mini-batches of 25 samples per iteration and for total of 70 epochs in the ID case, and 100 epochs in the 2D and 3D cases.
the Adam optimizer from the Optax library (DeepMind
Data Generation The reference solution at the second time-step is computed by two applications of the physics operator the initial condition, where the iterative solver is replaced by direct solver. Hence, the iterates [1|]] == converge to the reference solution. 205 samples are generated for training, with a traintvali split of 200:5.
PRDP PRDP controlled using the validation metric of equation 58 computed at t = 2, and refinement is applied to the iterative linear solver Px that steps from ( to ula), For the ID case, parameter values Tsepp 0.98, Tstop 0.9. 8 3, exponential smoothing is performed over a time window of 8. The 2D case uses fewer epochs, hence a less aggressive refinement using Tsep 0.9 with smoothing window of 6 was suitable. In the 3D case, which uses 100 epochs, used a less aggressive refinement with Tseep = 0.97, similar to the 1D case, along with Tstop = 0.9, 5 = 3, exponential smoothing performed a time window of 6.
F.3 BURGERS EQUATION AUTOREGRESSIVE NEURAL EMULATOR
We solve the discretized Burgers equation presented in section D.5 on the unit interval with 0.001 and time step size A/ 0.01. The spatial domain discretized into 257 points in-
32

---

Published as a conference paper at ICLR 2025
1=0, t=0
i=1, t=0
i=2, t=0
|=3, t=0
t=0
i=5, t=0
3
1=0,
i=1,
1=2, t=1
|=3,
1=0, t=2
t=2
1=2, t=2
1=3, t=2
1=4, t=2
|=5, t=2
-2
-3
1
0
1
1
Figure 17: Example of six different trajectories over three temporal snapshots for the two- dimensional heat equation.
i=0, t=0
i=1, t=0
i=2, t=0
i=3, t=0
i=4, t=0
i=5, t=0
1.00
0.75
0.50
t-1
i-1,t=1
i-3, t-1
i-4, t-1
i=
0.25
0.00
-0.25
1=0, t=
t=2
144,
i= t=2
0.75
-1.00
Figure 18: Example of six different trajectories over three temporal snapshots for the three- dimensional heat equation.
IC
IC
0.75
space
solutii oolutonn
0.50
0.25
012
5
10
15
20
012
5
10
15
20
012
10
15
20
time step
time step
time step t
Figure 19: Three example trajectories of the one-dimensional Burgers equation.
33

---

Published as a conference paper at ICLR 2025
cluding the boundaries, which results in N 256 degrees of freedom when applying a periodic boundary condition. Since the system matrix is now asymmetric, we use the GMRES linear solver (B.3) with restart set to 2. Although much higher restarts are possible and preferred in prac- tice, we used fewer restarts to emphasize the difference between network accuracies trained using different number of GMRES iterations since test problem was relatively easy solve.
Data Generation The initial conditions are generated as a truncated Fourier series as described in section F.2. The first 20 sine and cosine modes defined over the unit interval are summed, with their amplitudes sampled from a uniform distribution (11,,, 1). Qualitative example trajectories are given in figure 19
Training Procedure We train neural network fo to learn the Burgers stepper which represents the physics P. We use a classic ResNet architecture (He et al., 2016) consisting of 6 blocks, each with 32 hidden channels, and ReLU activation. The convolutions implement the physics' periodic boundary condition by circular padding. The weights are randomly initialized. We use our own implementation of the architecture using the Equinox library (Kidger Gareia, 2021).
The loss function for training the network is defined similarly to the mixed-chain setup of section F2. The ResNet steps from an initial condition to the first time step solution, and the physics Px steps from the first to the second time-step solution. The nMSE error between this solution at the second time-step against a reference solution forms the loss function for training.
The reference solution at the second time-step is computed by two applications of the physics, with a direct solver replacing the iterative solver. 205 samples are generated for training, with a validaton split of 200:5.
We use the Adam optimizer from the Optax library, with an exponential learning rate decay using an initialization of 10 10-3, decay rate of 0.7, and 100 transition steps. We train in mini-batches of 25 samples per iteration.
PRDP PRDP is controlled using the validation metric of equation 58 computed at t 5, and refinement is applied to the iterative linear solver Px that steps from [1) to ! This presents an interesting application scenario, where the training objective uses a relatively short temporal rollout of 2 steps while PRDP favors generalization over a longer rollout at the fifth time step. We use PRDP parameter values Ssep 0.9. Tstop 0.9, and 8 Due to high oscillations in the validation metric, we used exponential smoothing with a longer window of 15.
F.4 NAVIER STOKES NEURAL-HYBRID CORRECTOR LEARNING
In this example. we train network to learn to correct for discretization errors in the solver described in section D.7. Such heural-hybrid methods apply to correction learning setups as investigated by Um et al. (2020). Our problem setup is visualized in figure 23.
We distinguish between spatially- coarse solver, Ps.K. and a spatially-fine solver, P, (following the terminology of source and reference from Um et al. (2020)). In our test problem, Pa. operates on a solution manifold that has half the spatial resolution of P,. We use 97 spatial degrees of freedom for P, (i.e., a total of 291 degrees of freedom considering the three components of u ), and 48 for Pa.K (i.c., a total 144). Both solvers are set up for v = 0.0001 and At = 0.1 and employ the GMRES solver (section B.3) with restart set to 8.
Data generation To generate initial conditions for the data, we initialize a solution field in the manifold of P, with each component sampled from the standard normal distribution. This random field is passed through low pass filter, then normalized to have zero mean and standard deviation of 1, and finally passed through an incompressbbility projection, giving uy Data trajectories {u, y are produced by repeated applications of P, (uses a fully-converged GMRES) ! These tra- jectories are then downsampled to the source manifold {ul] 17 205 traeectories are generated for training, with traidtinln split of 200:5. Qualitative example trajectories are given in figure 20.
34

---

Published as a conference paper at ICLR 2025
1=0, Up. t=0
1=0, U1
Up.
41
1=2, Ub.
1=2, 01.
Ug
U1
1.0
1.0
0.8
0.8
i=0, Un.
i=0, U1
Un.
i=1, Up
i=2,
i=2, U1.
0.6
0.6
0.4
0.4
i=0, Up. t-2
i-0, U1
Up.
i-1, Un t-2
1-2, Up. t-2
i-2, U1
0.2
0.0
0.0
0
0
1
1 0
1 0
Figure 20: Three example trajectories of the Navier-Stokes scenario with the two velocity compo- nents each.
Figure 21: The compute graph and the reverse pass associated with the Poisson inverse problem.
Training Procedure The neural-hybrid model operates in the source manifold. We use a classic ResNet architecture (He et al., 2016) consisting of 3 blocks, each with 64 hidden channels. The convolutions implement the physics' periodic boundary condition by circular padding. The weights randomly initialized. We implement the network architecture using the Equinox library (Kidger & Garcia, 2021). )..
The loss function is defined the sum of MSE errors of the first and second time-step solutions. We use the Adam optimizer from the Optax library, with a cosine learning rate decay schedule using an initialization of 10 and decay steps 800. We train in mini batches of 25 samples per iteration and for a total of 100 epochs, i.e., 800 update steps.
PRDP PRDP is controlled using the validation metric of equation computed on the x-velocity at = 5. We used parameter values Tsee 0.98, Tstop 0.9, and 5 3. For the exponential smoothing, we used window size of Since this problem setup has 2 executions of Px (refer figure 23), three options arise for applying progressive refinement a) progressively refining only first execution, b) progressively refining only the second execution, and c) progressively refining both executions. We opted for in order to gain savings from both solvers. In general, correction learning setups would have as many physics executions as the number temporal unrolling steps in the loss function.
35

---

Published as a conference paper at ICLR 2025
Figure 22: In the mixed-chain approach an emulator prediction is chained with an application of the differentiable solver, effectively the training loss computed over two step unrolling. Reference data is computed with full refinement, only the physics of the second step solver is refined with PRDP. The yellow reversely pointing error indicates the differentiable physics over the approximate solver.
XX
1-10
(a) Sequential Correction Sctup
(b) Two-ster stepstep supervised unrolling
Figure 23: The neural hybrid emulator is at the place of "NN" in the right figure. Each pagproppgation-thouughh
also triggers the differentiation over the physics. Note that even there is no differentiation over the first neural-hybrid emulator application (i.e., differentiation over the first coarse solver usage), it is also affected by PRDP because we also limit exactness in the primal execution. Here, could think of it "Progressively Refined Physics" instead of "Progressively Refined Differentiable Physics".
36

---

Published as a conference paper at ICLR 2025
G SUPPLEMENTARY EXPERIMENTAL RESULTS
G.1 PRDP SAVINGS IN TERMS OF COMPUTE TIME
In figure 24, we present wall-clock training times with PRDP and with fully converged differen- tiable solver. For the smaller ID problems with small system matrices of the order of z 100 x 100, other overhead the compute graph dominates the training runtime. As such, the reduction in inner iterations does not translate to savings in wall clock time. Starting with 2D problems leading to larger matrices, decreases training time converge to the reductions inner iterations. For exam- ple, in the 3D heat emulator scenario, 81% reduction in inner iteration count aligned with wall clock time savings of 78%. As the share iterative linear solves in the total training cost increases, e.g. in nigh-dimensional problems, so do the savings enabled by PRDP.
G.2 NEURAL NETWORK PERFORMANCE. PRDP VS. BASELINE
The PRDP method is aimed at reducing the end-to-end computational cost of training neural net- works with physics solvers the loop, without significantly affecting the trained neural network's performance. This is depicted qualitatively in figure 6. In table 1, we provide the mean validation er- ror values at the last epoch of training along with standard deviation over seeds used for network initialization and stochastic minibatching.
30
3.0
50
Converged
Training Time [S]
25
2.5
100
40
PRDP
80
15
1.5
60
20
10
20
10
0.0
Heat 1D
Burgers (1D)
Heat 2D
Navier Stokes (2D)
Heat 3D
Figure 24: While all experiments reduced the number of inner iterations across the outer optimiza- tion (see section 4 bar charts), savings training time become noticeable the larger and more com- plex the setup gets. Once the problem is of sufficient size, linear solve times dominate the training compute graph, and hence, PRDP can deliver the highest amount of savings wall clock time. Eventually, savings in training compute time converge against the reductions in inner iterations.
Table Numeric values of the validation accuracy of the neural emulators and neural teral-hybrid emula- tors according to the experiments of section 4. The mean and standard deviation are computed based on ten different random seeds that affected neural network initialization and stochastic minibatching.
Outer training task
Inner physics
Val. error with Converged Physics
Val. error with PRDP (x10-2)
(x10-2)
Neural Emulator
Heat 1D
2.1 0.39
2.1 0.38
Neural Emulator
Heat 2D
1,4 0.08
1.1 0.06
Neural Emulator
Heat 3D
0.4 0.4 0.09
0.5 + 0.10
Neural Emulator
Burgers
5.5 + 1.38
6.2 : 2.13
-obral-Hybrid Emulator
Navier Stokes
0.6 + 0.05
0.7 0.04
37

---

Published as a conference paper at ICLR 2025
Validation error of trained network
100
Converged PRDP
% Savings in Solver Iterations
10
60
Total savings
PR savings
40
IC savings
10-²
104
10a
10a
10a
Network Parameters
Netwark Parameters
Figure 25: As expected with neural networks, the validation error falls over the number of param- eters due larger network indicating a higher expressiveness. However, PRDP savings remain unaffected.
G.3 PRDP SAVINGS VS. NEURAL NETWORK EXPRESSIVENESS
While could not identify single cause for the IC savings observed when solving the outer problem of learning in the high-dimensional, nonconvex weight space of neural network, we hy- pothesized several contributing factors in the last paragraph of section 2.3. One hypothesis that the inherently approximative nature of machine learning models plays role. Specifically, limited model expressiveness such as insufficient parameterization might contribute to IC savings. In this view, more expressive model would result in reduced IC savings.
To test this hypothesis, we conducted an ablation study by increasing the parameter count of MIP emulator for the 1D Heat example in section 4.2 by order of magnitude and retraining with PRDP. As expected, the larger model achieved higher accuracy, with convergence rate of approximately n -1/5 However, the impact on IC savings was minimal. Interestingly, the largest model exhibited slightly higher savings, likely due to triggering one fewer refinement level during training. Nonetheless, PRDP consistently yields savings of around 80% terms of solver iterations across the varying network sizes.
G.4 PRDP BASED ON A VALIDATION METRIC ON TRAINING LOSS
We motivated Progressively Refined Differentiable Physics (PRDP) by showing how different physics refinement levels K correspond distinct plateaus performance metrics other than train- ing loss. This phenomenon is evident in figure 3 (a) and figure (a). The PRDP control algorithm leverages this behavior by economically increasing the refinement levels, as visibly pronounced in figure 6 (a) and (b).
As introduced in section 3, the PRDP control algorithm relies on performance metric that triggers refinement when it plateaus. In the main text, we used validation error for this purpose, as it exhibited clear plateaus for K < Kmaxs and since hyperparameter tuning is typically done using a validation set. In this section, ablate this choice by instead using the training loss as the guiding metric for PRDP. The results, shown figure 26, reveal that training loss can also serve as valid performance metric for PRDP.
Interestingly, when using training loss, PRDP adopts a slower refinement strategy since training loss often continues to decrease after the validation loss has plateaued. While this approach yiclds slightly higher PR savings, the overall IC savings and final accuracy remain comparable.
Despite these findings, recommend caution when applying PRDP with training loss the guid- ing metric. For more complex problems, such as those with multi-modalee or spurious minima, the training loss may be less reliable than validation loss. Additionally, shown in figure 26, PRDP guided by validation loss results in smoother convergence for both training and validation losses. This smoother progression is likely consequence of the uniform refinement schedule, mak- ing validation loss the preferable choice for broader applications.
38

---

Published as a conference paper at ICLR 2025
12
avg.)
10-1
2-step MMSE
101
Physics refinement (lnner iterations)
10
10-2
10"
10
10
10
10
50
100
50
100
50
100
Epoch
Eooch
Epoch
107
(minibatch avg.
10
Physics refinement (lhner Iterations)
2-step niASE
10
10
10-2
10a
10-3
10-1
10-1
10-1
50
100
100
100
Epoch
Eaoch
Epoch
(a) Neural emulator learning for 2D heat equation: PRDP based on validation error (top) vs. training loss (bottom)
[minibatch avg.)
Physi refinement Ninner iterations)
10
2-step NMSC
10
10
10-1
10-1
10-1
10
50
100
50
100
50
100
Epoch
Eooch
Epoch
10-2
Loss (minibatch avg.1
Physics refinerent (In-er (terations)
101
2-step nMSC
10a
10-4
10-1
10
50
100
100
Epoch
Eooch
Epoch
(b) Neural emulator learning 3D heat equation: PRDP based on validation error (top) vs. training loss (bottom)
39

---

Published as a conference paper at ICLR 2025
(minibatch avg.)
2-step NMSE x-velocity
10
Physics refinement (Inner itarations)
109
100
10-1
10-2
50
100
50
100
50
100
Epoch
Eooch
Epoch
10
(minibatch avg.)
2-step MMSE x-velocity
Physics refinement (inner perations)
100
10>
30
1
10-1
10:2
50
100
50
100
50
100
Epoch
Eaoch
Epech
(c) Neural emulator learning for Burgers equation: PRDP based on validation error (top) vs. training loss (bottom)
(minibatch avg.)
2-step nMSE x-valozity
:0>
15
10-1
Physics refinement (lnner Iterations)
10
10-1
10-2
10-2
50
100
50
100
50
100
Epoch
Eooch
Epech
12.5
(minibatch aug.)
10-1
2-step IMSE x-velocity
10"
Physics refinement ((nner iterations)
10.0
7.5
10-1
10:2
2.5
10-4
50
100
50
100
50
100
Epoch
Eooch
Epoch
(d) Neural emulator learning for Navier Stokes equation: PRDP based on validation emror (top) vs. training loss (bottom)
Figure 26: PRDP based on training loss also successfully trains the network but leads to a slower refinement strategy since the training loss oftentimes continues to reduce even if the validation metric plateaued.
40


Reviews:
[{"review_id": "Hq0WmBKZSv", "reviewer": "Reviewer_Skvr", "summary": "This article presents a method to reduce the computational costs of end-to-end training with a linear PDE solver in the pipeline. In particular, the scope of this article is about linear solves that are big enough to require an iterative solver. It posits that the level of accuracy, that is needed of the forward model evaluation and its gradient, increases during training. The PRDP borrows inspiration from bi-level optimization schemes: the method starts with an inaccurate linear solver (where the number of iterations is stopped too early) and progressively increases the accuracy (or number of iterations in the solver) as the training starts to plateau. The authors provide an algorithm to schedule the number of iteration in function of the validation loss. The authors show computational savings from two major mechanisms: the progressive refinements (the early gradient updates in the training cost less iterations of the solver) and incomplete convergence (where the iterative solver reaches the desired accuracy without needed to train until the number of iterations is sufficient). The computational savings are up to 86% (case of the 2D heat equation if I understand correctly 72% (IC) + 14% (PR) = 86%, but the text only reports 81%), for more complicated examples, the savings go down to 59%.", "questions": "Please add a computational experiment that scales where an iterative solver is necessary (e.g. 3D heat equation).\nPlease consider harder baselines for the linear emulator learning, such as Helmholtz equation.\nPlease add a discussion about potential mechanisms specific to nonlinear neural emulator learning that would explain the reduced computational savings.\n\nMinor:\nIf needed, please correct the reported saving of the 2D heat equation to 72%+14%=86% as well as the maximum saving in the conclusion.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "This article is clearly written and introduces a promising solution to computational costs of training campaigns that include a linear PDE solver, end to end.\nThe main result, that progressively increasing the accuracy of the solver in end-to-end training converges to similar performance as a fully converged solver while substantially reducing computing costs, is original.\nThe benefit from incomplete convergence is very interesting. Although unexpected, the computational experiments show that they originate most of the computational savings.\nIt is interesting to see that the computational saving are larger for the 2D heat equations than for the 1D heat equation.\nIt is an interesting finding that unrolled differentiation which inefficiently accurately differentiate a sequence of iterative approximation performs similarly to the smarter implicit differentiation.", "weaknesses": "The paper suffer weak baselines. The main use case of the method is for iterative solvers of large linear models, however, the authors use 1D and 2D examples which would likely be solved efficiently by a direct solver. The authors should at least provide one 3D example with the simple heat equation. The heat equation itself is also a weak baseline, more complex linear models such as Helmholtz equation could be considered.\n\nIt is worrying that the benefit of the method diminish as the problem get harder as in the nonlinear neural emulator learning. Some discussion is needed about the potential harmful mechanisms that limited the computational savings in that case.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "KdIZWHyvvY", "reviewer": "Reviewer_5MWA", "summary": "This paper introduces a framework to progressively refine the resolution of physics solvers during neural network training. The authors demonstrate that training a neural network with a physics solver scheduled to increase in iterations $K$ over training can significantly reduce computational costs, especially by using fewer iterations in the early training phases (*progressive refinement* savings). Additionally, they observe that a neural network can be effectively trained even when the physics solver has not fully converged, eliminating the need for an extensive number of solver steps to achieve high accuracy (*incomplete convergence* savings). To automate this process and determine optimal parameters for these refinements, the authors propose an algorithm that monitors validation set metrics, incrementally increasing solver refinement when performance plateaus. The framework is validated across four use cases: a linear inverse solver, linear neural emulator learning, nonlinear neural emulator learning, and a neural-hybrid emulator.", "questions": "* How sensitive are the parameters of PRDP? Did you experiment with many different parameters for each problem before achieving the results, or was it relatively straightforward?\n* Do you expect the method to achieve similar savings in training time for domains with complex boundary conditions and irregular meshing?\n* Do you think your method could work where the physics solver contains incomplete physics (e.g., the parameters are not perfectly calibrated, and some terms of the equations could be missing)?\n* How would the method be affected by noise in the observations?\n* How well do you think this could help reduce the training time of neural GCMs [1]?\n\n[1] Kochkov et al. Neural general circulation models for weather and climate. Nature, 2024.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* The paper is well-written, with clear explanations of the intuitions and motivations behind the method.\n* The algorithm is straightforward and effectively delivers the intended results, as seen in the reduction of validation loss with the progressive refinement of the physics solver, particularly notable in Figure 4 for the 2D Heat and 2D Navier-Stokes cases.\n* The savings in training time and computational resources are substantial.\n* The appendix is thorough and well-organized, with especially valuable details on iterative linear solvers and detailed derivations for each problem.", "weaknesses": "* Overall, the technical contribution of the paper is somewhat limited, with the main novelty being the proposed algorithm for iterative refinements.\n* All physics solvers employed rely on iterative linear solvers. It would have been interesting to see if the method also applies with other physics solvers.\n* With the exception of the final example (the neural-hybrid approach), the other examples appear to be simplified or illustrative cases without clear, concrete applications.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "wt3hpen9nF", "reviewer": "Reviewer_1fLE", "summary": "This paper proposes to use progressively refined differentiable physics, termed as PRDP, to increase the training efficiency while not harnessing the accuracy. The key finding lies in the fact that the full accuracy of the neural network is achievable through insufficiently converged solvers. Several experiments are conducted to validate the effectiveness of PRDP in reducing training time.", "questions": "The experiments report the improved efficiency by adopting progressive refinement and incomplete convergence. Does these strategies influence the accuracy?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The topic this paper wants to tackle seems interesting. It seems intuitive that, considering the noiseness of neural network training and approximative nature of deep models, the physics solver does not need to fully converge for the network to achieve maximum possible accuracy. This paper proposes to use an adaptive strategy to progressively refine the physics solver and thus improve the training efficiency. Several experiments are conducted to verify the efficacy of the proposed method.", "weaknesses": "- This paper should provide more background information about *differentiable physics* to make readers better understand the core contribution of the proposed method. I am not an expert of this field, and I find this paper a little bit hard to follow, and also unaware of the broader context this paper lies in.\n- The experiment settings in this paper are not clearly presented. Considering that this is paper submitted to ICLR, I want to know what is the role of the neural networks in each experiment.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ebmMWv4Nqz", "reviewer": "Reviewer_h4zx", "summary": "PRDP proposes an algorithm to reduce computational cost of training neural networks when a nested optimization problem is required during training. The authors first exhibit source of savings (IC and PR), before introducing their method. Specifically, this method consists in considering partially solving the inner problem in for the firsts steps of training and then refining this inner solution to more precise one as the training goes through. Moreover, the method considers not fully solving linear systems, since it may not greatly improve the performances. Finally, they evaluate their method on several PDE problems and instances.", "questions": "-\tIs this method only applicable in the context of Physical systems? It seems to me that this method could be more general and this being used in a broader range of applications, as soon as an iterative process incurs in the forward pass?  \n-\tCould the author also provide a comparison of the differences of performances w/ and w/o PRDP? (see for example Fig 1 where it looks like there is a very little difference in performances, thus making me wondering how much is this loss)\n-\tWhy is the algorithm based on validation losses? In what does these losses consist in? At inference, these validation values are not available? \n-\tWhat are the applications at inference? Once trained, what would be some application of NNs? Could this method be applied on new physical systems/PDEs/boundary or initial conditions/discretization? \n-\tIn section 2, what does the subscript h stands for? Are the parameters $\\theta$ the neural network parameters to be optimized through training? \n-\tOn page 2, last paragraph it is stated that experiments are conducted on real world application, is the Navier-Stokes this example? These are synthetic data; in which sense do you consider this example as “real-world application”? \n-\tIn example 2.2, are the $\\theta$ parameters optimized with the outer step? This means that one wants to optimize the forcing terms? The application would be to find the forcing term associated to a recorded and given trajectory? \n-\tFor the IC savings, I was wondering if the authors have tried to experiment if without NN, the performances would be better? My guess is that the introduction of a NN in the framework prevent the performances from being optimal, thus allowing for IC savings. What if the neural network size increases/its expressiveness improves? Are the performances better? \n-\tThe main claim of the paper is computation savings. Could the author provide training times for their experiments? And a comparison of this solving time with standard methods? \n-\tWhat happen on cases where the losses plateau then decreases again? This situation could arise in some training of neural networks especially, when using a learning rate scheduler. How would behave the method in this context? \n-\tIn section 4.4 why don’t you compare your results (performances, training times) with the method from Um et al.? Since the setting is the same and your method is supposed to improve training times, it would be interesting to evaluate the benefit with other methods existing. Moreover, in the related work several methods are cited, and could be used as a comparison.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is very complete. \n- Several experiments are proposed and isolate and illustrate well the 2 kinds of cost reduction proposed. \n- The proposed algorithm is easy to understand and well described. \n- A lot of details about training and the physical systems studied are proposed in the appendices. \n- The authors discuss impact and limitations of their work", "weaknesses": "- The paper is very technical and I think pseudo code or a scheme would help in the comprehension of the method. For example, it is unclear to me where is the neural network used in the global framework/experiments, what are inner/outer optimization problems in the examples... Despite understanding the overall idea and performances of the proposed algorithm, I think more explicit and easy-to-understand notations would help the comprehension. A more detailed example would help for a precise comprehension of the framework proposed (see questions section).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)