Solving Inverse Physics Problems with Score Matching
Authors: Benjamin Holzschuh
Simona Vegetti
Nils Thuerey
Venue: NeurIPS 2023 poster (2023)
Keywords: inverse problems
diffusion models
learned corrections
score matching
TL;DR: We propose a novel methodology for solving inverse problems that involve the temporal evolution of physical systems.

Abstract: We propose to solve inverse problems involving the temporal evolution of physics systems by leveraging recent advances from diffusion models. 
Our method moves the system's current state backward in time step by step by combining an approximate inverse physics simulator and a learned correction function. 
A central insight of our work is that training the learned correction with a single-step loss is equivalent to a score matching objective, while recursively predicting longer parts of the trajectory during training relates to maximum likelihood training of a corresponding probability flow.
We highlight the advantages of our algorithm compared to standard denoising score matching and implicit score matching, as well as fully learned baselines for a wide range of inverse physics problems. The resulting inverse solver has excellent accuracy and temporal stability and, in contrast to other learned inverse solvers, allows for sampling the posterior of the solutions. Code and experiments are available at https://github.com/tum-pbs/SMDP.

Solving Inverse Physics Problems with Score Matching
Benjamin Holzschuh 1.2+
Simona Vegetti 2
Nils Thuerey'
Technical University of Munich, 85748 Garching, Germany Max Planck Institute for Astrophysics, 85748 Garching, Germany
Abstract
We propose to solve inverse problems involving the temporal evolution of physics systems by leveraging recent advances from diffusion models. Our method moves the system's 's current state backward in time step by step by combining proximate inverse physics simulator and a learned correction function. A central insight of our work is that training the learned correction with a single-step loss is equivalent to a score matching objective, while recursively predicting longer parts of the trajectory during training relates to maximum likelihood training of corresponding probability flow. We highlight the advantages of our algorithm com- pared to standard denoising score matching and implicit score matching, well fully learned baselines for a wide range of inverse physics problems. The resulting inverse solver has excellent accuracy and temporal stability and, in contrast other learned inverse solvers, allows for sampling the posterior of the solutions. Code and experiments are available at withttts:// //tithut.bt..... cww/ coom/tum-ppe Sum-pbs/SKDDP
ap-
Introduction
Many physical systems are time-reversible on a microscopic scale, For example, a continuous material can be represented by collection of interacting particles [Gur82 BLL02] based on which we can predict future states. We can also compute earlier states, meaning we can evolve the simulation backward in time [Mar+96]. When taking macroscopic perspective, we only know average quantities within specific regions [Far93], which constitutes a loss of information, and as consequence, time no longer reversible. In the following, we target inverse problems to reconstruct the distribution of initial macroscopic states for given end state. This problem is genuinely tough ZDDG6, Góm++8) Del+ 18 LP22). and existing methods lack tractable approaches to represent and sample the distribution of states.
Our method builds on recent advances from the field of diffusion- based approaches [Soh+15) HJA20, Son+21b]: Data samples x E R are gradually corrupted into Gaussian white noise via stochastic differential equation (SDE) dx = f(x, t) + (t))WW where the deterministic component of the SDE f IR' R>0 + R is called drift and the coefficient of the D D-dimensional Brownian motion W denoted by R22 + R20 called diffusiin.. If the score Vx log Pe (x) of the data distribution px(x) of corrupted samples at time I. is known, then the dynamics of the SDE can be reversed in time, allowing for the sampling of data from noise. Diffusion models are trained to approximate the score with neural network s9. which can then be used as plug-in in estimate for the reverse-time SDE.
However, in our physics-based approach, we consider an SDE that describes the physics system as dx = P(x) dt + gW, where P. R,D - H is a physics simulator that replaces the drift term of diffusion models. Instead of transforming the data distribution noise, we transform a simulation state at t = 0 to a simulation state at = T with Gaussian noise perturbation. Based on a given
"Correspondence to: benjamin, holzschuh@tum.de de
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

end state of the system at t= T, we predict a previous state by taking small time step backward in time and repeating this multiple times. Similar the reverse-time SDE of diffusion models, the prediction of the previous state depends on an approximate inverse of the physics simulator, learned update 80, and small Gaussian perturbation.
The training of is similar to learned correction approaches for numerical simulations (Um+20) Koc+2I LCT22]: : The network sy learns corrections to simulation states that evolve over time according to physics simulator so that the corrected trajectory matches a target trajectory. In our method, we target learning corrections for the "reverse" simulation. Training can either be based on single simulation steps, which only predict single previous state, or be extended to rollouts for multiple steps. The latter requires the differentiability of the inverse physics step (Thu+2T).
Importantly, we show that under mild conditions, learning 8e is equivalent to matching the score vx log px()) of the training data set distribution time f. Therefore, sampling from the reverse-time SDE of our physical system SDE constitutes a theoretically justified method to sample from the correct posterior.
While the training with single steps directly minimizes a score matching objective, we show that the extension to multiple steps corresponds to maximum likelihood training of a related neural ordinary differential equation (ODE). Considering multiple steps is important for the stability of the produced trajectories. Feedback from physics and neural network interactions at training time leads to more robust results.
In contrast to previous diffusion methods, we include domain knowledge about the physical process in the form of an approximate inverse simulator that replaces the drift term of diffusion models [Son+2Ib] ZC21]. In practice, the learned component corrects any errors that occur due to numerical issues, e.g., the time discretization, and breaks ambiguities due to a loss of information in the simulation over time.
Figure I gives an overview of our method. Our central aim is to show that the combination of diffusion- based techniques and differentiable simulations has merit for solving inverse problems and to provide theoretical foundation for combining PDEs and diffusion modeling. In the following. we refer to methods using this combination as score matching via differentiable physics (SMDP). The main contributions of our work are: (1) We introduce reverse physics simulation step into diffusion models to develop probabilistic framework for solving inverse problems. (2) We provide the theoretical foundation that this combination yields learned corrections representing the score of the underlying data distribution. (3) We highlight the effectiveness of SMDP with set of challenging inverse problems and show the superior performance of SMDP compared to a range of stochastic and deterministic baselines.
2 Related Work
Diffusion models and generative modeling with SDEs Diffusion models (Soh+15) HJA20
have been considered for a wide range of generative applications, most notably for image DN21
video [Ho+22; Höp+22 YSM22), audio synthesis [Che+2l). uncertainty quantification (CSY22
Chu+22 22n+22 KVE21 Ram+20], and as autoregressive PDE-solvers (KCT23). However, most approaches either focus on the denoising objective common for tasks involving natural images the
KVE21
synthesis process of solutions does not directly consider the underlying physics. Models based on Langevin dynamics [Vinll] SE19] discrete Markov chains (Soh+15) HJA20) can be unified in a contennon framework using SDEs [Son+21b). Synthesizing data by sampling from neural SDEs has been considered by, e.g.. [Kid+21| [Son+21b). Contrary to existing approaches, the drift in our method is an actual physics step, and the underlying SDE does not transform a data distribution to noise but models the temporal evolution of physics system with stochastic perturbations.
Methods for solving inverse problems for (stochastic) PDEs Differentiable solvers for physics dynamies can be used to optimize solutions of inverse problems with gradient- based methods by backpropagating gradients through the solver steps [hu++1l] Learning- based approaches directly learn solution operators for PDEs and stochastic PDEs, i.e., mappings between spaces of functions, such as Fourier neural operators [Li+1]] DeepONets 21+21] or generalizations thereof that include stochastic forcing for stochastic PDEs, e.g., neural stochastic PDEs (SIGG22]. Recently, there have been several approaches that leverage the learned scores from diffusion models data-driven

---

marginal likelibood Py
Purtuubed fesining trajectories
Time
Xo
muve window
Optimization
Reference
Prediction
Training
Ss.
X2
X2
Fit probability flow to trajectories dx
Apf
Sliding window sizre S
reverse physics step
function
X3
No 34. E4
X3
Inference
Campeian
portqubatiee
Po
(xa)
Kelx
gradients gradient flow
pr
approximates
V. log
X4
data flow
(a) Training and inference overview.
(h) Sliding window and optimizing sg-
Figure I: Overview of our method. For training, we fit neural ODE, the probability flow, to the of perturbed training trajectories (a, top). The probability flow is comprised of reverse physics simulator P that an approximate inverse of the forward solver P well as a correction function 8g. In many cases, we can obtain P-1 from by using negative step size At by learning surrogate model from data. For inference, we simulate the system backward in time from Xy to Xy by combining P the trained and Gaussian noise in each step (a, bottom). For optimizing our approach moves sliding window of size S along training trajectories and reconstructs the current window (b). Gradients for / are accumulated and backpropagated through all prediction steps.
regularizations for linear inverse problems [Ram+20] Son+22 KVE21 CSY22 Chu+22 and general
noisy inverse problems [Chu+23]. Our method can applied general non-linear inverse physics problems with temporal evolution, and we do not require to backpropagate gradients through all solver steps during inference. This makes inference significantly faster and more stable.
Learned corrections for numerical errors Numerical simulations benefit greatly from machine
learning models [Tom+17] Mor+ 20 Li+21]. By integrating neural network into differential
equation solvers, it is possible to learn to reduce numerical errors [Um+200] Koc+21]
B BWW22 or
guide the simulation towards a desired target state (HTK20) Li+22]. The optimization
8o with
the 1-step and multi-step loss propose in section 3.1 is conceptually similar to learned correction approaches. However, this method has, to our knowledge, not been applied to correcting the "reverse" simulation and solving inverse problems.
Maximum likelihood training and continuous normalizing flows Continuous normalizing flows (CNFs) are invertible generative models based neural ODEs [Che+ KPB20 Pap+21], which are similar to our proposed physics-based neural ODE. The evolution of the marginal probability density of the SDE underlying the physies system described by Kolmogorov's forward equation [Øks03), and there is corresponding probability flow ODE [MRO20, Son+21b), When the score represented by sg. this constitutes CNF and can typically be trained with standard methods [Che+ 18b] and maximum likelihood training [Son+21a]. Huang et al. (HLC21) show that minimizing score-matching loss is equivalent maximizing a lower bound of the likelihood obtained by sampling from the reverse-time SDE. A recent variant combines score matching with CNFs [2C2 and employs joint learning of drift and corresponding score for generative modeling. To the best of our knowledge, training with rollouts of multiple steps and its relation to maximum likelihood training have not been considered so far.
Method Overview
Problem formulation Let (Ω.F, P) be probability space and ((C) = (W), (D) be D-dimensional Brownian motion. Moreover, let Xq be Fe-mauurable TR random variable that is distributed as po and represents the initial simulation state. We consider the time evolution of

---

the physical system for 0 S T modeled by the stochastic differential equation (SDE)
dx = P(x)dl + 9()dd
with initial value Xy and Borel measurable drift P R2D - IR and diffusion g [0, T + R>0-
This SDE transforms the marginal distribution Po of initial states time 0 to the marginal distribution Pr of end states at time T. We include additional assumptions appendix
Moreover, we assume that we have sampled N trajectories of length M from the above SDE with fixed time discretization 0 S to < ti < ... < < S for the interval [0, 2] and collected them in training data set {xx;
1.N (n) ) ====n For simplicity, we assume that all time steps are equally spaced, i.e. fm+1 fm = At. Moreover, in the following we use the notation Xem for 0 S l m M to
refer to the trajectory (Xe,, Xter11 Xt..). We include additional assumptions in appendix A
Our goal is infer initial state Xo given simulation end state XM,, i.e., we want sample from the distribution Po( XM), or obtain a maximum likelihood solution.
3.1 Learned Corrections for Reverse Simulation
In the following, we furthermore assume that have access to reverse physics simulator p-1 RDD + RD, which moves the simulation state backward in time and is an approximate inverse the forward simulator P [HK In our experiments, from the forward simulator by using a negative step size At or by learning a surrogate model from the training data. We train neural network sx(x, parameterized by 0 such that
either obtain the reverse physics simulator
Xm Xm++1 At p (xm+1) + 88 (Xm+1 (m+1)
In this equation, the term $8 (xm+1, +m++1) corrects approximation errors and resolves uncertainties from the Gaussian perturbation g())dW. Below, we explain our proposed 1-step training loss and its multi-step extension before connecting this formulation to diffusion models in the next section.
1-step loss optimizing 89 is the 1.2 distance between Xm and the prediction via (2). For the entire training data set, the loss becomes
For pair of adjacent samples (Xm: Xm++1) ) on data trajectory, the 1-step loss for
M Σ
Lsingle (0)
E
Xm xm+11 Al [ - (xm=1) so (Xm+1, (mmt1)
M
Computing the expectation can be thought of as moving window of size from the beginning of each trajectory until the end and averaging the losses for individual pairs of adjacent points.
Multi-step loss As each simulation state depends only on its previous state, the 1-step loss should be sufficient for training sa. However, practice, approaches that consider loss based on predicting more successful for training learned corrections [Bar+19] Um+20
longer parts of the trajectories Koc+21], For that purpose, we define a hyperparameter S, called sliding window size, and write X S to denote the trajectory starting at x; that is comprised of x; and the following S states. Then, we define the multi-step loss as
M- Σ
1
Ematti (0)
Ex
ExO.M
[|xmmm551 8mm +
(4)
M
M==
where 82148-1 the predicted trajectory that is defined recursively by
8445 = X
and
8145-1-3 = 8its-j + At (appp-(xxe (p-1(x(4s_j)) + so ((445-))
(5)
3.2 Learning the Score
Denoising score matching Given distribution of states for 0 < << T, we follow Son+2 and consider the score matching objective
1
(0) = 2 Ex-m [1s)((x))) Vx [xlogp((x)1 dt,
(6)

---

ie.. the network s9 is trained to approximate the score Vx log (x). In denoising score matching [VinIT] Soh+15 HIA2 the distributions pt are implicitly defined by a noising process that given by the forward SDE dx = [(x t)dt + g())WW where W is the standard Brownian motion. The DR Rz0 + RDD is called drift, and g R20 -> R2o is called diffusion.. The process transforms the training data distribution Po to a noise distribution that is approximately Gaussian pr- For affine functions/ f and g. the transition probabilities are available analytically, which allows for efficient training of sg. It can be shown that under mild conditions, for the forward SDE, there is corresponding everse-time SDE dx = [[(x, 1) v(1)²Vx log (x))dl + g())dW [And82]. In particular, this means that given a marginal distribution of states pr. which is approximately Gaussian, can sample from pr and simulate paths of the reverse-time SDE to obtain samples from the data distribution po-
functioo f:
Score matching, probability flow ODE and 1-step training There is a deterministic ODE [Son+21b), called probability flow ODE, which yields the same transformation of marginal probabili- ties from pr to po as the reverse-time SDE. For the physics-based SDE (T. it is given by
1
dx = P(x)
2 g log Px
df.
For At -> 0, we can rewrite the update rule
(2) of the training as
dx [- (xx- so (x, t)] dt.
(8)
Therefore, we can identify (x) with -P(x) and s0( 1) with (1) (x.). We show that for the 1-step training and sufficiently small At, we minimize the score matching objective 6
Theorem 3.1. Consider data set with trajectories sampled from SDE ① and let P- (x) -x().
Then the 1-step loss 3 is equivalent to minimizing the score matching objective (6) as At - 0.
Proof. See appendix A.1
Maximum likelihood and multi-step training Extending the single step training to multiple steps does not directly minimize the score matching objective, but we can still interpret the learned correction in probabilistic sense. For denoising score matching, it is possible train so via maximum likelihood training [Son+2la), which minimizes the KL-divergence between po and the distribution obtained by sampling Xr from pr and simulating the probability flow ODE (8) from t= T to / = O. We derive similar result for the multi-step loss.
Theorem 3.2. Consider data set with trajectories sampled from SDE ① and let (x) -x().. Then the multi-step loss (4) maximizes a variational lower bound for maximum likelihood training of the probubility flow ODE as Al + O.
Proof. See appendix A.2
To conclude, we have formulated a probabilistic multi-step training to solve inverse physics problems and provided theoretical basis to solve these problems with score matching. Next, outline additional details for implementing SMDP.
3.3 Training and Inference
We start training so with the multi-step loss and window size S = 2, which is equivalent to the 1 I-step loss. Then, we gradually increase the window size S until maximum Smax- For > 2, the unrolling of the predicted trajectory includes interactions between sy and the reverse physics simulator P-1, For inference, we consider the neural SDE
dx = - (x) C s t) dt + gt)dd,
(9)
which we solve via the Muler-Maruyama method. For C 2, we obtain the system's reverse-time
SDE and sampling from this SDE yields the posterior distribution. Setting = 1 and excluding the noise gives the probability flow ODE (8). We denote the ODE variant by SMDP ODE and the SDE variant by SMDP SDE. While the SDE can used to obtain many samples and to explore the posterior distribution, the ODE variant constitutes unique and deterministic solution based on maximum likelihood training.
5

---

GRID, 1-step
Training trajectories
Reverse-time SDE
dx =(x)dtt + g(t)aw
dx - ((x) (es(x, s1dtt w(o)
GRID multi-step
-1
Score Vx (xg
Probability Flow
Exploding
trajectories
MLP, 1-step
appreximattdd by so(x,
dx p- (xx(x
so(x, d)
T Q
75 T
Time
Time t
T
Time
(a) Training
(b) Inference
(c) 1-step vs. multi-step
Figure Overview of our ID toy SDE. (a) Training with a data set of trajectories and known temporal dynamics given by P(x) == =xign(x)x² and 0.1. We estimate the score Vx log (x) with our proposed method using an MLP network for s t). Negative values (blue) push down the trajectories, and positive ones (red) push them up. Together with the dynamics, this can be used to reverse the system as shown in (b) either with the reverse-timg SDE or the probability flow ODE. A successful inversion of the dynamics requires the network to be robust and extrapolate well (c). Inference using GRID trained with the 1-step loss causes trajectories to explode, as the network does extrapolate well. Training GRID with the multi-step loss solves this issue.
4 Experiments
We show the capabilities of the proposed algorithm with a range of experiments. The first experiment section 4.1 (uses simple ID process to compare our method to existing score matching baselines. The underlying model has known posterior distribution which allows for an accurate evaluation of the performance, and we use it to analyze the role of the multi-step loss formulation. Secondly, in section 4.2] we experiment with the stochastic heat equation. This is particularly interesting test case as the diffusive nature of the equation effectively destroys information over time. In section 4.3 we apply our method to scenario without stochastic perturbations in the form of gnoyancy-drived Navier-Stoke Stokes flow with obstacles. This case highlights the usefulness of the ODE variant. Finally, in section 4.4 we consider the situation where the reverse physics simulator P-1 is not known. Here, we train a surrogate model P for isotropic turbulence flows and evaluate how well SMDP works with learned reverse physics simulator.
4.1 1D Toy SDE
As first experiment with known posterior distribution we consider a simple quadratic SDE of the form: de [A sign(z)>>) dl. + A2dW, with A 7 and d2 = 0.03. Throughout this experiment, Po a categorical distribution, where we draw either or with the same probability. The reverse-time SDE that transforms the distribution pr of values at T = 10 to Po is given by
dr [A sign(r)>² Xラ V, logPe(=)) dt A2dw.
(10)
In figure 2a
we show paths from this SDE simulated with the Euler-Maruyama method. The trajectories approach 0 increases. Given the trajectory value at = 10, it is no longer possible to infer the origin of the trajectory at (
This experiment allows us to use an analytic reverse simulator: P (r)(x) = A sign(x)x². This is challenging problem because the reverse physics step increases quadratically with æ, and sæ has to control the reverse process accurately to stay within the training domain, or paths will explode to infinity. We evaluate each model based on how well the predicted trajectories forr match the posterior distribution. When drawing पT randomly from I-0.1, 0.1]. we should obtain trajectories with to being either -1 or with the same likelihood. We assign the label or if the relative distance of an endpoint is < 10% and denote the percentage in each class by p-1 1 and p1. As some trajectories miss the target, typically p-1 + + #1 < Hence, we define the posterior metric Q as twice the minimum of p_1 and pl, i.e., Q = 2 2-n--miinnpp.. min(p-1, P1) so that values closer to one indicate better match with the correct posterior distribution.

---

Training The training data set consists of 2500 simulated trajectories from 0 to T and At = 0.02. Therefore each training trajectory has a length of M 500. For the network sp(a, t). we consider multilayer perceptron (MIP) and, as a special case, grid-based discretization (GRID). The latter not feasible for realistic use cases and high-dimensional data but provides means for an in-depth analysis of different training variants. For GRID, we discretize the domain [0, 7) x [-1 1.255, 1.25 to obtain a rectangular grid with 500 x 250 cells and linearly interpolate the solution. The cell centers are initialized with O. We evaluate sœ trained via the 1-step and multi-step losses with Smax 10. Details of hyperparameters and model architectures are given appendix
Better extrapolation and robustness from multi-step loss See figure 2c for overview of the differences between the learned score from MLP and GRID and the effects of the multi-step loss. For the 1- 1-step training with MLP, observe clear and smooth score field with two tubes that merge to one at æ = 0 as increases. As a result, the trajectories of the probability flow ODE and reverse-ting SDE converge to the correct value. Training via GRID shows that most cells do not get any gradient updates and remain 0. This is caused by a need for more training data in these regions. In addition, the boundary of the trained region is jagged and diffuse. Trajectories traversing these regions quickly explode. In contrast, the multi-step loss leads to a consistent signal around the center line at 0, effectively preventing exploding trajectories.
Evaluation and comparison with baselines As a baseline for learning the scores, we consider implicit score matching (Hyv05, ISMJ. Additionally, variance reduction [Son+19 19 SSM-VRJ as a variant of ISM. We train all methods with the same network architecture using three different data set sizes. As can be seen table the 1-step loss, which is conceptually similar to denois-
consider sliced score matching with
ing score matching. compares favorably against ISM and SSM-VR. All methods perform well for the reverse-time SDE, even for very little training data. Using the multi-step loss consistently gives signifi- cant improvements at the cost of slightly increased training time. Our proposed multi-step training performs best or is on par with the baselines for data set sizes and inference types. Because the poste- rior metric Q is very sensitive to the score where the paths from both starting points intersect, evaluations are slightly noisy.
Method
Probability flow ODE
Reverse- time SDE
Data set size
Data set size
100%
10%
1%
100%
10%
1%
multi-step
0.97
0.91
0.81
0.99
0.94
0.85
1-step ISM SSM-VR
0.78
0.44
0.41
0.93
0.71
0.75
0.19
0.15
0.01
0.92
0.94
0.52
0.17
0.49
0.27
0.88
0.94
0.67
Table 1: Posterior metric Q for 1000 predicted trajecto- ries averaged over three runs. For standard deviations, see table the appendix.
Comparison with analytic scores We perform further experiments to empirically verify Theorem 3.1b by comparing the learned scores of our method with analytic scores in appendix C
10
Head diffusson
SMDP ODE
0.2
ResNet
FNO
BNN
10
10a
Sy only ODE
only SDE
SMDP ODE
SMDP SDE
Reconstruction MSE
Spectral error
(a) Inference
(b) Evaluation
Figure 3: Stochastic heat equation overview. While the ODE trajectories provide smooth solutions with the lowest reconstruction MSE, the SDE solutions synthesize high-frequency content, signifi- cantly improving spectral error. The "sø only" version without the reverse physics step exhibits significantly larger spectral error. Metrics in (b) are averaged over three runs.
7

---

4.2 Stochastic Heat Equation
The heat equation Ou = a&u plays fundamental role in many physical systems. For this experiment, we consider the stochastic heat equation, which slightly perturbs the heat diffusion process and includes an additional term g(t) Ç, where is space-time white noise, see Pardoux [Par21 Chapter 3.2]. For our experiments, we fix the diffusivity constant a and sample initial conditions at t t= O from Gaussian random fields with
at
4 at resolution 32 x 32. We simulate the heat diffusion with noise from 1 = until = 0.2 using the Euler-Maruyama method and a spectral solver Ph with fixed step size At 6.25 x 10 and 0.1. Given simulation end state XT. we want to recover a possible initial state Xp. In this experiment, the forward solver cannot be used to infer Xy directly in single step or without corrections since high frequencies due to noise are amplified, leading to physically implausible solutions. We implement the reverse physics simulator p-1 by using the forward step of the solver Pa i.e. A...p-2 P=1 R (x).
Training and Baselines Our training data set consists 2500 initial conditions with their corre- sponding trajectories and end states at / 0.2. We consider small ResNe/-like architecture based on an encoder and decoder part as representation for the score function sy (x. t). The spectral solver implemented via differentiable programming in JAX (SC20), appendix D As bascline methods, we consider a supervised training of the same ResNer- architecture as sp x, t). a Bayesian neural network (BNN) as well as Fourier neural operator (FNO) network [Li+21]. We adopt L2 loss for all these methods, i.e., the training data consists of pairs of initial state X and end state Xy-
Additionally, we consider variant of our proposed method for which we remove the reverse physics step p-1 such that the inversion of the dynamies has to be learned entirely by s9, denoted by "sø only". We do not compare to ISM and SSM-VR in the following as the data dimensions are too high for both methods train properly, and we did not obtain stable trajectories during inference.
Reconstruction accuracy fitting the data manifold We evaluate our method and the baselines by considering the reconstruction MSE on test set 500 initial conditions and end states. For the reconstruction MSE, sim- ulate the prediction of the network forward in time with the solver Pa, to obtain correspond- ing end state, which we compare to the ground truth via the L2 distance. This metric has the disadvantage that it does not measure how well the prediction matches the training data mani- fold. I.c., for this case, whether the prediction resembles the properties of the Gaussian random field. For that reason, we additionally compare
Reconstruction MSE
-x10
10
10-5
10-3
6x10
4x
10:1
2
16
32
16
32
(a)
ODE trajectories
(b) SDE trajectories
Figure 4: Multi-step Smax vs. reconstruction MSE averaged over 5 runs.
the power spectral density of the states as the spectral loss. An evaluation and visualization of the reconstructions are given in figure which shows that our ODE inference performs best regarding the reconstruction MSE. However, its solutions are smooth and do not contain the necessary small-scale structures. This is reflected in a high spectral error. The SDE variant, on the other hand, performs very well in terms of spectral error and yields visually convincing solutions with only a slight increase in the reconstruction MSE This highlights the role of noise as a source of entropy in the inference process for SMDP SDE, which is essential for synthesizing small-scale structures. Note that there is natural tradeoff between both metrics, and the ODE and SDE inference perform best for each of the cases while using an identical set of weights.
Multi-step loss is crucial for good performance We performed ablation study the maximum
window size Smax in figure 4 for the reconstruction MSE. For both ODE and SDE inference, increasing Smax yields significant improvements the cost of slightly increased training resources. This also highlights the importance using a multi-step loss instead of the 1-step loss (Smax = 2) for inverse problems with poor conditioning.
We perform further experiments regarding test-time distribution shifts when modifying the noise scale and diffusivity, see appendix D which showcase the robustness of our methods.

---

4.3 Buoyancy-driven Flow with Obstacles
Next, we test our methodology on more challenging problem. For this purpose, we consider deterministic simulations of buoyancy-driven flow within fixed domain Ω C [0. [) x [0. 1] and randomly placed obstacles. Each simulation runs from time 0.0 to = 0.65 with a step size of
At= = 0.01. SMDP is trained with the objective of reconstructing a plausible initial state given end state of the marker density and velocity fields at time = 0.65, as shown in figure 55 and figure (5b We place spheres and boxes with varying sizes different positions within the simulation domain that do not overlap with the marker inflow. For each simulation, we place one to two objects each category.
Score matching for deterministic systems During training, we add Gaussian noise to each simulation state x, with o, = Vtt. In this experiment, no stochastic forcing is used to create the data set, i.e., 9 = 0. By adding noise to the simulation states, the 1-step loss still minimizes score matching objective in this situation, similar to denoising score matching; see appendix [A.3 for a derivation. In the situation without stochastic forcing, during inference, our method effectively alternates between the reverse physics step, small perturbation, and the correction by so(x. xx l), which projects the perturbed simulation state back the distribution pt. We find that for the SDE trajectories, C = 2 slightly overshoots, and C = 1 gives an improved performance. In this setting, the "gg only" version of our method closely resembles a denoiser that learns additional physics dynamics.
Training and comparison Our training data set consists of 250 simulations with corresponding trajectories generated with phiflow [Hol+20]. Our neural network architecture for sg x, t) uses dilated convolutions [Sta+21]], appendix H for details. The reverse physics step P-1 is implemented directly in the solver by using negative step size - At for time integration. For training, we consider the multi-step formulation with Smax 20. We additionally compare with solutions from directly optimizing the initial smoke and velocity states at t = 0.35 using the differentiable forward simulation and limited-memory BFGS [LN89] LBFGS]. Moreover, we compare with solutions obtained from diffusion posterior sampling for general noisy inverse problems [Chu+23] DPSJ with pretrained diffusion model on simulation states at 0.35. For the evaluation, we consider a reconstruction MSE analogous to section 4.2 and the perceptual similarity metric LPIPS. The test set contains five simulations. The SDE version yields good results for this experiment but is most likely constrained in performance by the approximate reverse physics step and large step sizes. However, the ODE version outperforms directly inverting the simulation numerically (P-1 only), and when training without the reverse physics step (sg only), as shown in 5c
4.4 Navier-Stokes with Unknown Dynamics
As a fourth experiment, we aim to learn the time evolution of isotropic, forced turbulence with a similar setup as Li et al. [Li+2T]. The training data set consists of vorticity fields from 1000 simulation trajectories from 0 until T = 10 with At = 1, spatial resolution of 64 x 64 and
e 0.60
Reconstruction
n.66
Reconstruction
LBEGS
DPS
5p only
r-1 only
SMDP ODE
SMDP SDE
Ground bruth
0.6
6 0.55
0.15
Reference
SMDP ODE
Reference
SMDP ODE
10
0.2
10-3
only
p
only
So only
only
LPIPS (r 0.35)
Rec. MSE (t 0.35)
(a) Marker density
(b) Velocity field (x)
(c) Evaluation
Figure 5: Buoyancy flow case. Ground truth shows the marker density and velocity field in the a--direction at different points of the simulation trajectory from the test set (a, b). We show recon- structions given the simulation end state t = 0.65 and provide an evaluation of the reconstructed trajectories based on perceptual similarity (LPIPS) and the reconstruction MSE for three runs (c).

---

0.6
Learned reverse step
Xg
Simultaneous training of jir se
Reconstruction
0.4
16
Pretrained jr1 SMDP ODE
15
0.2
f-1
Xg Reference
Xg only
X SMDP ODE
14
0.1
SMDP SDE
MSE
Spectral error
(a) Reconstructed trajectories at t
(b) Evaluation
Figure 6: Turbulence case. Comparison of reconstructed trajectories (a) and evaluation of MSE and spectral error for different training variants (b). Our proposed ODE and SDE inference outperforms the learned surrogate model p- Metries are averaged over three runs.
viscosity fixed 10-5. As before, objective is to predict trajectory &o:M that reconstructs the true trajectory given an end state XM- In this experiment, we pretrain a surrogate for the reverse physics step p-1 by employing the FNO architecture from [Li+2T] trained on the reverse simulation. For pretraining p-1 we use our proposed training setup with the multi-step loss and Smax = 10 but freeze the score se (x, t) 0. Then, train the time-dependent score sa(x, 6) while freezing the reverse physics step. This approach guarantees that any time-independent physics are captured by P-1 and (x can focus learning small improvements to P-1 as well as respond possibly lime-dependent data biases. We give additional training details in appendix
Evaluation and training variants For evaluation, consider the MSE and spectral error of the reconstructed initial state Xg compared to the reference Xg. As baselines, during inference, we employ only the learned surrogate model p-1 without 5g. In addition that, evaluate variant for which we train both the surrogate model and sp(x, t) at the same time. As the two components resemble the drift and score of the reverse time SDE, this approach is similar DiffFlow (ZC21]] which learns both components in the context of generative modeling. We label this approach simultaneous training. Results are shown in figure 6 Similar to the stochastic heat equation results in section 4.2 the SDE achieves the best spectral error, while the ODE obtains the best MSE. Our proposed method outperforms both the surrogate model and the simultaneous training of the components.
5 Discussion and Conclusions
We presented combination of learned corrections training and diffusion models in the context of physical simulations and differentiable physics for solving inverse physics problems. We showed its competitiveness, accuracy, and long-term stability in challenging and versatile experiments and motivated our design choices. We considered two variants with complementary benefits for inference: while the ODE variants achieve the best MSE, the SDE variants allow for sampling the posterior and yield an improved coverage of the target data manifold. Additionally, we provided theoretical insights that the 1-step is mathematically equivalent to optimizing the score matching objective. We showed that its multi-step extension maximizes variational lower bound for maximum likelihood training.
Despite the promising initial results, our work has limitations that open up exciting directions for future work: Among others, requires simulating the system backward in time step by step, which can be costly and alleviated by reduced order methods. Additionally, we assume that At is sufficiently small and the reverse physics simulator is accurate enough. Determining good balance between accurate solutions with few time steps and diverse solutions with many time steps represents an important area for future research.
Acknowledgements BH, SV, and NT acknowledge funding from the European Research Council (ERC) under the European Union's Horizon 2020 research and innovation programme (LEDA: grant agreement No 758853, and SpaTTe: grant agreement No 863850). SV thanks the Max Planck Society for support through a Max Planck Lise Meitner Group. This research was partly carried out on the High Performance Computing resources of the FREYA cluster at the Max Planck Computing and Data Facility (MPCDF) in Garching operated by the Max Planck Society (MPG).
10

---

References
[And82]
B. Anderson. "Reverse-time diffusion equation models". In: Stochastic Processes and their Applications 12.3 (1982), pp. 313-326.
[Bar+19]
Y. Bar-Sinai, S. Hoyer, J. Hickey, and M. Brenner. "Learning data-driven discretiza- tions for partial differential equations". In: Proceedings of the National Academy of Sciences 116.31 (2019), pp. 15344-15349.
[BLL02]
X. Blanc, C. Le Bris, and P.-L. Lions. "From molecular models to continuum mechan- ies". In: Archive for Rational Mechanics and Analysis 164.4 (2002), pp. 341-381. [BWW22] J. Brandstetter, D. Worrall, and M. Welling. "Message passing neural PDE solvers". In: International Conference on Learning Representations (2022).
[Che+18a] R. T. Chen, Y. Rubanova, J. Bettencourt, and D. K. Duvenaud. "Neural ordinary differ- ential equations". In: Advances in neural information processing systems 31 (2018).
[Che+18b] T. Q. Chen, Y. Rubanova, J. Bettencourt, and D. Duvenaud. "Neural Ordinary Dif- ferential Equations". In: Advances in Neural Information Processing Systems. 2018, pp. 6572-6583 URL: httta::://pprcceeddngg 69386f6bb1df
neurips cc/ wwt.cc/ppppppe/.../ /
[Che+21] N. Chen, Y. Zhang. H. Zen, J. Weiss, al. "WaveGrad: Estimating Gradients for Waveform Generation". In: 9th International Conference on Learning Representations. OpenReview.net, 2021. URL: /httpss//oppnreeeee
[Chu+22]
H. Chung, B. Sim, D. Ryu, and C. Ye. "Improving Diffusion Models Inverse Prob- lems using Manifold Constraints" In: NearIPS. 2022. URL: h papp cc /apeer////// / 2022 / hash c44e5777ccoreea11111990002220000 Conference html
Abstract
[Chu+23]
H. Chung. J. Kim, M. T. McCann, M. L. Klasky, and J. C. Ye. "Diffusion Posterior Sam- pling for General Noisy Inverse Problems". In: The Eleventh International Conference on Learning Representations. OpenReview.net, 2023. URL: https: /ttps://openrevieu net/
[CSY22]
H. Chung, B. Sim, and J. C. Ye. Cfome-Closer-Diftuse Faster: Accelerating Conditional Diffusion Models for Inverse Problems through Stochastic Contraction". In: IEEE/CVF
Conference on Computer Vision and Pattern Recognition, CVPR 2022, New Orleans,
LA. USA, June 18-24, 2022. IEEE, 2022. pp. 12403 12403-12412.10 DOL: 10. 1199/CVPR52688 2022. 2022.01209 URL: https:://de org/ 10. SL00//CPR52288 2022. 01209
[Del+18]
S. Delaquis, M. Jewell, I. Ostrovskiy, M. Weber, et al. "Deep neural networks for energy and position reconstruction in EXO-200". In: Journal of Instrumentation 13.08 (2018),
POS023.
[DN21]
Dhariwal and A. Q. Nichol. "Diffusion Models Beat GANs on Image Synthe- sis". In: Advances in Neural Information Processing Systems. 2021, pp. 8780- 8794. URL: https / proceedings neurips cc / paper / 2021 hash / 19dd23dieeSSrabbcddd/a2222222 attp://aaaaaaaaaaaaaaaaaaattttte Abstract. html
[Far93]
S. J. Farlow. Partial differential equations for scientists and engineers. Courier Corpora-
tion, 1993.
[G6m+18]
R. Gómez-Bonee Bombarelli, J. N. Wei, D. Duvenaud, J. M. Hemnández-Lobato, et al. "Auto- matic chemical design using data-driven continuous representation of molecules". In: ACS central science 4.2 (2018), pp. 268-276.
[Gur82]
M. E. Gurtin. An introduction to continuum mechanics. Academic press, 1982.
[HJA20]
Ho, A. Jain, and Abbeel. "Denoising Diffusion Probabilistic Models". In: Advances in Neural Information Processing Systems. 2020. URL: https : / procedings neurips cc paper / 2020 hash 445bcfees544f0006671111000000000 Abstract. html
[HKT22]
P. Holl, V. Koltun, and N. Thuerey. "Scale-invariant Learning by Physics Inversion", In: Advances in Neural Information Processing Systems 35 (2022), pp. 5390-5403.
[HI.211]
C. Huang, J. H. Lim, and A. C. Courville, "A Variational Perspective on Diffusion- Based Generative Models and Score Matching". In: Advances in Neural Information Process- ing Systems. 2021, pp. 22863-22876. URL: http:: / neurips...... cc/

---

[Ho+22]
J. Ho. T. Salimans, A. A. Gritsenko, W. Chan, et al. "Video Diffusion Models". In: NeurIPS. 2022. URL: http::/paperrs nips
hash/39235c56arr 13fo055aaddc5eeeeddee66666 cote html
[Hol+20]
P. Holl, V. Koltun, K. Um, and N. Thuerey. "phiflow: A differentiable pde solving framework for deep learning via physical simulations". In: NeurIPS Workshop. Vol. 2. 2020.
[H6p+22]
T. Höppe, A. Mchrjou. S. Bauer, D. Nielsen, and A. Dittadi. "Diffusion Models for Video Prediction and Infilling". In: CoRR abs/2206.0769) (2022). DOL: 10. 4 arXiv 2206. 07696 URL: tittps://dot org/ 0085000485500 arXiv
[HTK20]
P. Holl, N. Thuerey, and V. Koltun. "Learnnng Control PDEs with Differentiable Physics". In: 8th International Conference on Learning Reprexentations. OpenRe- veew.net, 2020. URL::
[Hyv05]
A. Hyvärinen. "Estimation of Non-Normalized Statistical Models by Score Matching". In: J. Mach. Learn. Res. 6 (2005), pp. 695-709. URL: http: /jmr org/ httgppaprr// v6 v6/hyvarinen05a. html
[KCT23]
G. Kohl, L. Chen, and N. Thuerey. "Turbulent Flow Simulation using Autoregressive Conditional Diffusion Models". In: arXiv:2309.011 01745 (2023).
[Kid+21]
P. Kidger, J. Foster, X. Li, and T. J. Lyons. "Neural SDEs as Infinite-Dimne Dimensional GANs". : In: Proceedings of the 38th International Conference on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event. Vol. 139. Proceedings of Machine Learning Research. PMLR, 2021, pp. 5453-5463. URL:
mlr. press/
v139/kidger21b.c html
[Koc+21]
D. Kochkov, J. A. Smith, A. Alieva, Q. Wang, et al. "Machine learning accelerated computational fluid dynamics". In: CoRR abs/2102.01010 (2021). URL: https : // arxiv org/abg/////ooooooo 2202. 01010
[KPB20]
Kobyzev, S. J. Prince. and M. A. Brubaker. "Normalizing flows: An introduction and review of current methods". In: IEEE transactions on pattern analysis and machine intelligence 43.11 1 (2020), pp. 3964-3979.
[KVE21]
B. Kawar, G. Vaksman, and M. Elad. "SNIPS: Solving Noisy Inverse Problems Stochas- tically". In: Advances in Neural Information Processing Systems. 2021, pp. 21757- 21769. URL:
https / proceedings neurips cc / paper 2021 hash /
wwtrcc-bbbttat.... html
[LCT22]
B. List, L.W.. Chen, and N. Thuerey. "Learned turbulence modelling with differentiable fluid solvers: physics- bassed loss functions and optimisation horizons". In: Journal of Fluid Mechanics 949 (2022), A25.
[Li+21]
7. Li, N. B. Kovachki, K. Arizzadenesheli, B. Liu, et al. "Fourier Neural Operator for Parametric Partial Differential Equations". In: 9th International Conference on Learning Representations, ICI. 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021. URL: http::ttt:::///tott http::::///peeevieeee net. htt.//oormm11-ccP550temmmmm
[Li+22]
S. Li, Z. Huang. T. Du, H. Su, et al. "Contact Points Discovery for Soft-Body Ma- nipulations with Differentiable Physics". : In: International Conference on Learning Representations (2022).
[LN89]
D. C. Liu and J. Nocedal. "On the limited memory BFGS method for large scale optimization". In: Math. Program. 45.1-3 (1989), pp. 503-528. DOI: 10 1007 / BFO1589116 URL: htt:::tttss:/o/o org/10. 1007 1007/BF015891166
[LP22]
J. Lim and D. Psaltis. "MaxwellNet: Physics-driven deep neural network training based on Maxwell's equations". In: Apl Photonics 7.1 (2022), p. 011301.
[Lu+21]
L. Lu, Jin, G. Pang, Z. Zhang, and G. E. Karniadakis. "Learning nonlinear operators via DeepONet based on the universal approximation theorem of operators". In: Nat. Mach. Intell. 3.3 (2021). pp. 218 229. DO1: 10. 10. 1038//42286-6 0177556--21--11 021-00302-3 5 URL:
https://doi org/ 10. 10008//42256--- 02300-
[Mar+96]
G. J. Martyna, M. E. Tuckerman, D. Tobias, and M. L.. Klein. "Explicit reversible inte- grators for extended systems dynamics". In: Molecular Physics 87.5 (1996), pp. 1117 1157.
[Mor+ 18]
J. Morton, A. Jameson, M. J. Kochenderfer, and F. Witherden. "Deep Dynamical Modeling and Control of Unsteady Fluid Flows". In: Advances in Neural Information Processing Systems. 2018
12

---

[MRO20]
D. Maoutsa, S. Reich, and M. Opper. "Interacting Particle Solutions of Fokker- Planck Equations Through Gradient-Log- Lognenstt Estimation". In: Entropy 22.8 (2020), p. 802. DOL: 000.3390/e2200 e22080022 URL.: https: h /oi. org/ 1///0..3330 ww3.3//22200000
[Øks03]
B. Øksendal. "Stoocasttcc differential equations". In: Stochastic differential equations. Springer, 2003, pp. 15-88.
[Pap+21]
G. Papamakarios, E. Nalisnick, D. J. Rezende, S. Mohamed, and B. Lakshminarayanan. "Normalizing flows for probabilistic modeling and inference". In: The Journal of Ma- chine Learning Research 22.1 (2021), pp. 2617-2680.
[Par21] [Pfa+20]
f. Pardoux. Stochastic Partial Differential Equations: An Introduction. Springer, 2021.
T. Pfaff, M. Fortunato, A. Manchez-Gonzale.. and P. Battaglia. "Learning mesh-based simulation with graph networks". In: International Conference on Learning Representa- tions (2020).
[Ra
Z. Ramzi, B. Remy, F. Lanusse, J. Starck, and P. Ciuciu. "Denoising Score-Matching for Uncertainty Quantification Inverse Problems". In: CoRR abs/201 1.08698 (2020). URL: URL: hhtts://aaxxiii org/
[SC20]
S. Schoenholz and F. D. Cubuk. "Jax md: framework for differentiable physics". In:
Advances in Neural Information Processing Systems 33 (2020), pp. 1 1442--1144.
[SE19]
Y. Song and S. Ermon. "Generative Modeling by Estimating Gradients of the Data Distribution". In: Advances in Neural Information Processing Systems. 2019, pp. 11895- 11907. URL: https / proceedings neurips cC / paper 2019 / hash / Abb.aatttbtttat..ttmm
[SLG22]
Salvi, M. Lemercier, and A. Gerasimovics. "Neural Stochastic PDEs: Resolution-
Invariant Learning of Continuous Spatiotemporal Dynamics". In: NeurIPS. 2022.
URL: http : //pperss nips cc paper 5, files / paper 2022 / hash / 091 49116620a04228966664444 wcw.://o55555.411009090000----a- www.eettttoooooooeeeeeoe hhtmll
[Soh+15]
J. Sohl-Dickstein, E. A. Weiss, N. Maheswaranathan, and S. Ganguli. "Deep Unsuper- vised Learning using Nonequilibrium Thermodynamics". In: Proceedings of the 32nd International Conference on Machine Learning. Vol. 37. JMLR Workshop and Con-
ference Proceedings, JMLR. org. 2015, pp. 2256-2265. URL: http://proceed
mlr. mlr.prres/v37/ hhch///shl-ddaa www.ediiiss.e.n.n..cmmmmmmmm html
[Son+19]
Y. Song, Garg, Shi, and S. Ermon. "Sliced Score Matching: A Scalable Approach to Density and Score Estimation". In: Proceedings of the Thirty-Fifth Conference on Uncertainty in Artificial Intelligence. Vol. 115, Proceedings of Machine Learning Research. AUAI Press, 2019, pp. 574-584. URL: http wwwpeeiiiinnmmmmmm..mmeeeemmm press/ v115/song20a. html
[Son+21a]
Y. Song, C. Durkan, I. Murray, and S. Ermon. "Maximum Likelihood Training of Score- Based Diffusion Models". In: Advances in Neural Information Processing Systems.
2021, pp. 1415-1428, URL: hhtpp::::: pttp://ed/o/ce/o/pp////oe neurips. htcc//appr/22222
Awweraatttttt..aaae.e html
[Son+21b] Y. Song, Sohi-Dickstein, P. Kingma, A. Kumar, et al. "Score-Based Generative Modeling through Stochastic Differential Equations". In: 9th International Conference on Learning Representations. OpenReview.net, 2021. URI.: https:::s : / openreview
net het//forumIdd-xIII11200
[Son+22]
Y. Song, L. Shen, L. Xing, and S. Ermon. "Solving Inverse Problems in Medical Imaging with Score-Based Generative Models". In: The Tenth International Conference on Learning Representations. OpenReview.net 2022. URL: https: //pprrevieeee i
[Sta+21]
K. Stachenfeld, D. B. Fielding, D. Kochkov, M. Cranmer, et al. "Learned Simulators for Turbulence". In: International Conference on Learning Representations. 2021. N. Thuerey, P. Holl, M. Mueller, P. Schnell, et al. Physics- Deep Learning. WWW,
[Thu+21]
2021. URI.: ht pccpppppe/ppyyyccccaeeeeeeeaaar
org.
[Tom+17]
Tompson, K. Schlachter, Sprechmann, and K. Perlin. "Accelerating Eulerian Fluid Simulation With Convolutional Networks". In: 5th International Conference on Learning Representations. Workshop Track Proceedings. OpenReview.net, 2017. URL: https://openreviers net/ http:pawwoooo..ttetooooommmoooop
13

---

[Um+20]
K. Um, R. Brand, Y. ( Fei, P. Holl, and N. Thuerey. "Solver-in-the-Loopp Learning from Differentiable Physics Interact with Iterative PDE-Solvers". In: Advances Neural Information Processing Systems. 2020. URI.: https::/ neurips cc/ .html
[VinlI]
P. Vincent. "A Connection Between Score Matching and Denoising Autoencoders". In: Neural Comput. 23.7 (201 1), pp- 1661-1674.
[YSM22]
R. Yang. Srivastava, and S. Mandt. "Diffusion Probabilistic Modeling for Video Generation". In: CoRR abs/2203 09481 (2022). D01: 10. 48550/arXiy 2ww.:2223..0000
URL: URL:/httss:// //dii. org 10. 48540/ arivv 2203. 09481
[2C21]
Q. Zhang and Y. Chen. "Diffusion normalizing flow". In: Advances in Neural Informa-
tion Processing Systems 34 (2021), pp. 16280-16291.
[ZDG96]
K. Zhou, Doyle, and K. Glover. "Robust and optimal control", In: Control Engineering Practice 4.8 (1996), pp. 1189-1190.
14

---

Appendix
A Proofs and Training Methodology
Below we summarize the problem formulation from the main paper and provide details about the training procedure and the derivation of our methodology.
Let (R,,, P) be probability space and (1) (Wi()), (I)) be D-dimensional Brownian motion. Moreover, let Xo be Fo-measurable IR -valued random variable that is distributed as po and represents the initial simulation state. We consider the time evolution of the physical system for 0 s < T modeled by the stochastic differential equation (SDE)
Problem setting
dx d( + g())dW
(11)
with initial value Xy and Borel measurable drift P RD - RD and diffusion [0. T| + R20- This SDE transforms the marginal distribution Pa of initial states time the marginal distribution Pr of end states at time T.
Moreover, assume that have sampled N trajectories of length M from the above SDE with fixed time discretization 0 S Lo < l < < La S for the interval [0, T and collected them in assume that all time steps are equally spaced,
training data {xx, i.e, fm+11 Im == Al. Moreover, in the following we use the notation X for 0 S E < m S M to refer the trajectory (X1, 1te )
(n) Ad === For simplicity,
Assumptions Throughout this paper, we make some additional assumptions to ensure the existence of unique solution to the SDE 11 and the strong convergence of the Euler-Maruyama method. In particular:
Finite variance of samples from Po: Ex Xo [ oc
Lipschitz continuity of P: 3K; Vx, y RD ||((x) P(y) S Ki|[x yll2
Lipschitz continuity of g: 3K2 > 0 Ot,, E [0,7] (g(t) g(s)l I Kalt s|
Linear growth condition: 3K3 > 0' E RD: ||P(x)]l2 & Ks(1 + +||x||))
g bounded: 3K₁
E [0, T] [g(t)) < K₁
Maruuaaaa Method Using Euler-Maruyama steps, we can simulate paths from SDE [I]
similar to ordinary differential equations (ODE). Given an initial state X we let X X,, and define recursively
X t
XA + AtP (..)) VAt gtmm) stm!
(12)
where 2t,, are i.idd. with 2tm.. solution of the Euler-Maruyan Method as X := XAA Let X, denote the solution of the SDE
(). For t; < t < f,41, we define the piecewise constant
(TT. Then the Euler- Maruyama solution X converges strongly X,,
Lemma A.1. [Strong convergence of Euler-Maruyann method] Consider the piecewise constant
solution X of the Euler-Maruyane method. There is constant C such that
sup E|||X, - CVAL. 0s+ST
(13)
Proof. See Kloeden et al. [Klo+92 p. 10.2.2]
A.1 1-step Loss and Score Matching Objective
(x)) VN Theorem A.2. Consider data set {xx Frmm with trajectories sampled from SDE 111) Then the 1-step loss
- Σ m==0
Lsingle (0) = M
B
xm fxm+1 + At(-P(x,,, m x fm (
(14)
15

---

is equivalent to minimizing the score matching objective
/
'
-(0)
Ex [1 log
(15)
where 50(x;t)=so(xx t) as At +
Proof.
"e": Consider 6* such that Sg- (x. 1) = Vx (x) which minimizes the score matching
objective J().. Then fix a time step and sample x₁ and Xt+ from the data set. The probability flow solution x based equation [4 is
x == Xe+ + At [-P(xx+ac) + sa(xi=At, + At)]
(16)
the same time, we know that the transformation of marginal likelihoods from Pe+Ar to Pc follows the eeeers-ttmee SDE [And82]
dx = (P(x) + ( (x) dt + g()dW,
(17)
which runs backward in time from to O. Denote by x the solution of the Fuler-Maruyama method at time t initialized with X at time + At.
Using the triangle inequality for squared norms, we can write
lim IE |||x
x111 S 2 lim E |||x | + lim E ||| x111
(18)
Because of the strong convergence of the Fuler-Maruyama method, we have that for the first term of the bound in equation (18)
lim E |||x x At+0
0
(19)
independent of 0. At the same time, for the Fuler-Maruyana discretization, we can write
& Xx+ + At [- + g² + At) log (o)p:ta(((++a))
(20)
+g(t + At) At2+4b
(21)
where 2++ is a standard Gaussian distribution, i.e., 2++ No 7). Therefore, we can simplify the second term of the bound in equation (T8)
2
lim E At
x
(22)
|| At)² Vvx logpe
lim E Expaar A0-0
(23)
5x(xxt At) g(t At)
(24)
If 0* minimizes the score matching objective, then sq- (x, Vx log pe (x), and therefore the above is the same as
E. gt + At)
(25)
At
+0
Combining equations (88)
19
and
25
yields
lim
E [||x
x
0.
(26)
Additionally, since g is bounded, we even have
EfllvAt + At) x| < ElI K2 = A
(27)
For Lsingle (0), using the above bound (7 and strong convergence of the Euler- Euler- Maruyama method, we can therefore derive the following bound
Lsinglo (0)
1
E
Σ
[mm + Al (xxm4)) so (x bm []][[]]
(28)
M
<
M 2
CV At + E. K. 1
(29)
M
which implies that Caingle (0) + 0 as At 0
16

---

"=": With the definitions from "=", let 0 denote minimizer such that Leingle (0) - 0 as At 0, i.e., we assume there is a sequence At,, Atg.... with lim At, and a sequence 01, 02,,,... where 6n is global minimizer to the objective a (0) that depends on the step size Aln If there is 0* such that 88- (x, = Vx log Pa (x)), then singie (fn.) S siugle (f*). From "=" we know that lim,- ssingle (0*) = 0 and therefore
limn- Linggle (0n) O. This implies that each summand of Caingle (0) also goes zero as Al + 0, i.e. lim [||x x 0. Again, with the triangle inequality for squared norms, have that
lim E ||| x11 2 lim E ||
" +2 lim E
[||x:
x x1:1
(30)
At-0
Ar-1
Ar-0
By the strong convergence the Euler- Mareyyma method and 0", obtain
lim
E
[|]*
x=1 =0
(31)
Ar-0
At the same time, for fixed At > we can compute
E
[|| x x
(32)
R M (1 lt + At)² V log Pt (x(++()))
(33)
5x + At) Al (( + At) 2|
(34)
At + Bx) N g t- At) V Po
(35)
( At) + 21
(36)
For fixed Xi+Ae. the distribution over 2 ~~ (0, [0 I) in equation (36) correspond to non-
central chi chi-squared distribution [JKB95] Chapter 13.4], whose mean can be calculated as B2~N(O.)
|| g(t At)s/2 V (x so(xI4Ant At)]
2
(37)
= =|| g(t + At) [vx log 80(X Ar) + D.
(38)
For each At > 0, the above is minimized if and only if Vx Vx 50(X++Ar, At).
=
A.2 Multi-step Loss and Maximum Likelihood Training
We now extend the 1-step formulation from above to multiple steps and discuss its relation to maximum likelihood training. For this, consider our proposed probability flow ODE defined by
dx
[P(x) + (x) dt.
(39)
and for < <t define Pt, probability flow with network s0(x, t) equation (7) backward in time until 4,. We can choose two arbitrary time points 1, and by with < ty because we do not require fixed start and end times of the simulation.
,DE as the distribution obtained by sampling x ~ Pt, and integrating the
The maximum likelihood training objective of the probability flow ODE can be written as maximizing
Pty. ODE
Exx, ~Pe llg Pt.
(x..)-
(40)
Our proposed multi-step loss based on the sliding window size S, which is the length of the sub-trajectory that we aim to reconstruct with the probability flow ODE (7). The multi-step loss is defined as
Fs 5+1
1
Liaati (0) =
Exo
Σ
|| - 8 8 +
(41)
M
m==
17

---

where we compute the expectation by sampling Xo:m from the training data set and - is the predicted sub-trajectory that defined recursively by
X X and 8++5-1-j X ( + so( -it-s-sj------- (++5-))] (42)
In the following, we show that the multi-step loss maximum likelihood training objective (40)
maximizes variational lower bound for the
Theorem A.3. Consider data set (Xomm y with trajectories sampled from SDE (I). Then the multi-step loss (4) maximizes variational lower bound for the maximum likelihood training objective of the probability flow ODE (40) as At +
Let (x) denote the solution of the probability flow ODE integrated backward from time ty t with initial value Xt,
For the maximum likelihood objective, we can derive a variational lower bound
Ex,,
log Pt.
t.. ODE (xx..) Exx, log
[xx,
[Pt.
ta ODE,
(xt, (xx,,) (x),)))]]]
(43)
Excrs
log
Exxe, xx+=
Pt. (xx,)
Pti
1s. ODE,
(xe, (x,)
(44)
Pr, (Xe, (x,))
Pc, (xx,)
2 Exs, Exe, xx:s
|x
log
Pti
ts ODD
(xi, (x(,))
(45)
Pt, (2,,, (xa,)
Pt, (I<e,)
-Ex,
Exc, 1xs
log
+ log
is ,DD (P
(xts (x,;))
(46)
Pt. (X1, xxx;)
where the inequality is due to Jensen's inequality. Since only Pi: the same as maximizing
,ODE (x, (xe,) depends on 8 this is
Ex, Ex,, |xc, [log (pt: ,DDE (xx, (xx, )x
(47)
The probability flow ODE is likelihood-free. which makes it challenging to optimize. Therefore, we relax the objective by perturbing the ODE distributions by convolving them with a Gaussian kernel G. with small > 0, see, c.g., Kersting et al. [Kr++20 Gaussian ODE filtering]. This allows us to model the conditional distribution p!? ,DE Xt, Gaussian distribution with mean p4
f, ODE (xx,)
and variance 6² = Then maximizing minimizing
47
reduces to matching the mean of the distribution, i.e.,
Ex, Exy, jx, [||x,,
Me, (x,, 112
(48)
independent of E > O. Since this is true for any time step tj > t and corresponding simulation state (X,,, x2, (x,,) and on. Then, we can
x,, given X we can pick the pairs (xx,, xt-1). (X1, optimize them jointly by considering the sum of the individual objectives up to a maximum sliding window size
Exsd Σ
||xt
ODE Pe
(x 112
(49)
As At + O, we compute the terms 1o ODE (x,, ) on a single trajectory starting at x,, with sliding window covering the trajectory until Xx, via the Euler method, i.e., we can define recursively
8 X and 8 k + At [ +
(50)
By varying the starting points of the sliding window X1,, this yields our proposed multi-step loss (f).
A.3 Denoising Score Matching for Deterministic Simulations
So far, we have considered physical systems that can modeled by an SDE, i.e., equation (I) While this problem setup is suitable for many scenarios, would also like to apply similar
18

---

+ AtP(x.)) + V
+ A P(x.)
At P(x.)
Forwards backwarda
Training data
Xea
X4
X1.t
X
Xtul
Inference
X,
Xan
X4
Xais
X1
Xt,1
(P(x...)
solsa fi-1)]
At ((x)..)) sy(xancstes) (enl)
(8
At Pixt..) + V Biai
+ 2h.
(a)
SDE for training daza and reverso-time SDE inference
ODE for training data and SDE for inforence
ODE for training data and (e) physics + demoising for inferencs
Figure 7: Variants of training and inference for different physical systems. (a) shows the SDE and reverse-time SDE setup with the Fuler-Maruyama discretization when the system modeled by an SDE. The diffusion term g() is absorbed in the Gaussian random variable 2, NO, g(t) r) and network (x t). In (b), we assume that the temporal evolution of the training data is deterministic. i.e. we model the physical system without the diffusion term. However, for inference, we consider the reverse-time SDE of the same form as in (a), where the diffusion coefficient g(() is chosen as a hyperparameter that depends on the noise scale added to the data. Then, in (c), we split the Euler step for the backward direction into a physics-only update, adding the Gaussian noise and a denoising step by se (x,,t)). t).
methodology when the system is deterministic, i.e., when we can write the problem as an ordinary stochastic equation
dx P(x)dt.
(51)
In the case of chaotic dynamical systems, this still represents hard inverse problem, especially when information lost due noise added to the trajectories after their generation.
The training setup based on modeling the physics system with an SDE is shown in figure 7a. Figure 7p and 7 illustrate two additional data setup and inference variants for deterministic physical systems modeled by the ODE (5T). While for the experiments in sections 3.1 and 3.2 in the main paper, our setup resembles (a), for the buoyancy-driven flow in section 3.3 and the forced isotropic turbulence in section 3.4 in the main paper, we consider (c) as the system is deterministic.
For this variant, the update by -x() and (x, is separated into two steps. The temporal evolution from 44+1 to tú is then defined entirely by physics. We apply an additive noise to the system and the update step by se (x, t), which can be interpreted as denoising for a now slightly perturbed state Xt.- In this case, we show that the network s9(x, t) still learns the correct score Vx log pr (x) during training using denoising
matching. We compare the performance of variants (b) and (c) for the
buoyancy-drive flow in appendix E
When separating physics and score updates, we calculate the updates as
R, = Xt++1 AtP((
(52)
xnoie E2 Vtt gt:)) 24,
(53)
quoise + At (t.)) so (o .ta); X,, X
(54)
where 21, ~ N(0 I). If the physics system is deterministic and At is small enough, then we can approximate Xt, Re Xe, and for the moment, assume that we can write
noise &
Xt: + Atgg(t)) 2t.-
(55)
In this case, we can rewrite the 1-step loss Calngle(() from I4 to obtain the denoising score matching loss
CDSM(() : E E(xx, x [[|xe -
Quoise Xt;
At (ttg sy (Xcooise 12]
(56)
which is the same minimizing
Exx,; x411))
||10 (xproi (2))
1
(xt: *e )
(57)
At (tt))
19

---

Now, the idea presented in Vincent [ is that for score matching, we can consider joint distribution Pt, (x x of sample X₁, and corrupted sample x,, Using Bayes' rule, we can write Pt, (Xt,, xx,) = Pr [x, X1, Pt, (xt,) The conditional distribution Pa (-1 xx,)) for the corrupted sample is then modeled by Gaussian with standard deviation v g(;), i.e., we can write
x x + gt;) for (0, similar equation (55) Moreover, can define the distribution of corrupted data (a as
(x(x)) = Po (x(())x)e
(58)
If a is small, then 4o 2 Pt, and (L Pt.) - Oas a + 0. Importantly, in this case, we can directly compute the score for Po (-1 x) as
Vx Po (& x) σ² x).
(59)
Moreover, the theorem proven by Vincent [Vin11] means that we can use the score of the conditional distribution pa (-Ix) to train 98 (x, t) to learn the score go (x), i.e.
arg min Ex [x) 4;) Vx log gr (x)
(60)
arg min E Ex-p,, E-po x)) []|sg(X, tx) - Vxlogoo (x [xxx)
(61)
By combining 61 and (59). this exactly equals the denoising loss LDSM (0) in 57 As go = Pt. we also obtain that Vx log go (x) Vx log Pt, (x), so the network sy tj) approximately learns the correct score for Pt,-
We have assumed (55 that the only corruption for &noise is the Gaussian noise. This is not true, as we have
noise 4
Xx, + g(4) 24, + (Xi,, - PIP( xx,),
(62)
so there is additional source of corruption, which comes from the numerical errors due to the term Xt.t1 At P ( Xr. The conditional distribution (-]x)) is only approximately Gaussian. Ideally, the effects of numerical errors are dominated by the Gaussian random noise. However, even small eтrоrs may accumulate for longer sequences of inference steps. To account for this, we argue that the multi-step loss (0) should be used. During training, with the separation of physics update and denoising, the simulation state is first progressed from time f++1 to time i; using the reverse physics solver. This only yields perturbed version of the simulation time t, due to numerical inaccuracies. Then a small Gaussian noise is added and, via the denoising network so(x, t). the simulation state is projected back the distribution pt,, which should also resolve the numerical errors. This is iterated, as discussed in section 2 in the main paper, depending on the sliding window size and location.
20

---

B Architectures
ResNet We employ simple ResNet-like architecture, which is used for the score function so(x. t) and the convolutional neural network baseline (ResNet) for the stochastic heat equation in section 3.2 as well in section 3.4 again for the score Sy (x,t).
For experiments with periodic boundary conditions, we apply periodic padding with length 16, i.e., the underlying 2-dimensional data dimensions are N N, the dimensions after the periodic padding are (N + 16) x (N + 16). We implement the periodic padding by tiling the input three times in x- and y-direction and then cropping to the correct sizes. The time is concatenated constant channel to the 2-dimensional input data when this architecture represents the score sy (x, t).
additional
The encoder part of our network begins with single 2D-convolttion encoding layer with 32 filters, kernel size 4, and no activation function. This is followed by four consecutive residual blocks, each consisting of 2D-convolution, LeakyReLU, 2D-convoltioo, and Leaky ReLU. All 2D convolutions have 32 filters with kernel size four and stride The encoder part ends with a single 2D convolution with one filter, kernel size 1, and no activation. Then, in the decoder part, begin with transposed 2D convolution, 32 filters, and kernel size 4. Afterward, there are four consecutive residual blocks, analogous the residual encoder blocks, but with the 2D convolution replaced by transposed 21) convolution. Finally, there is final 2D convolution with one filter and kernel size of 5. Parameter counts of this model and other models are given in table 2
UNet We use the UNet architecture with spatial dropout as described in MMue+22] Appendix A. for the Bayesian neural network baseline of the stochastic heat equation experiment in section 3.2. The dropout rate is set to 0.25. We do not include batch normalization and apply the same periodic padding as done for our ResNet architecture.
The FNO-2D architecture introduced in Li+21 with kmax.j 12 Fourier modes per channel used as a baseline for the stochastic heat equation experiment section 3.2 and the learned physies surrogate model in section 3.4.
FNO
Dil-ResNet The Dil-ResNet architecture described in [Sta+21), Appendix A. This architecture represents the score (x, t) in the buoyancy-driven flow with obstacles experiment in section 3.3. We concatenate the constant time channel analogously the ResNet architecture. Additionally, positional information is added to the network input by encoding the x--ossitonn and y-position inside the domain in two separate channels. Architecture
Parameters
ResNet
330 754
UNet
706 035
Dil-ResNet
336 915
FNO
465 377
Table 2: Summary of architectures.
21

---

C ID Toy SDE
For the ID toy SDE discussed in section 3.1, we consider the SDE given by
dr = A sign(>)a"] dt + Azdw,
(63)
with A 7 and A2 0.03. The corresponding reverse-time SDE is
dx
[A sign(x) X V logpe(s)) di + X2dw.
(64)
Throughout this experiment, Po is categorical distribution, where we draw either 1 or - with the same probability. In figure & show trajectories from this SDE simulated with the Buler-Maruyama method. Trajectories start or -1 and approach 0 as increases.
Neural network architecture We employ a neural network so(z, t) parameterized by 0 to approx- imate the score via the 1-step loss, the multi-step loss, implicit score matching [Hyv05] ISM] and sliced score matching with variance reduction [Son+19] SSM-VRJ. In all cases, the neural network is simple multilayer perceptron with elu activations and five hidden layers with 30, 30, 25, 20, and then 10 neurons for the last hidden layer.
We use the Adam optimizer with standard hyperarrameter' described in the original paper [KB 15 The learning rate, batch size, and the number epochs depend on the data set size (100% with 2 500 trajectories, 10%, or 16) and are chosen to ensure convergence of the training loss.
1.0
3.0
0.5
6.5
-4.5
10
(a) A2 0
(b) )2 0.03
Figure 8: Trajectories from SDE 663 with X2 0 (a) and X2 0.03 (b).
Training 1-step loss For the I-step loss and all data set sizes, train for 250 epochs with learning rate of 10e-3 and batch size of 256. the first phase, we only keep every 5th point of a trajectory and discard the rest. Then, again train for 250 epochs with the same batch size and learning rate of 10e-4 but keep all points. Finally, we finetune the network with 750 training epochs and learning rate of 10e-5.
Training- multi-step loss For the multi-step loss and 100% of the data set, we first train with the 1 1-step loss, which resembles a sliding window size of 2. We initially train for 1 000 epochs with batch size of 512 and a learning rate of 10e-3, where and discard the rest. Then, with a decreased learning rate of 10c-4, we begin training with a sliding window size of S and increment every 1 000 epochs by one until Smax 10. In this phase, we train on all points without any removals.
keep only every 5th point on trajectory
Training ISM For ISM. we compute the partial derivative Ösp (x)(x): (x):/xxx /xx; using eeverse-mode automatic differentiation JAX (jax.jacrev). For 100% and 10% of the data set, we train for 2 000 epochs with learning rate of 10e-3 and batch size of 10 000. Then we train for an additional 2 000 epochs with learning rate 10e-4. For 1%, increase the number of epochs 20 000.
Training SSM-VR For sliced score matching with variance reduction use the same training setup for ISM.
Son+19 SSM-VRJ we
22

---

75
75
10
50
50
25
0.0
0.0
-0.5
-25
-0.5
-2
-10
-1.0
2
6
10
-75
2
6
10
-75
(a) ISM learned score.
(b) Multi-step learned score.
1.0
1.0
0.0
0.5
10
8
10
8
(c) reverse-time SDE trajectories.
(d) Multi-step reverse-time SDE trajectories.
1.0
0.0
0.0
-0.5
-0.5
-1.0
-1.0
10
10
8
(e) ISM probability flow trajectories.
(f) Multi-step probability flow trajectories.
Figure 9: Comparison of Implicit Score Matching (ISM, left) and proposed training with the multi-step loss (Multi-step, right). Colormap in (a) and (b) truncated to [-75, 75].
Method
Probability flow ODE
Reverse-ting time SDE
Data set size
Data set size
100%
10%
1%
100%
10%
1%
multi-step
0.9710.04
0.91=10.05
0.8110.01
0.99-0.01
0.94
0.85
1-step
0.78
0.4410.13
0.41-0.13
0.93=0.05
0.71
0.75+0.11
ISM
0.19:10.05
0.15+0015
0.01
0.94110.1 0.92-0.05
0.52:10.22
SSM-VR
0.17=0.16 0.49=:022 0.27+0.47
0.88=0.06
0.94 0.67+
Table 3:
Posterior metric Q for 000 predicted trajectories averaged over three runs.
Comparison We directly compare the learned score for the reverse-time SDE trajectories and the probability flow trajectories between ISM and the multi-step loss figure D trained on the full data set. The learned score of ISM and the multi-step loss in figure Da and figure 9b are visually very similar, showing that our method and loss learn the correct score. Overall, after finetuning both ISM and the multi-step loss. the trajectories of the multi-step loss more accurate compared to ISM. For example, in figure Pe a trajectory explodes to negative infinity. Also, trajectories from the multi-step loss end in either -I or 1, while ISM trajectories attenuated and do not fully reach -1 or exactly, particularly for the probability flow ODE.
Results of Table in Main Paper We include the standard deviations of table 1 from the main paper table 3 above. The posterior metric Q is very sensitive the learned score sg (xx L). Overall, proposed multi-step loss gives the most consistent and reliable results.
23

---

For the quadratic SDE equation 63 the analytic score non-trivial, therefore direct comparison of the learned network s,(x, and the true score Vx (ogp(()) is difficult. However, we can consider the SDE with affine drift given by
Empirical verification of Theorem 3.1
dr = Ardr + gdW
(65)
with 0.5 and g = 0.04. Because this SDE is affine and there are only two starting points and 1, we can write the distribution of states starting in zig Gaussian with mean ((1:20) Jige -
and variance (²; zo) / (1 e 2, see [ks03). Then, the score at time and position x conditioned on the starting point To is (x - [u(:; To (o²)) zo) See figure 10 for a visualization of the analytic score and comparison with the learned score. The learned score from the 1-step
training matches the analytic score very well in regions where the data density is sufficiently high.
1.0
1.0
50
0.5
0.5
0.0
0.0
0
-05
-0.5
-50
-1.0
-1.
(a) Paths density
(b) Analytic score field
1.0
1.0
50
0.5
0.5
0.5
-0.5
-50
-1.0
8
10
0
2
10
(c) 1-step training
(d) Multi-step training
Figure 10: 1-step training score matches analytic score. (a) shows some paths sampled from SDE equation (65) and contour of the density. (b) is the analytic score field, and (c) and (d) are visualizations of the learned score with 1-step and multi-step training. Scores from the multi- step training correspond to more narrow high-density regions. This implics that during inference, trajectories are pulled more strongly to modes the training data set distribution than for the 1-step training.
24

---

D Stochastic Heat Equation
Spectral solver The physics of the 2-dimensional heat equation for x E Rdddd can be computed analytically. The solver PAi(x (x) simulates the systems x forward in time by fixed At using the (shifted) Fourier transformation F. In particular, we can implement the solver with
(x) = F (A(At) F(x)),
(66)
where denotes clement-wise multiplication and A(At);; E ddddd is matrix with entries A A(At)i, := exp (-At min(i, (i, j. - 1..j-71 i)). The power spectrum of x is scaled down by A(At), and higher frequencies are suppressed more than lower frequencies (for At 0). If noise is added to
x, then this means that especially higher frequencies are affected. Therefore the inverse transforma-
tion (when At > 0) exponentially scales contributions by the noise, causing significant distortions for the reconstruction of x.
Spectral loss We consider a spectral error based on the two-dimensional power spectral density. The radially averaged power spectra s and 82 for two images are computed as the absolute values of the 2D Fourier transform raised to the second power, which are then averaged based on their distance (in pixels) to the center of the shifted power spectrum. We define the spectral error as the weighted difference between the log of the spectral densities
S2) = Σ wkl log($1.)) log(s2,*)
(67)
with a weighting vector w E Rd and wi = 1 for h S 10 and w. = otherwise.
Training For inference, we consider the linear time discretization t,, = nAt with At = 0.2/32 and t32 0.2. During training, we sample a random time discretization S to < t/2 < .... < tg t32 for each batch based on fn via t ~ (t,,, 2t/2, t, + At/2) for n = 1,....31 31 and adjust the reverse physics step based on the time difference 1 111-1- 1. In the first training phase, we consider the multi-step loss with sliding window size of S = 6, 8, ...., 32 steps, where we increase S every two epochs. We use Adam to update the weights 0 with learning rate 10-1, We finetune the network weights for 80 epochs with an initial learning rate of 10-1, which we reduce by factor of 0.5 every 20 epochs.
8o only version For the 1-step loss, this method is similar to Rissanen et al. [RHS22], which proposes a classical diffusion-like model that generates data from the dynamics of the heat equation. Nonetheless, the implementation details and methodology are analogous to the multi-step loss training, except that the reverse physics step P-1 is not explicitly defined but instead learned by the network sx(x, t) together with the score at the same time. We make use of the same ResNet architecture as the default variant. Except for the reverse physics step, the training setup is identical. Although the network s0 (x, t) is not trained with any noise for a larger sliding window with the multi-step loss, we add noise to the simulation states for the SDE inference, while there is no noise for the ODE inference.
Baseline methods All other baseline methods are trained for 80 epochs using the Adam optimizer algorithm with an initial learning rate of 10-4, which is decreased by a factor of 0.5 every 20 epochs. For the training data, we consider solutions to the heat equation consisting of initial state xo and end state X7.
Test-time distribution shifts equation experiment. We train the score network 8p for a specific combination of diffusivity a and noise g and vary both parameters for testing. We modify both the simulator and test ground truth for the updated diffusivity and noise. See figure I1 Overall, for small changes of the parameters, there is very little overfitting. Changes in the reconstruction MSE and spectral error can mainly be explained by making the task itself easier or harder to which our network generalizes nicely, e.g., less noise or higher diffusivity leads to smaller reconstruction error.
We have tested the effects of test-time distribution shifts for the heat
Additional results
We show visualizations for the predictions of different methods in figure 12 and
figure 1
25

---

10a
10
10a
MSE
10
10
10a
10
10.
10a1
10"
Spectral error
10a
10a
-30%
-20%
-10%
1.0
+10%
+20% +30%
-30%
-20%
-10%
+10% +3%%%%%%%00%
Diffusivity
Noise
(a) Reverse-time SDE
10
10a
104
MSE
10
10
10
10
109
Spectral error
101
10
10
-30%
-20%
-10%
1.0 +10%
20%
30%
-30%
-20%
-10%
+10%+20%++
Diffusivity
Noise
(b) Prohability flow ODE
Figure 11: Test-time distribution shifts for noise and diffusivity. The correction network so is trained on diffusivity 1.0 and noise g = 0.1. During testing, we vary the diffusivity of the data and simulator as well as the noise for inference. Especially the low spectral errors indicate that the network generalizes well the new behavior of the physics.
E Buoyancy-driven Flow with Obstacles
We semi- semi-Lagrangee Lagrangian advection for the velocity and MacCormack advection for the hot marker density within a fixed domain C. [0, 1] [0, 1]. The temperature dynamics of the marker field are modeled with Boussinesq approximation.
Training We train all networks with Adam and learning rate 10-4 with batch size 16. We begin training with a sliding window size of S = 2, which we increase every 30 epochs by until Smax 20.
Separate vs. joint updates We compare joint update of the reverse physics step and corrector function 8g, see figure and a separate update of reverse physies step and corrector function, see figure 7c. An evaluation regarding the reconstruction MSE and perceptual distance is shown in
26

---

(a) Ground truth
(b) Input
(c) SMDP-ODD
(d) ResNet
(e) FNO
8g only ODE
(g) 80 only SDE
(h) BNN
SMDP SDE
Figure 12: Predictions of different methods for the heat equation problem (example of 2).
(a) Ground truth
(b) Input
(c) SMDP ODE
(d) ResNet
(e) FNO
(f) 89 only ODE
(g) 80 only SDF
(b) BNN
SMDP SDE
Figure 13: Predictions of different methods for the heat equation problem (example 2 of 2). Neither the BNN nor the "sg only" model can produce small-scale structures.
figure 14 Both training and inference variants achieve advantages over "p-1 only" and "sø only"
27

---

approaches, Overall, there are no apparent differences for the ODE inference performance but slight benefits for the SDE inference when separating physies and corrector update.
Limited-memon BFGS We use numerical optimization of the marker and velocity fields at t = 0.35 match the target smoke and velocity fields at 0.65 using limited-memory BFGS [LN89] LBFGS] and the differentiable forward simulation implemented in phiflow (Hol+20]. Our implementation directly uses the LBFGS implementation provided by torch.optim. LBFGS [Pa++ 19]. As arguments, we use history_size
10, max max_iter = and line_search ffarcch__n_ strong wolle. Otherwise, we leave all other arguments to the default values. We optimize for steps which takes 240 seconds per sample on single NVIDIA RTX 2070 gpu.
Diffusion posterior sampling DPS An additional baseline for this problem diffusion posterior sampling for general noisy inverse problems [Chu+23 DPSJ. As first step, we pretrain diffusion model on the data set of marker and velocity fields at 0.35. We use the mask for obstacle positions additional conditioning input to the network to which no noise is applicd. Our architecture and training closely resemble Denoising Diffusion Probabilistic Models (HJA20) DDPM]. Our network consists of ca. 18.44 million parameters trained for 100k steps and learning rate 1 10 using cosine annealing with warm restarts (Tp- = 50000, Mauin = x 10 9). The measurement operator A is implemented using differentiable forward simulation. We consider the Gaussian version of DPS, i.e., Algorithm in Chu+23) with N 1000. We fix the step size Çi at each iteration to 1/lly A [I: For each inference step, required backpropagate gradients through the diffusion model and the forward simulation. Inference for a single sample requires ca. 5000 seconds on single NVIDIA RTX 2070 gpu.
We provide more detailed visualizations for the buoyancy-driver flow case in figure 16 and figure [7 These again highlight the difficulties of the reverse physics simulator to recover the initial states by itself. Including the learned corrections significantly improves this behavior.
A Additional results
In figure 15 we also show an example of the posterior sampling for the SDE. It becomes apparent that the inferred small-scale structures of the different samples change. However, in contrast to cases like the heat diffusion example, the physics simulation this scenario leaves only little room for substantial changes in the states.
28

---

0.8
SMDP SDE
MSE (t (1=00) 0.35)
SMDP ODE
0.6
0.4
0.2
0.0
g-1
only
only
joint update
separate update
(a) Reconstruction MSE 0.35).
SMDP SDE
LPIPS 0.35)
SMDP ODE
10-1
10-2
--1 only
So only
joint update
separate update
(b) Perceptual distance LPIPS (t 0.35).
Figure 14: Comparison of separate and joint updates averaged over three
Ground truth
SMDP-SDE
Figure 15: Comparison of SMDP-SDE predictions and ground truth for buoyancy-drice driven flow at
t 0.36.
29

---

SMID ODE
(a) Marker field
SMDP Ground
Fp-1
0.65
0.55
t= 0.45
0.35
0.25
0.15
(b) Velocity field (x-direction)
Figure 16: Predictions for buoyancy-driven flow with obstacles (example of 2).
30

---

ODE
SMDF SDE
9-11 only
(a) Marker field
Ground
7p-1
r= 0.65
0.55
t= 0.45
0.35
0.25
0.15
(b) Velocity field (x-direction)
Figure 17: Predictions for buoyancy-driven flow with obstacles (example 2 of 2).
31

---

Ground truth
only
100 stops
200 steps
stops
4000 steps
steps
750 steps
1000 steps
2000 steps
Figure 18: Steps of Langevin dynamics for 10-5
Ground truth
P
only
100 sleps
200 steps
300 sleps
400 steps
500 steps
750 steps
1000 steps
2000 steps
Figure 19: Steps with Langevin dynamics for x 10-5 and an additional step with At 80 1)
which smoothes the images.
Isotropic turbulence
Training For the physics surrogate model P-1, we employ an FNO neural network, see appendix with batch size 20. We train the FNO for 500 epochs using Adam optimizer with learning rate 10 which we decrease every 100 epochs by factor of 0.5. We train s,(x, t) with the ResNet architecture, see appendix B for 250 epochs with learning rate 10 44 decreased every 50 epochs by factor of 0.5 and batch size 6.
Refinement with Langevin Dynamics Since the trained network sy sa approximates the score Vx log Pr (x), it can be used for post-pre post-processing strategies [WTI]] SE19]. We do fixed point iteration at a single point time based on Langevin Dynamics via:
++ x + Vx log + V2ez,
(68)
for number of steps T and 2 1 cf. figure 18 and figure 19 For prior distribution we-
x ~ Iy, and by iterating (68), the distribution of x equals pt for and T + Do. There are
some theoretical caveats, i.e.. Metropolis-Hasting update needs to be added in 68) and there are additional regularity conditions SE19
Additional results We show additional visualizations of the ground truth and reconstructed trajec-
tories figure and figure?
32

---

Ground truth
A only
SMOP ODE
SMOP-S08
Figure 20: Predictions for isotropic turbulence (example 1 of 2).
Cround truch
only
SHDP SMDP-GDE DDD
SMOP SDE
Figure 21: Predictions for isotropic turbulence (example 2 of 2).
33

---

References
[And82]
B. Anderson. "Reverse-time diffusion equation models". In: Stochastic Processes and their Applications 12.3 (1982), pp. 313-326.
[Chu+23]
H. Chung, J. Kim, M. T. McCann, M. L. Klasky, and J. C. Ye. "Diffusion Posterior Sam- pling for General Noisy Inverse Problems". In: The Eleventh International Conference on Learning Representations. OpenReview.net, 2023. URL: httpp:://ppemiii net. net/ pett/pdfid==0DDD2Ziia 2id==DD9GGGGAT
[HJA20]
Ho, A. Jain, and P. Abbeel. "Denoising Diffusion Probabilistic Models", In: Advances in Neural Information Processing Systems. 2020. URL: https p/pooeeddg/s neurips cc / paper / 2020 / hash 4cbbcfe66444f0066771111000776666 Abstract.html
[Hol+20]
P. Holl, V. Koltun, K. Um, and N. Thuerey. "phiflow: A differentiable pde solving framework for deep learning via physical simulations". In: NeurlPS Workshop. Vol. 2. 2020.
[Hyv05]
A. Hyvärinen. "Estimation of Non-Normalized Statistical Models by Score Matching",
In: J. Mach. Learn. Res. 6 (2005), pp. 695-709. URL: http://jmmii: org/ atg/papers/ v6 w6/hyvarinen05a. html
[JKB95]
N. L. Johnson, S. Kotz, and N. Balakrishnan. Continuous univariate distributions, volume 2. Vol. 289. John wiley sons, 1995.
[KB15]
D. P. Kingma and J. Ba. "Adam: A Method for Stochastic Optimization". In: 3rd International Conference Learning Representations, 2015. URL: htttp://axivv org/abs/1412 6980
[Ker+20]
H. Kersting, N. Krämer, M. Schiegg, C. Daniel, et al. "Differentiable Likelihoods for
Fast Inversion of "Likelihood-Free Dynamical Systems". In: Proceedings of the 37th International Conference on Machine Learning, ICMI. 2020, 13-18 July 2020, Virtual Event. Vol. 119. Proceedings of Machine Learning Research. PMLR, 2020, pp. 5198- 5208. URL: http::
mlr press/ vt///vv19//// 19/ 1ete////eeiing0aaa 2tmml
[Klo+92]
P. E. Kloeden, E. Platen, P. E. Kloeden, and E. Platen. Stochastic differential equations. Springer, 1992.
[Li+21]
Z. Li, N. B. Kovachki, K. Azizzadenesheli, B. Liu, et al. "Fourier Neural Operator for Parametric Partial Differential Equations". In: 9th International Conference Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021. URL: htttp:://ppenreviiiiee wwtpreevviiwneeeeeeeeeeeee nww/ffrumm wtforrmmiiaccPPPNoota..m
[LN89]
D. Liu and J. Nocedal. "On the limited memory BFGS method for large scale optimization". In: Math. Program. 45.1-3 (1989), pp. 503-528. DOI: 10 1007 / org/10. 10. 9007//B001589166
BF01589116 URL: httt:/ttts:///ae
[Mue+22]
M. Mueller, R. Greif, F. Jenko, and N. Thuerey. "Leveraging Stochastic Predictions of Bayesian Neural Networks for Fluid Simulations". In: arXiv preprint arXiv:2205.01222 (2022).
[Øks03]
B. Øksendal. "Stochastic differential equations". . In: Stochastic differential equations. Springer, 2003, pp. 65-84.
[Pas+19]
A. Paszke, S. Gross, Massa, A. Lerer, et al. "PyTorch: An Imperative Style, High- Performance Deep Learning Library". In: Advances in Neural Information Processing Systems. 2019, pp. 8024-8035. URL: http:tt://oooccediiiigg http:///bbbbaaeeeeffffffaaaaaaae
neur wwwuuipp.ccpppppae Abstract html
[RHS22]
S. Rissanen, M. Heinonen, and A. Solin. "Generative Modelling With Inverse Heat
Dissipation". In: 22XXv:22066.3397 (2022).
[SE19]
Y. Song and S. Ermon. "Generative Modeling by Estimating Gradients of the Data Distribution" In: Advances in Neurul Information Processing Systems. 2019, pp. 11895- 11907. URL.: httpss : proceedings neurips cc paper 2019 / hash /
1e8e2255447d55311119000677777777 Abstracttstttt.....mm html
[Son+19]
Y. Song. S. Garg, J. Shi, and S. Ermon. "Sliced Score Matching: A Scalable Approach to Density and Score Estimation". In: Proceedings of the Thirty-Fifth Conference on Uncertainty in Artificial Intelligence. Vol. 115. Proceedings of Machine Learning Research. AUAI Press, 2019, pp. 574-584. URL: http: http::/prooceeinggg...... mlr. mlr.pr v115/song20a. html
34

---

[Sta+21]
K. Stachenfeld, D. B. Fielding, D. Kochkov, M. Cranmer, et al. "Learned Simulators for Turbulence". In: International Conference on Learning Representations. 2021.
[VinlI]
P. Vincent. "A Connection Between Score Matching and Denoising Autoencoders". In: Neural Comput. 23.7 (201 1), pp- 1661-1674.
[WTH]]
M. Welling and Y. W. Teh. "Bayesian Learning via Stochastic Gradient Langevin Dynamics", In: Proceedings of the 28th International Conference on Machine Learning, ICML 2011, Bellevue, Washington, USA, June 28- July 2011. Omnipress, 201
pp. 68 1-8-68
35


Reviews:
[{"review_id": "sSr3rt1TTH", "reviewer": "Reviewer_SQFk", "summary": "The paper proposes an approach to sample from the posterior distribution based on score-based diffusion models with a particular focus on inverse problems in physics. \n\nThe proposed approach, to my understanding, has two novel contributions:\n- they use reverse-physics simulations and augment them with score estimates that supposedly leads to superior performance in sampling,\n- they propose a multi-step training regime where scores are estimated in sequential manner for a given time-horizon, in contrary to the standard single-step score estimation.\n\nThe method was tested on different synthetic data experiments, and both the proposed approaches seem to produce meaningful improvements. ", "questions": "- How is the ground truth obtained in the case of SDEs?\n- In Section 3.4 (Fig. 6), why is the spectral error worse for SDEs compared to ODEs while it is the other way around in other experiments?", "limitations": "- Releasing the code could be crucial as there are many delicate details one might need to get right for the method to work.", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- The multi-step training of the score function seems to be effective in synthetic data applications.\n- The theoretical results, although not particularly novel, provide a complete picture of the proposed methods. \n- The experiments presented validate the proposed ideas well.", "weaknesses": "- One of the novelties claimed in the paper is the use of score-function to simply _refine_ the outputs of a reverse-physics simulator. Can't one simply view the reverse-physics simulator as (non-learned) _part of_ the parametric score model? In this case, the claim simply becomes that a physics-informed model to approximate the score is better than one that is oblivious to the physics? Isn't this an unsurprising statement? This is the underlying motivation behind physics-informed neural nets (PINNs), a relatively large area of research. Could the authors clarify?\n- In some of the experiments, the authors use the LPIPS metric to evaluate the quality of the solution. This does not make much sense as LPIPS is designed to evaluate the \"perceptual\" quality of the image, which has nothing to do with the physical accuracy (the metric one cares about).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their valuable review and helpful suggestions.\n\n- **Refinement of model outputs**: We agree that the reverse-physics simulator and the score network can be unified in a single joint model. Nonetheless, a clear distinction between the simulator and score network makes sense from a methodological viewpoint, as the learned score is data-driven with little to no inductive bias at all, whereas the simulator includes only inductive biases and conservation laws. Also, during inference, the update step by the simulator can be regarded as a step between different points in time, whereas the score network corrects the simulation state by projecting it towards a region with higher likelihood at the same point in time. One can expect an improved performance from a physics-aware score-matching approach, but nonetheless this connection has not been made before in the form we present it. As such, we believe our results and theory provide an important basis for future work at the intersection of both methods.\n- **LPIPS metric**:  Our primary metric for simulation data is the L2 distance. However, the L2 distance can become very large when the prediction does not match the ground truth on a per pixel basis, even though the prediction might still be close to the ground truth visually. In these situations the LPIPS distance can be more informative as it is less sensitive to the exact pixel positions. \n- **Ground truth for SDEs**: We draw individual samples from the testing dataset at time 0 and simulate them forward in time. The prediction is compared to this sample via the spectral error and reconstruction MSE, which are both less sensitive to the exact pixel values of the sample. \n- **Spectral error in Fig. 6**: In this specific plot, there is a mistake and the colors for the SDE and ODE should be switched. We noticed this mistake shortly after submission but could not fix it any more. We apologize for the confusion.  \n"}, {"review_id": "sYeUPCghim", "reviewer": "Reviewer_h6fX", "summary": "This work describes a new method for solving inverse problems for time-evolved physical systems. It does this by combining two components: a time-independent reverse physics simulator (based on a priori domain knowledge or learned) and a learned time-dependent correction term. The posterior for the initial state can then be obtained by sampling from a related SDE (or a related ODE). This method is referred to as score matching via differentiable physics (SMDP). The authors argue that the learned correction term corresponds to the score (the gradient of the log-likelihood of the trajectory), providing a probabilistic interpretation of their model.\n\nThe method is numerically evaluated on four different physical systems, where it is shown to yield results superior to various baselines. This is the case for both the SDE and ODE variants of the method. The authors also discuss relations with existing methods for inference based on score matching as well as generative methods such as diffusion models.\n", "questions": "– How is the problem in Section 2 motivated? Is this supposed to be a discretization of an SDE?\n\n– Why is the step *added* to x_m in eq. (1) in order to obtain x_{m+1}? If this is stepping backwards, shouldn't it be subtracted? The conclusion on line 113 would make more sense in that case.\n\n– In Section 2, N trajectories indexed by i from 1 to N are introduced, but never make another appearance. Instead, eq. (2) and (3) refer to an expectation over trajectories. How are we to think of the trajectories, as samples or as a random vector?\n\n– Some more discussion as to the difference between the SDE and ODE variants would be useful. If I understand correctly, the former sampled the posterior states, while the latter provides an approximation of its mode (calculating the maximum a posteriori through a normalizing flow-type construction).\n\n– Figure 2 is hard to parse. In (a), are we observing the true score or its approximation by the correction term? Also, it is not clear what is shown in (c), especially the trajectories at the bottom. Are these exploding trajectories occurring for GRID or for MLP? Are they for one step or multiple steps? The discussion of these results (lines 178–186) are similarly hard to follow.\n\n– The two sentences on lines 212–215 seem to contradict one another (if the forward solver cannot be used, how are we implementing the reverse step using the forward solver?).\n\n– Why are the results of the SDE not shown in Figure 5?\n\n– Why are we adding Gaussian noise to each state in Section 3.3? If this is a deterministic system, should we not be considering deterministic methods to solve it?\n\n– The conclusion on lines 302–303 contradicts the results in Figure 6(b) (ODE obtains better spectral error, not MSE).\n\n– In various places, the authors use a period as a thousands separator. This should be a space in an English-language context.\n", "limitations": "The authors briefly discuss limitations (and possible future directions) in the last section. I don't believe there are any potential negative societal impact from the work that should be considered.\n", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "The paper is well laid out, clearly explaining the task considered and the proposed solution. Various quantities and models are well defined and motivated. The numerical section is also well presented, with appropriate comparison to baseline methods such as implicit score matching and sliced score matching for one of the tasks and other neural network models for another.\n", "weaknesses": "There is little discussion of previous work on the inverse problem studied in this work. Indeed, it is only towards the end of the work that the authors discuss a set of related work, with little or no discussion of methods applied to the inverse problem in question. If it is the case that this particular problem has not been studied and therefore there are no state-of-the-art methods, this should be clearly spelled out.\n\nAnother issue is that the central theoretical result (that of the equivalence of the correction term to the score function of the system) is not empirically verified. This should be done since the theoretical result only holds in the continuous limit, so it is not clear how this applies in the discrete case. While this may not be possible in the more complex tasks studied, it should be doable for the 1D toy problem considered in Section 3.1. Due to the simplicity of the model, it should be possible to calculate the score function and compare it to the learned correction term.\n\nAnother issue is that the tasks are very close to the actual application, and a more detailed description may be necessary for readers who are not domain experts. For example, it is not clear what is meant by “We use semi-Lagrangian advection for the velocity and MacCormack advection for the hot marker density.”.\n\nOverall, the novelty of the proposed method is not very high. This is a learned correction to a standard method of reverse time-stepping a simulation. Perhaps a more careful study of the learned correction term and its properties would help here, but as it stands, the results are incremental.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their valuable review and helpful suggestions.\n\n- **Previous work for inverse problem tasks**: We have updated the related work section of the paper to include additional references to prior work on inverse problems, particularly the heat equation. The buoyancy-driven flow task is new with no prior state-of-the-art methods, as we wanted to design a suitable problem that is high-dimensional and includes non-trivial physics simulators with difficult-to-learn dynamics (caused by placements of different obstacles). We have included additional baselines (see global response) for general inverse problem approaches to further highlight the effectiveness of our approach compared to state-of-the-art methods. As also suggested by other reviewers, we extended the related work and moved problem-specific details to the appendix. We restructured the paper by placing the related work section right after the introduction.   \n- **Empirical verification of Theorem 1**: We have included new results in our global response where we analytically compute the correct score and compare it to the learned score of our methods (1-step, multi-step). This demonstrates that the 1-step training learns the correct score very accurately, whereas the multi-step also learns an approximation to the score, but overall weighs the prior implicit in the training dataset higher, which can be seen visually in the learned score field representation.\n- **Novelty of method**: We want to highlight that a diffusion-based approach to reverse a simulation in a probabilistic way has to the best of our knowledge not been considered in prior work. By theoretically linking learned correction training with the 1-step and multi-step loss to the score matching objective and maximum likelihood training, this approach can be very efficiently and reliably used for a large number of specific scientific downstream applications that rely on inverting highly non-linear dynamics and uncertainty quantification. \n\n- **Answers to questions**: While the training setup can be applied in more general settings (where the noise is not Gaussian), in our experiments and theory section, we consider Gaussian noise. In this case this is an SDE with a step size $\\Delta t$ that depends on specific applications. Trajectories should be thought of as vectors (the size of the vector grows with the sliding window). We will expand the discussion of differences between ODE and SDE inference in the main paper. Your understanding is correct that the ODE inference generates a maximum likelihood solution whereas the SDE samples from the entire posterior. Fig. 2(a) shows the approximation of the score. We have included a comparison of the actual (analytic) score of an SDE and the learned correction in the global response pdf. Exploding trajectories in (c) occur in the GRID model trained with 1-step loss. For GRID multi-step training, there are no exploding trajectories. In line 212-215, we mean that we cannot use a single step backward because of exploding pixel values, but instead small steps + corrections are possible. In Fig. 5 SDE results are not shown because of space limitations, but they are shown in the supplementary material (Figure 8). For deterministic systems, there is often still noise in the measurement and another practical reason is that adding small noise helps performance in the case of multimodal posteriors. There was an error in Fig. 6b where the colors of ODE and SDE were swapped. The sentences in the text are correct, we apologize for the confusion. \n\n"}, {"review_id": "YbYdC23PoT", "reviewer": "Reviewer_riEs", "summary": "This paper proposes a diffusion-based unrolled strategy for learning solve ordinary differential equations (stochastic or not). After introducing the problem at stake and proposing two training strategies for solving it (namely a 1-step loss approach and a multi-step approach), the authors draw a theoretical parallel with denoising score matching and with probability flow ODEs. More precisely, the authors claim an equivalence between training an architecture with the proposed 1-step loss and training the same architecture with a score matching objective; and the equivalence between minimizing the multi-step loss and maximizing a variational lower bound. Eventually, the authors investigate the performance of the proposed approach in 4 setups, including deterministic and stochastic problems, ranging from a simple 1D experiment to a navier-stokes simulation.", "questions": "The following points list my concerns / questions / suggestions for the authors.\n\n**1. Main concerns**\n\n**1.1 Difficulties to follow the article** My main concern is that the paper is very difficult to follow. While overall well written, I tend to lose track of what the authors are aiming to do. For instance, in \"method overview\", the problem is not clearly stated. In my view, equation (6) of the main is the model that the authors aim at solving everywhere, but it comes in the middle of a section explaining score matching. In fact, the problem setting description from the appendix is rather clear while the first paragraph of 2 is not clear. There is a numerical update rule, but where does it come from? What is P? What is x? How are they related? Maybe if (6) had come in the introduction, the \"Problem formulation\" from section 2 would have been clearer.\n\n**1.2 Comparison with baselines** While the theoretical explanations and intuitions link quite clearly denoising score matching with the problem of interest, the case where a network $s_{\\theta}(x, t)$ is trained as a denoiser with a time embedding in a diffusion fashion does not seem to be present in the experiments (in ISM and SSM, $s_\\theta(x, t)$ are not trained as denoisers if I understand correctly; please correct me if I'm wrong). Moreover, these architectures do not scale to higher dimensional problems (see line 225). As such, I wonder whether these baselines provide fair comparisons. Given the embedding of $s_\\theta(x, t)$, I would suggest to use a method relying on denoising diffusion such as [HJA20]. If not applicable, maybe Diffusion Schrodinger Bridge would be a strong baseline to compare to [1]. Furthermore, given the similarities between the problem the authors want to tackle and the one from [1], I think adding a brief explanation on how the problems / approaches differ might be welcome.\n\n**1.3 PDE litterature** I have difficulties with the relations to other works. In my opinion, section 4 comes too late and should be inserted way earlier first - probably before section 2. More precisely: in \"Learned corrections for numerical errors\", a more detailed review of learning-based methods for solving PDEs / stochastic PDEs would be helpful. \n\n**1.4 Lack of references to physics informed diffusion approaches** In general I believe that a large part of the imaging inverse problems literature is not mentioned. While this is not the core topic of the paper, it remains interesting in its own right since many papers have proposed methods for incorporating a measurement operator (or P(x) in the authors' words) within diffusion models, thus making the diffusion process \"physics aware\", which is precisely what the authors want to do here. A cornerstone reference, which the authors included, is [Chu+22]. However, the authors mention that this work performs uncertainty quantification, and state \"either focus on the denoising objective common for tasks involving natural images, or the synthesis process of solutions does no directly consider the underlying physics.\": I disagree with this, see e.g. Figure 4 of the paper. If this reference does not convince the authors, here are other references where underlying physics / acquisition procedures are taken into account in a diffusion process: [2, 3, 4, 5]. Note also that an extensive literature in the inverse imaging literature has focused on a similar approach to your multistep loss, via architectures known as unfolded architectures incorporating the physics model inside the architecture [6].\n\n**1.5 Proof of Theorem 2.1** I wonder whether the proof of $\\Rightarrow$ is correct. My concern is with the particular sentence: \"let $\\theta^*$ denote a minimizer such that $\\mathcal{L}(\\theta) \\to 0$ as $\\Delta t \\to 0$. Note that at least one minimizer exists as we can choose $s_{\\theta^*}(x, t) = \\nabla_x \\operatorname{log} p_t(x)$.\" While I agree that a minimizer to this convex functional exists regardless of the nature of  $s_\\theta(x, t)$, I am not sure you can assume that the minimum value of the functional tends to 0 without any further assumption on the very nature of $s_\\theta(x, t)$: take for instance a simplistic model that is not powerful enough to approximate $\\nabla_x \\operatorname{log} p_t(x)$...\n\n**2. Additional painpoints**\n\n**2.1 Title** IMO, there is a mismatch between the title and the article: the authors do not propose a denoising score matching method, but a radically different approach that they claim to be equivalent to score matching. \n\n**2.2 Spectral loss** Why is a spectral loss necessary (line 241)? Aren't l2 (or l1, often more efficient) sufficient? Choosing a spectral loss seems unusual to me, maybe linking to some other papers using a similar loss would be welcome?\n\n**2.3 Cumberstone notations** Notations are sometimes difficult to follow, maybe clarifying them would be useful. Some examples: around line 167, $p_1$ and $p_{-1}$ clash with $p_0$ and $p_T$. Around line 63, $\\Delta t = t_j-t_k$, but then the authors use the convention $(t_m)_{0 \\leq m \\leq M}$. Why not replacing $j$ and $k$ with $m+1$ and $m$? etc...\n\n**2.4** Shouldn't the term inside the norm in eq. (2) be updated $x_{m+1}-x_m - \\Delta t \\cdots$?\n\n**2.5** line 110 of the supplementary, should (9) not be (8) instead?\n\n**2.6** Eq. (38) in supplementary: is the brownian term not missing?\n\n**2.7** Table 1 from supplementary: DilatedConv --> Dil-ResNet?\n\n**2.8** Line 426 of supplementary: 100% of what?\n\n**2.9** line 565 of supplementary: define the exact expression for $s_1$ and $s_2$.\n\n\n**References:**\n\n[1] De Bortoli, Valentin, James Thornton, Jeremy Heng, and Arnaud Doucet. \"Diffusion Schrödinger bridge with applications to score-based generative modeling.\" Advances in Neural Information Processing Systems 34 (2021): 17695-17709.\n\n[2] Chung, Hyungjin, Jeongsol Kim, Michael T. Mccann, Marc L. Klasky, and Jong Chul Ye. \"Diffusion posterior sampling for general noisy inverse problems.\" arXiv preprint arXiv:2209.14687 (2022).\n\n[3] Zhu, Yuanzhi, Kai Zhang, Jingyun Liang, Jiezhang Cao, Bihan Wen, Radu Timofte, and Luc Van Gool. \"Denoising Diffusion Models for Plug-and-Play Image Restoration.\" In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 1219-1229. 2023.\n\n[4] Kawar, Bahjat, Michael Elad, Stefano Ermon, and Jiaming Song. \"Denoising diffusion restoration models.\" Advances in Neural Information Processing Systems 35 (2022): 23593-23606.\n\n[5] Rout, Litu, Negin Raoof, Giannis Daras, Constantine Caramanis, Alexandros G. Dimakis, and Sanjay Shakkottai. \"Solving Linear Inverse Problems Provably via Posterior Sampling with Latent Diffusion Models.\" arXiv preprint arXiv:2307.00619 (2023).\n\n[6] Adler, Jonas, and Ozan Öktem. \"Learned primal-dual reconstruction.\" IEEE transactions on medical imaging 37, no. 6 (2018): 1322-1332.", "limitations": "Yes.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- The paper investigates an original idea proposing to link traditional 1-step or multi-step losses in trajectory estimation for dynamical systems with diffusion and score matching. To the best of my knowledge, this idea is new in the literature and makes a link between traditional training approaches and diffusion processes. \n- The authors perform a large number of experiments to support their claims.", "weaknesses": "- Albeit well written, the article is difficult to follow due to not always well organised sections and very large amount of information.\n- The chosen baselines do not seem very relevant, because not relying on standard losses for training denoising diffusion architectures for inverse problems.\n- There is a potential problem with one of the theoretical results (Theorem 2.1)\n- The literature review is not sufficient; while important and recent works are appropriately cited, there lacks references to physics informed diffusion approaches for inverse problems.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their valuable review and helpful suggestions.\n\n- **Structure of sections and extended literature review**: In line with feedback from other reviewers, we have moved the related work section to appear after the introduction. We also removed some of the experimental details and pushed them to the appendix. With the new space we expanded the method overview section to include more explanations, similar to the extended method overview in the appendix. We also already discuss the physical system as an SDE (equation 6) in the introduction now. We also expanded the discussion of additional learning-based methods for solving PDEs / stochastic PDEs and physics-informed diffusion approaches in the related work. We hope that this improves the structure of the paper and facilitates following the main goals of the article while only making minor changes to the content. \n\n- **Additional baselines and physics informed diffusion approaches**:  We have included additional baselines for the buoyancy-flow problem, which represents the most challenging task. In particular, we include direct numerical optimization with the BFGS method and the differentiable solver as well one of the recent physics informed diffusion approaches [2] mentioned by the reviewer. Our method compares very favorably to these baselines for challenging tasks, see the pdf of the global response. The solutions obtained by classical optimization with BFGS and the differentiable solver attain a good reconstruction MSE, although still outperformed by SMDP. However, as can be seen in the visualizations, the solutions are far away from more natural flows contained in the training dataset. In additional to that, SMDP provides a significant speedup ($\\sim$ 100x) since it does not rely on expensive gradient computations during inference. Also, our approach obtains better performance than [2], Algorithm 1 for this task. We believe that there are several reasons why our method performs significantly better than [2] in this situation. The training dataset is quite small for standard diffusion models as it comprises only 250 simulations. When training a DDPM to generate flow fields for a specific point in time, there are only 250 samples, which might not be sufficient. Additionally, the computation of the forward simulator $\\mathcal{A}$ by the differentiable solver is very slow (>60 minutes for inference of a single simulation) and gradients backpropagated through many simulation steps are not very helpful anymore for the optimization. If we run [2], Algorithm 1 on this task, the initial state is too noisy and gradients are extremely high or contain NaNs. We therefore sample the first 900 steps with the standard DDPM algorithm and switch to [2] after that. However, this could degrade the performance of the method. \n\n- For the toy problem, where we compare with ISM/SSM, we are mainly interested in methods that learn the score to a given arbitrary SDE. For the toy problem, it would be trivial to learn the mapping between the distribution of initial states at $t=0$ and the Gaussian distribution at $t=10$ with a standard diffusion model or the diffusion Schroedinger Bridge method, but that was not the main objective of this task.  \n\n- **Proof of Theorem 2.1**: We included an additional assumption for this theorem that the hypothesis space corresponding to the model architecture $s_\\theta(x,t)$ includes the correct score $\\nabla_\\mathbf{x} \\log p_t(x)$. This should fix the mentioned issue. \n\n- **Spectral loss**: Using a spectral error for scientific data is not uncommon, see e.g. Um et al. (2020). The spectral error facilitates evaluating how the prediction matches the statistical properties of the ground truth on different scales, which is not possible with the l2 error.\n\n- **Notation issues**: We thank the reviewer for the detailed comments. We have fixed all mentioned issues in the main paper and supplementary where appropriate.  \n\nWe believe these updates address the issues raised by the reviewer, and hence we kindly ask the reviewer to consider raising their score. \n"}, {"review_id": "LgIju3mFlh", "reviewer": "Reviewer_og5h", "summary": "The authors propose diffusion-based inverse problem solvers involving the temporal evolution of physics systems. The method utilize a combination of score function and an inverse physics simulator, which corresponds to reverse of drift term in diffusion models, to moves the system’s state backward in time. They demonstrate the effectiveness of their method in a wide range of inverse physics problems.", "questions": "In the evaluation of this methods, is it acceptable to not compare it with conventional method such as finite element methods?", "limitations": "It is suggested to include the future direction of this work.", "rating": 4, "confidence": 2, "soundness": 4, "presentation": 4, "contribution": 2, "strengths": "The authors propose multi-step loss to capture long-range dependence in physics system, which has potential implication for general diffusion models.\nThe experiment are conducted extensively.", "weaknesses": "The method is modification of diffusion model with adoption of inverse physic simulator in order to address nonlinear drift term of the physics system. In this regard, its novelty is limited and its applicability may be restricted to certain scenarios. \n\nMinor/errata\n\nHats are missing in equation (4)\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their valuable review and helpful suggestions.\n\n- **Conventional methods as baseline**: We evaluated an additional baseline which represents a more traditional method for solving inverse problems for the buoyancy-driven flow with obstacles experiment. This baseline is based on classical gradient-based optimization for differentiable physics and the L-BFGS algorithm, see [1]. Results can be seen in the pdf attached to the global response.\n- **Future directions**: We now included future directions and extensions of this work in our conclusion section. In particular, we are interested in enhancing the trajectories and obtaining solutions by including gradient-based optimization during inference. \n\n[1] Thuerey et al. \"Physics-based deep learning\" arXiv:2109.05237          \n"}, {"review_id": "pWSRPA0cXe", "reviewer": "Reviewer_mtJA", "summary": "The paper proposes using score matching to learn the backward process of a given forward SDE, notably coming from a physics application. The paper states that this can be used to simulate backward the distribution of the initial condition given the end state, by starting from the end state and drawing trajectories backward according to the backward SDE. This can be of interest in situations where the physics of a system is not invertible, as the authors point out is the case for most of the macroscopic state governing equations. \n\nThe authors propose two inference procedures, a SDE (sampler) and an ODE (deterministic) and test in 4 different settings, with different levels of complexity.", "questions": "* There is a practical difference between training the proposed model and the standard denoising diffusion model [2]. In [2] the losses are calculated by sampling from $p(x_t | x_0)$ which can be written as $\\mu_t(x_0) + \\epsilon$ where $\\epsilon$ is some gaussian noise. In the context of the paper this is not always possible and one need to rely on a given set of paths. How does this impact the training? Is it possible to generate the same kind of training function that does not depend of unrolling the paths? \n\n* The values in Table 1 are counter-intuitive as far as I'm concerned. Why do the methods '1-step' and 'SSM-VR' seem to achieve the best posterior metric Q for smaller datasets in both ODE and SDE?\n\n* Are the situations where using an inverse step solver $P^{-1}$ instead of $P$ motivated only by the numerical cost of running $P$ ?\n\n\n\n\n[2] Ho, J., Jain, A., & Abbeel, P. (2020). Denoising Diffusion Probabilistic Models. In Advances in Neural Information Processing Systems (pp. 6840–6851). Curran Associates, Inc..", "limitations": "NA", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "*The paper exploits an interesting analogy between the current score matching framework used for generative modelling and SDEs coming from physics and the inverse problems associated to them. The idea is that the forward process do not need to be a diffusion and therefore we can learn the backward process for a given SDE to sample a trajectory backward. This is indeed an interesting analogy that can lead to several interesting applications in the future. \n\n*The numerical examples are clear and nicely illustrated. The numerical examples show that the current framework works better than simply running the simulator backwards or simply approaching it by a neural network.", "weaknesses": "* In line 70 the paper states that it's goal is to sample from the posterior distribution $p(x_0 | x_M^*)$ where $p$ corresponds to the forward physical process. But, during the remainder of the paper, one of the proposed algorithms suggested (SMDP-ODE) can not be considered as a sampler from $p(x_0 | x_M^*)$, since it is a function of $x_M^*$. As shown by [1], what holds is that the ODE pushes forward the distribution of $p(x_M)$ to a distribution that converges (weakly) to $p(x_0)$ so it is not clear what is the actual point of starting the ODE from a given (fixed) $x_M^*$. In the conclusion the authors do touch on the point that the ODE variant is not a sampler but I feel that the presentation generates a lot of confusion for the reader during most parts of the paper. \n\n* As I understand it, only the toy problem presents a multimodal posterior. One would expect this kind of technique to be particularly useful in settings where multimodality of the posterior is present, so I would expect the authors to focus more on those cases. This is also reflected by the metrics being used (comparisons to the \"true\" $x_0$, either $RMSE$ or $LPIPS$), which would arguably make less sense in the case of multimodality and even in the case where the mode of the posterior distribution does not match the $x_0$ that produced the fixed $x_M^*$.  \n\n* The paper do not present any comparison with other inverse problems solvers, focusing on score matching based approaches only. Even though this in comprehensible to a degree, I feel there should be at least one of the non-score matching based approaches for solving inverse problems.\n\n\n[1] Yang Song, et al. \"Score-Based Generative Modeling through Stochastic Differential Equations.\" International Conference on Learning Representations. 2021.\n\nRemarks:\n* The posterior distribution is clearly defined, as being the distribution given by\n$\\int p(x_{0:M-1}, x_M^*) dx_{1:M-1}$, therefore I find it strange to say \"a\" posterior distribution in line 136.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their valuable review and helpful suggestions.\n\n- **SMDP-ODE and $p(\\mathbf{x}_0|\\mathbf{x}_t)$**: We updated section 2.3 to more clearly distinguish the different design choices for the ODE and SDE sampler. Since the inference of the ODE method is similar to the trajectories of the multi-step loss based on maximum likelihood training it can be regarded as a maximum likelihood solution.  \n- **Multimodal posterior**: As correctly stated by the reviewer, our proposed method is capable of sampling from multimodal posteriors similar to diffusion models. We have tested this extensively for the 1D toy problem. In more challenging high-dimensional problems, it is often difficult to obtain a ground truth posterior to which we can compare the posterior we obtain by sampling from the SDE if there is multimodality. Additionally, even if the ground truth posterior is known, defining a good metric to compare both high-dimensional posteriors for simulation data is non-trivial. Nonetheless, the heat diffusion experiments make a step in that direction as in our experiments the information from higher frequencies is lost and must be generated by the network. As demonstrated by the low spectral error, our network correctly synthesizes structures on smaller scales that match the ground truth data distribution. \n- **Additional baselines**: We have included two additional baselines for the buoyancy-driven flow with obstacles experiment. The first baseline is based on classical optimization with differentiable physics and the L-BFGS algorithm [1] . The second baseline is a hybrid approach that combines a standard diffusion model with gradient based optimization [2]. Both of them represent very strong baselines and highlight the effectiveness of our method in challenging high-dimensional inverse problems.\n- **Standard denoising diffusion model training and relying on a given set of paths**: In principle it is not necessary to rely on a specific set of paths. In standard diffusion models it is easy to sample points at a specific time index from the subspace of paths that connect the data distribution at t=0 to the noise distribution at t=T due to the simplicity of the underlying SDE. In our context, just sampling points that are available from the dataset would correspond to this approach and resembles the 1-step loss. However, it is easy to generate the paths for standard diffusion models during training (the points can be sampled easily) and an infinite number can be considered during training in theory whereas it is not so easy in our scenario since the paths can neither be easily generated nor expanded. In that sense, the multi-step loss increases the number of points in the training, by generating additional possible paths based on individual points in the training dataset based on the physics simulator and current score model parameter. If we consider extremely large training dataset as done for standard diffusion models in computer vision, we expect that the 1-step training and multi-step loss would attain a similar performance. Such large high-quality and diverse datasets are however rarely available for more specific scientific applications.\n- **Counter-intuitive values in Table 1**: The training and evaluation of this task is somewhat noisy. The posterior metric is very sensitive to the score field at a specific region (where the paths from -1 or 1 merge or separate). Because of this, even though the training converges and losses decrease, it is possible that the metric Q is close to 0, because there is an imbalance between the predicted classes (in these cases the evaluation gives either paths for -1 or 1 and has a problem with the multimodality in this example). The multi-step loss is much less noisy than the other training methods here.\n- **Inverse step solver $\\tilde{\\mathcal{P}}^{-1}$**: For all experiments where the physics simulator is not learned (experiments 1-3) we use the existing implementation of the forward solver to derive the reverse solver by changing $\\Delta t$ to $-\\Delta t$. This is not motivated by the numerical cost. When pretraining the inverse solver in experiment 4, we specifically use a solver pretrained for the inverse problem.\n\n[1] Thuerey et al. \"Physics-based deep learning\" arXiv:2109.05237\n\n[2] Chung, Hyungjin, Jeongsol Kim, Michael T. Mccann, Marc L. Klasky, and Jong Chul Ye. \"Diffusion posterior sampling for general noisy inverse problems.\" arXiv:2209.14687                                                                                                                                                                                         \n\n"}, {"review_id": "dHgNUpRTct", "reviewer": "Reviewer_3QqF", "summary": "The authors leverage the framework of score matching, which has become popular for training diffusion-based models for generative tasks, to reverse physical processes defined by forward stochastic differential equations (SDEs). Given a system state at time t=T, they propose to solve for the initial conditions at t=0 by iteratively applying a reverse-time diffusion process defined by a reverse physics simulator, a diffusion term, and the score of the data distribution. In the main contribution of the paper, the authors introduce both a 1-step and multi-step loss for training a network to learn the diffusion and score terms for solving the reverse SDE. They prove the equivalence of their proposed training objectives to vanilla denoising score matching (and a related variational objective in the multi-step case). Extensive experiments demonstrate the efficacy of their method compared to baselines, and ablation studies show the utility of the proposed multi-step objective.               ", "questions": "- What is your reasoning for training your model with the 1-step loss between $x_m$ and $x_{m-1}$ instead of $x_m$ and $x_0$ (as in typical diffusion-based generative models)? As I understand it, the forward and reverse physical processes are not time-dependent, and arbitrarily large time steps can be taken by simply increasing $\\Delta t$. Intuitively, this method makes more sense to me than your proposed approach, and would also remove the need for multi-step training. \n- How did you deal with the memory expense during multi-step training? Did you find that you needed to use smaller data and network sizes to fit the gradients in memory, or were there tricks and optimizations you used to reduce memory use?\n- One concern is that your model will overfit to the physical parameters at training and perform poorly if there is a test-time distribution shift. How robust is your proposed approach to test-time shifts in the SDE parameters (namely, the coefficients of the simulator and diffusion terms)? I understand that there is limited time for responses, but a small experiment would be appreciated and may convince me to raise my score.  ", "limitations": "The authors mention limitations in their conclusion.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "### Originality\n\nTo my knowledge, this is the first work that considers the drift term in the typical forward SDE used in diffusion-based generative models to be a realistic physical process. Other works [1,2] have considered different diffusion processes besides additive Gaussian noise, but even these use artificial processes such as synthetic blur and pixel masking. I appreciate that the authors of the current paper identify the connection between physically-defined differential equations and those used in training score-based models, and propose a novel method for solving physics problems.\n\n### Quality\n\n- The authors perform extensive experiments grounded in real-world physical processes and show the effects of varying numerous design choices in each experiment\n- The proposed approach demonstrates superior performance compared to baselines across various settings\n- Ablation studies show the utility of the multi-step objective over the 1-step variant\n- The figures are well-made, particularly Figure 1 which clearly lays out the key ideas of the proposed approach  \n\n### Clarity\n\n- The introduction does a good job of laying out the current state of diffusion-based generative models, the connection to physics processes, and the contribution of the current work\n- Clear descriptions of the problem setup, model training and inference, and hyperparameters are given for each experiment\n\n### Significance\n\nPhysics-based inverse problems arise in many fields such as astronomy, geophysics, and wireless communication, so finding better solutions is a highly significant problem with broad interest. Furthermore, quantifying the uncertainty of these solutions is crucial to the downstream decision-making process. The authors of the current paper provide a principled approach to both solve and provide uncertainty estimates (using multiple samples from the SDE) for inverse problems in physics.                    \n\n### References:\n\n[1] G. Daras, M. Delbracio, H. Talebi, A. Dimakis, and P. Milanfar, “Soft Diffusion: Score Matching with General Corruptions,” Transactions on Machine Learning Research, 2023, [Online].\n\n[2] A. Bansal et al., Cold Diffusion: Inverting Arbitrary Image Transforms Without Noise. 2022.", "weaknesses": "- I believe that the loss in Eq (2) is incorrect. The quantity within the squared L2 norm should be $x_m$ minus the quantity on the right-hand side of Eq (1). In its current form, Eq (2) does not match Eq (3) when Eq (3) is expanded with window S=2.  \n- While the thorough experimental details are appreciated, I believe some of that can be pushed to the appendix. More space should be dedicated to expanding on motivation and design choices. As it currently stands, section 2 reads as a constant flow of information with insufficient context for the proposed objectives.\n- I believe that the organization of the sections could be improved by introducing the typical score matching SDE subject matter first, identifying the differences between that formulation and the physics-based inverse problem formulation, then motivating your proposed approach within this context.\n- An obvious concern to me regarding the multi-step training is the huge memory expense arising from the recursive calls to $s_{\\theta}$. However, the authors did not address this point in the main paper.\n- The authors state that, for decreasing time step $\\Delta t$, the reverse physics simulator is equivalent to the negative of the forward simulator (line 113). However, it is not clarified whether this is a simplifying assumption, true in general, or true in the specific case of the reverse-time ODE. The authors expand on the specific choices of the reverse simulators in the experiments section, but the relationship between the forward and reverse simulators remains unclear to me. \n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their valuable review and helpful suggestions.\n\n- We restructured our paper based on the feedback we received. In particular, we moved the related work section forward so that it follows the introduction and moved parts of the experimental details to the appendix to improve overall readability.\n\n- **Memory expense**:  Possible solutions to save memory are gradient cutting and gradient checkpointing. We have tested the effects of gradient cutting for the heat equation task. In this experiment, we unrolled the entire simulation with the multi-step loss, but stop backpropagating the gradient after n steps. n = 32 corresponds the entire simulation trajectory. Our results show that there are no more performance benefits when backpropagating the gradient for more than 8 simulation steps for the ODE inference. \n| Gradient cutting after | Avg. reconstruction MSE  |\n| ----| ----------- |\n| 1   |  4.6e-5   |\n| 2   | 1.4e-5    |\n| 4   | 1.0e-5    |\n| 8   | 0.8e-5    |\n| 16 | 0.8e-5  |\n| 32 | 0.8e-5  |\n\n- For the evaluation of the tasks in the paper, we trained with no specific optimizations of the memory consumption. The most expensive task in terms of memory is the buoyancy-driven flow with obstacles experiment for which we used a single A100 GPU with 40GB of memory. Our neural networks are very small compared to architectures for diffusion models and generative modeling and have a size that’s comparable to typical learned correction approaches, cf. e.g., [Um et al. 2020, Kochkov et al. 2021].\n\n-  **Reverse-physics simulator**: In general there is no reverse simulator for larger time scales, as information can be lost and the initial state might not be possible to reconstruct. If we solve the underlying PDE iteratively this can be described by the update rule $x_{t+1} = x_{t} + \\Delta t$ PDEupdate$(x_t)$. Then, we may approximate $x_{t} \\approx x_{t+1} - \\Delta t$ PDEupdate$(x_{t+1})$. So locally we can identify the reverse simulator with the negative of the forward simulator. In our experiments with non-learned simulators, we use the implementation of the forward simulator to obtain the reverse simulator by using a negative step size $\\Delta t$.  We included additional comments about this in the main paper as well as expanding upon the relationship in the appendix.\n\n-  **1-step loss between $\\mathbf{x}_0$ and $\\mathbf{x}_m$**: In theory it would be possible to consider the 1-step loss between $\\mathbf{x}_0$ and $\\mathbf{x}_m$ and compute the trajectory in a single step. However, in practice there are numerical issues that need to be considered. In standard diffusion models, the score is easy to compute given sample and noise, but in our case, we still require an ODE solver. Our method is described from the viewpoint of Euler steps. However, other methods for time integration could be considered, which we leave to future work. Importantly, there are considerations regarding the implementation of the reverse-physics simulator. We found that many small steps + correction for each step is numerically significantly more stable than a single big step with one correction only. \n\n- **Test-time distribution shifts**: We have tested the effects of test-time distribution shifts for the heat equation experiment. Here we train the score network for a specific combination of diffusivity and noise and vary both parameters for testing (always updating both the simulator and test ground truth), see the pdf in the global response. Overall, for small changes of the parameters, there seems to be very little overfitting. Changes in the reconstruction MSE and spectral error can mainly be explained by making the task itself easier/harder to which our network generalizes nicely (e.g. less noise or higher diffusivity -> smaller reconstruction error).\n"}]

Author rebuttal:
We thank all reviewers for their helpful suggestions and comments. Based on the feedback, there are several updates that are of interest to all reviewers:

- **Additional baselines and physics informed diffusion approaches**:  We have included additional baselines for the buoyancy-driven flow problem, which represents the most challenging task. In particular, we include direct numerical optimization with the BFGS method and the differentiable solver as well as a recent physics informed diffusion approach [1] mentioned by one of the reviewers. Our method compares very favorably to these baselines for challenging tasks while inference is significantly faster at the same time, see the attached pdf.

- **Test-time distribution shifts**: We have tested the effects of test-time distribution shifts for the heat equation experiment. Here we train the score network for a specific combination of diffusivity and noise and vary both parameters for testing (always updating both the simulator and test ground truth), see the attached pdf. Overall, for small changes of the parameters (<30%), there seems to be very little overfitting. Changes in the reconstruction MSE and spectral error can mainly be explained by variations of the parameters making the task itself easier/harder to which our network generalizes nicely (e.g. less noise or higher diffusivity -> smaller reconstruction error).

- **Empirical verification of Theorem 1**: We have included new results in the attached pdf where we analytically compute the correct score and compare it to the learned score of our methods (1-step, multi-step). This demonstrates that the 1-step training learns the correct score very accurately, whereas the multi-step also learns an approximation to the score, but overall pulls trajectories during inference closer to the training data set distribution, which can be seen visually in the learned score field representation. 

- **Structure of sections and extended literature review**: In line with feedback from other reviewers, we have moved the related work section to appear after the introduction. We also removed some of the experimental details and pushed them to the appendix. With the new space we expanded the method overview section to include more explanations, similar to the extended method overview in the appendix. We also already discuss the physical system as an SDE (equation 6) in the introduction now. Moreover, we expanded the discussion of additional learning-based methods for solving PDEs / stochastic PDEs and physics-informed diffusion approaches in the related work. We believe this improves the structure of the paper and facilitates following the main goals of the article while only making minor changes to the content. 

[1] Chung, Hyungjin, Jeongsol Kim, Michael T. Mccann, Marc L. Klasky, and Jong Chul Ye. "Diffusion posterior sampling for general noisy inverse problems." arXiv preprint arXiv:2209.14687 


Meta-review: The paper has received mixed reviews (6-6-6-4-4-7). The reviewers have raised a number of concerns, many of which have been addressed by the authors during the author-reviewer discussion period. In particular, reviewer mtJA's concerns (borderline reject) about the presentation and the consolidation of the empirical evaluation have been, in my opinion, satisfactorily addressed. Reviewer og5h's concerns (borderline reject) regarding the novelty have been cleared out during the discussion, although the evaluation has remained the same due to a lack of familiarity with the topic. Overall, I believe the reasons to accept the paper outweigh the reasons to reject it. The authors are requested to address any remaining concerns in the final version of the paper.
Decision: Accept (poster)