Title: Solving Inverse Physics Problems with Score Matching
Venue: NeurIPS 2023
Accepted: yes
Keywords: inverse problems

diffusion models

learned corrections

score matching
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

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 .

1 Introduction

Many physical systems are time-reversible on a microscopic scale. For example, a continuous material can be represented by a 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 a macroscopic perspective, we only know average quantities within specific regions [Far93], which constitutes a loss of information, and as a consequence, time is no longer reversible. In the following, we target inverse problems to reconstruct the distribution of initial macroscopic states for a given end state. This problem is genuinely tough [ZDG96; Góm+18; 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 ∈ R D are gradually corrupted into Gaussian white noise via a stochastic differential equation (SDE) d x = f ( x , t ) dt + g ( t ) dW , where the deterministic component of the SDE f : R D × R ≥ 0 → R D is called drift and the coefficient of the D -dimensional Brownian motion W denoted by g : R ≥ 0 → R ≥ 0 is called diffusion . If the score ∇ x log p t ( x ) of the data distribution p t ( x ) of corrupted samples at time t 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 a neural network s θ , which can then be used as a plug-in estimate for the reverse-time SDE.

However, in our physics-based approach, we consider an SDE that describes the physics system as d x = P ( x ) dt + g ( t ) dW , where P : R D → R D is a physics simulator that replaces the drift term of diffusion models. Instead of transforming the data distribution to noise, we transform a simulation state at t = 0 to a simulation state at t = T with Gaussian noise as a perturbation. Based on a given

1

end state of the system at t = T , we predict a previous state by taking a small time step backward in time and repeating this multiple times. Similar to the reverse-time SDE of diffusion models, the prediction of the previous state depends on an approximate inverse of the physics simulator, a learned update s θ , and a small Gaussian perturbation.

The training of s θ is similar to learned correction approaches for numerical simulations [Um+20; Koc+21; LCT22]: The network s θ learns corrections to simulation states that evolve over time according to a 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 a single previous state, or be extended to rollouts for multiple steps. The latter requires the differentiability of the inverse physics step [Thu+21].

Importantly, we show that under mild conditions, learning s θ is equivalent to matching the score ∇ x log p t ( x ) of the training data set distribution at time t . 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+21b; ZC21]. In practice, the learned component s θ 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 1 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 a 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 a reverse physics simulation step into diffusion models to develop a 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 a 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+21], uncertainty quantification [CSY22; Chu+22; Son+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 or the synthesis process of solutions does not directly consider the underlying physics. Models based on Langevin dynamics [Vin11; SE19] or discrete Markov chains [Soh+15; HJA20] can be unified in a time-continuous 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 a physics system with stochastic perturbations.

Methods for solving inverse problems for (stochastic) PDEs Differentiable solvers for physics dynamics can be used to optimize solutions of inverse problems with gradient-based methods by backpropagating gradients through the solver steps [Thu+21]. 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+21], DeepONets [Lu+21], or generalizations thereof that include stochastic forcing for stochastic PDEs, e.g., neural stochastic PDEs [SLG22]. Recently, there have been several approaches that leverage the learned scores from diffusion models as data-driven

OX

Timet

X1

PT

X2

ion

X3

po

X4

Training

Reference

X2

X3

X4

1on

Prediction

2

PT

X

80（x,t)

approximates

score

V

7xlogpt

X

(a) Training and inference overview.

(b) Sliding window and optimizing s θ .

regularizations for linear inverse problems [Ram+20; Son+22; KVE21; CSY22; Chu+22] and general noisy inverse problems [Chu+23]. Our method can be applied to 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+18; Pfa+20; Li+21]. By integrating a neural network into differential equation solvers, it is possible to learn to reduce numerical errors [Um+20; Koc+21; BWW22] or guide the simulation towards a desired target state [HTK20; Li+22]. The optimization of s θ with the 1-step and multi-step loss we 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 on neural ODEs [Che+18a; 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 physics system is described by Kolmogorov's forward equation [Øks03], and there is a corresponding probability flow ODE [MRO20; Son+21b]. When the score is represented by s θ , this constitutes a CNF and can typically be trained with standard methods [Che+18b] and maximum likelihood training [Son+21a]. Huang et al. [HLC21] show that minimizing the score-matching loss is equivalent to maximizing a lower bound of the likelihood obtained by sampling from the reverse-time SDE. A recent variant combines score matching with CNFs [ZC21] 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.

3 Method Overview

Problem formulation Let (Ω , F , P ) be a probability space and W ( t ) = ( W 1 ( t ) , ..., W D ( t )) T be a D -dimensional Brownian motion. Moreover, let x 0 be a F 0 -measurable R D -valued random variable that is distributed as p 0 and represents the initial simulation state. We consider the time evolution of

the physical system for 0 ≤ t ≤ T modeled by the stochastic differential equation (SDE)

with initial value x 0 and Borel measurable drift P : R D → R D and diffusion g : [0 , T ] → R ≥ 0 . This SDE transforms the marginal distribution p 0 of initial states at time 0 to the marginal distribution p T of end states at time T . We include additional assumptions in appendix A.

Moreover, we assume that we have sampled N trajectories of length M from the above SDE with a fixed time discretization 0 ≤ t 0 < t 1 < ... < t M ≤ T for the interval [0 , T ] and collected them in a training data set { ( x ( n ) t m ) i M =0 } N n =0 . For simplicity, we assume that all time steps are equally spaced, i.e., t m +1 -t m := ∆ t . Moreover, in the following we use the notation x ℓ : m for 0 ≤ ℓ < m ≤ M to refer to the trajectory ( x t ℓ , x t ℓ +1 , ..., x t m ) . We include additional assumptions in appendix A.

Our goal is to infer an initial state x 0 given a simulation end state x M , i.e., we want to sample from the distribution p 0 ( · | x M ) , or obtain a maximum likelihood solution.

3.1 Learned Corrections for Reverse Simulation

In the following, we furthermore assume that we have access to a reverse physics simulator ˜ P -1 : R D → R D , which moves the simulation state backward in time and is an approximate inverse of the forward simulator P [HKT22]. In our experiments, we either obtain the reverse physics simulator from the forward simulator by using a negative step size ∆ t or by learning a surrogate model from the training data. We train a neural network s θ ( x , t ) parameterized by θ such that

In this equation, the term s θ ( x m +1 , t m +1 ) corrects approximation errors and resolves uncertainties from the Gaussian perturbation g ( t ) 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 For a pair of adjacent samples ( x m , x m +1 ) on a data trajectory, the 1-step loss for optimizing s θ is the L 2 distance between x m and the prediction via (2). For the entire training data set, the loss becomes

Computing the expectation can be thought of as moving a window of size two 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 s θ . However, in practice, approaches that consider a loss based on predicting longer parts of the trajectories are more successful for training learned corrections [Bar+19; Um+20; Koc+21]. For that purpose, we define a hyperparameter S , called sliding window size, and write x i : i + S ∈ R S × D to denote the trajectory starting at x i that is comprised of x i and the following S -1 states. Then, we define the multi-step loss as

where ˆ x i : i + S -1 is the predicted trajectory that is defined recursively by

3.2 Learning the Score

Denoising score matching Given a distribution of states p t for 0 < t < T , we follow [Son+21b] and consider the score matching objective

i.e., the network s θ is trained to approximate the score ∇ x log p t ( x ) . In denoising score matching [Vin11; Soh+15; HJA20], the distributions p t are implicitly defined by a noising process that is given by the forward SDE d x = f ( x , t ) dt + g ( t ) dW , where W is the standard Brownian motion. The function f : R D × R ≥ 0 → R D is called drift , and g : R ≥ 0 → R ≥ 0 is called diffusion . The process transforms the training data distribution p 0 to a noise distribution that is approximately Gaussian p T . For affine functions f and g , the transition probabilities are available analytically, which allows for efficient training of s θ . It can be shown that under mild conditions, for the forward SDE, there is a corresponding reverse-time SDE d x = [ f ( x , t ) -g ( t ) 2 ∇ x log p t ( x )] dt + g ( t ) d ˜ W [And82]. In particular, this means that given a marginal distribution of states p T , which is approximately Gaussian, we can sample from p T and simulate paths of the reverse-time SDE to obtain samples from the data distribution p 0 .

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 probabilities from p T to p 0 as the reverse-time SDE. For the physics-based SDE (1), it is given by

For ∆ t → 0 , we can rewrite the update rule (2) of the training as

Therefore, we can identify ˜ P -1 ( x ) with -P ( x ) and s θ ( x , t ) with 1 2 g ( t ) ∇ x log p t ( x ) . We show that for the 1-step training and sufficiently small ∆ t , we minimize the score matching objective (6).

Theorem 3.1. Consider a data set with trajectories sampled from SDE (1) and let ˜ P -1 ( x ) = -P ( x ) .

Then the 1-step loss (3) is equivalent to minimizing the score matching objective (6) as ∆ t → 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 a probabilistic sense. For denoising score matching, it is possible to train s θ via maximum likelihood training [Son+21a], which minimizes the KL-divergence between p 0 and the distribution obtained by sampling x T from p T and simulating the probability flow ODE (8) from t = T to t = 0 . We derive a similar result for the multi-step loss.

Theorem 3.2. Consider a data set with trajectories sampled from SDE (1) and let ˜ P -1 ( x ) = -P ( x ) . Then the multi-step loss (4) maximizes a variational lower bound for maximum likelihood training of the probability flow ODE (7) as ∆ t → 0 .

Proof. See appendix A.2

To conclude, we have formulated a probabilistic multi-step training to solve inverse physics problems and provided a theoretical basis to solve these problems with score matching. Next, we outline additional details for implementing SMDP.

3.3 Training and Inference

We start training s θ with the multi-step loss and window size S = 2 , which is equivalent to the 1-step loss. Then, we gradually increase the window size S until a maximum S max . For S > 2 , the unrolling of the predicted trajectory includes interactions between s θ and the reverse physics simulator ˜ P -1 . For inference, we consider the neural SDE

which we solve via the Euler-Maruyama method. For C = 2 , we obtain the system's reverse-time SDE and sampling from this SDE yields the posterior distribution. Setting C = 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 be used to obtain many samples and to explore the posterior distribution, the ODE variant constitutes a unique and deterministic solution based on maximum likelihood training.

(a) Training

(b) Inference

(c) 1-step vs. multi-step

GRID,1-step

75

0

GRID

V

multi-step

Exploding

trajectories

MLP,1-step

T

Q

Time t

1

75

Score

Vxlogpt(x)

0

approximated by s

Sθ(x,t)

-1

75

Time

t

4 Experiments

We show the capabilities of the proposed algorithm with a range of experiments. The first experiment in section 4.1 uses a simple 1D process to compare our method to existing score matching baselines. The underlying model has a 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 a 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 a scenario without stochastic perturbations in the form of a buoyancy-driven Navier-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 -1 for isotropic turbulence flows and evaluate how well SMDP works with a learned reverse physics simulator.

4.1 1D Toy SDE

As a first experiment with a known posterior distribution we consider a simple quadratic SDE of the form: dx = -[ λ 1 · sign( x ) x 2 ] dt + λ 2 dW , with λ 1 = 7 and λ 2 = 0 . 03 . Throughout this experiment, p 0 is a categorical distribution, where we draw either 1 or -1 with the same probability. The reverse-time SDE that transforms the distribution p T of values at T = 10 to p 0 is given by

In figure 2a, we show paths from this SDE simulated with the Euler-Maruyama method. The trajectories approach 0 as t increases. Given the trajectory value at t = 10 , it is no longer possible to infer the origin of the trajectory at t = 0 .

This experiment allows us to use an analytic reverse simulator: ˜ P -1 ( x ) = λ 1 · sign ( x ) x 2 . This is a challenging problem because the reverse physics step increases quadratically with x , 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 ˆ x 0: T match the posterior distribution. When drawing ˆ x T randomly from [ -0 . 1 , 0 . 1] , we should obtain trajectories with ˆ x 0 being either -1 or 1 with the same likelihood. We assign the label -1 or 1 if the relative distance of an endpoint is < 10 %and denote the percentage in each class by ρ -1 and ρ 1 . As some trajectories miss the target, typically ρ -1 + ρ 1 < 1 . Hence, we define the posterior metric Q as twice the minimum of ρ -1 and ρ 1 , i.e., Q := 2 · min( ρ -1 , ρ 1 ) so that values closer to one indicate a better match with the correct posterior distribution.

Training The training data set consists of 2500 simulated trajectories from 0 to T and ∆ t = 0 . 02 . Therefore each training trajectory has a length of M = 500 . For the network s θ ( x, t ) , we consider a multilayer perceptron (MLP) and, as a special case, a grid-based discretization (GRID). The latter is 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 , T ] × [ -1 . 25 , 1 . 25] to obtain a rectangular grid with 500 × 250 cells and linearly interpolate the solution. The cell centers are initialized with 0 . We evaluate s θ trained via the 1-step and multi-step losses with S max = 10 . Details of hyperparameters and model architectures are given in appendix C.

Better extrapolation and robustness from multi-step loss See figure 2c for an overview of the differences between the learned score from MLP and GRID and the effects of the multi-step loss. For the 1-step training with MLP, we observe a clear and smooth score field with two tubes that merge to one at x = 0 as t increases. As a result, the trajectories of the probability flow ODE and reverse-time 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 can quickly explode. In contrast, the multi-step loss leads to a consistent signal around the center line at x = 0 , effectively preventing exploding trajectories.

Evaluation and comparison with baselines As a baseline for learning the scores, we consider implicit score matching [Hyv05, ISM]. Additionally, we consider sliced score matching with variance reduction [Son+19, SSM-VR] as a variant of ISM. We train all methods with the same network architecture using three different data set sizes. As can be seen in table 1, the 1-step loss, which is conceptually similar to denois-

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 significant improvements at the cost of a slightly increased training time. Our proposed multi-step training performs best or is on par with the baselines for all data set sizes and inference types. Because the posterior metric Q is very sensitive to the score where the paths from both starting points intersect, evaluations are slightly noisy.

Comparison with analytic scores We perform further experiments to empirically verify Theorem 3.1 by comparing the learned scores of our method with analytic scores in appendix C.

(a) Inference

Reconstruction MSE

10

4

10

3

10

2

Spectral error

10 0

ResNet

FNO

BNN

s

only ODE

s

only SDE

SMDP ODE

SMDP SDE

(b) Evaluation

4.2 Stochastic Heat Equation

The heat equation ∂u ∂t = α ∆ u plays a 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 to α = 1 and sample initial conditions at t = 0 from Gaussian random fields with n = 4 at resolution 32 × 32 . We simulate the heat diffusion with noise from t = 0 until t = 0 . 2 using the Euler-Maruyama method and a spectral solver P h with a fixed step size ∆ t = 6 . 25 × 10 -3 and g ≡ 0 . 1 . Given a simulation end state x T , we want to recover a possible initial state x 0 . In this experiment, the forward solver cannot be used to infer x 0 directly in a 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 P h ( x ) , i.e. ˜ P -1 ( x ) ≈ -P h ( x ) .

Training and Baselines Our training data set consists of 2500 initial conditions with their corresponding trajectories and end states at t = 0 . 2 . We consider a small ResNet -like architecture based on an encoder and decoder part as representation for the score function s θ ( x , t ) . The spectral solver is implemented via differentiable programming in JAX [SC20], see appendix D. As baseline methods, we consider a supervised training of the same ResNet -like architecture as s θ ( x , t ) , a Bayesian neural network (BNN) as well as a Fourier neural operator (FNO) network [Li+21]. We adopt an L 2 loss for all these methods, i.e., the training data consists of pairs of initial state x 0 and end state x T .

Additionally, we consider a variant of our proposed method for which we remove the reverse physics step ˜ P -1 such that the inversion of the dynamics has to be learned entirely by s θ , 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 to train properly, and we did not obtain stable trajectories during inference.

Reconstruction accuracy vs. fitting the data manifold We evaluate our method and the baselines by considering the reconstruction MSE on a test set of 500 initial conditions and end states. For the reconstruction MSE, we simulate the prediction of the network forward in time with the solver P h to obtain a corresponding end state, which we compare to the ground truth via the L 2 distance. This metric has the disadvantage that it does not measure how well the prediction matches the training data manifold. I.e., for this case, whether the prediction resembles the properties of the Gaussian random field. For that reason, we additionally compare

2

4

8

16

32

(a) ODE trajectories

10

-

5

4

×

10

-

6

6

×

10

-

6

2

×

10

-

5

Reconstruction MSE

2

4

8

16

32

(b) SDE trajectories

10

-

5

10

-

3

10

-

1

the power spectral density of the states as the spectral loss . An evaluation and visualization of the reconstructions are given in figure 3, 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 a 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 an ablation study on the maximum window size S max in figure 4 for the reconstruction MSE. For both ODE and SDE inference, increasing S max yields significant improvements at the cost of slightly increased training resources. This also highlights the importance of using a multi-step loss instead of the 1-step loss ( S max = 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 a more challenging problem. For this purpose, we consider deterministic simulations of buoyancy-driven flow within a fixed domain Ω ⊂ [0 , 1] × [0 , 1] and randomly placed obstacles. Each simulation runs from time t = 0 . 0 to t = 0 . 65 with a step size of ∆ t = 0 . 01 . SMDP is trained with the objective of reconstructing a plausible initial state given an end state of the marker density and velocity fields at time t = 0 . 65 , as shown in figure 5a and figure 5b. We place spheres and boxes with varying sizes at different positions within the simulation domain that do not overlap with the marker inflow. For each simulation, we place one to two objects of each category.

Score matching for deterministic systems During training, we add Gaussian noise to each simulation state x t with σ t = √ ∆ t . In this experiment, no stochastic forcing is used to create the data set, i.e., g ≡ 0 . By adding noise to the simulation states, the 1-step loss still minimizes a 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, a small perturbation, and the correction by s θ ( x , t ) , which projects the perturbed simulation state back to the distribution p t . We find that for the SDE trajectories, C = 2 slightly overshoots, and C = 1 gives an improved performance. In this setting, the " s θ 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 s θ ( x , t ) uses dilated convolutions [Sta+21], see appendix E for details. The reverse physics step ˜ P -1 is implemented directly in the solver by using a negative step size -∆ t for time integration. For training, we consider the multi-step formulation with S max = 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, DPS] with a pretrained diffusion model on simulation states at t = 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 ( s θ 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+21]. The training data set consists of vorticity fields from 1000 simulation trajectories from t = 0 until T = 10 with ∆ t = 1 , a spatial resolution of 64 × 64 and

(b) Velocity field (x)

(c) Evaluation

t=0.65

Reconstruction

t=0.65

Reconstruction

Ground truth

t=0.35

t=0.35

Reference

SMDPODE

Reference

SMDPODE

1

t=0.20

t=0.20

1

sg only

only

sg only

only

(a) Reconstructed trajectories at

t

= 9

.

MSE

14

15

16

17

Spectral error

0.1

0.2

0.3

0.4

0.5

0.6

Learned reverse

step

1

Simultaneous

training of

1

and

s

Pretrained

1

SMDP SDE

Pretrained

1

SMDP ODE

(b) Evaluation

viscosity fixed at ν = 10 -5 . As before, our objective is to predict a trajectory ˆ x 0: M that reconstructs the true trajectory given an end state x M . In this experiment, we pretrain a surrogate for the reverse physics step ˜ P -1 by employing the FNO architecture from [Li+21] trained on the reverse simulation. For pretraining ˜ P -1 we use our proposed training setup with the multi-step loss and S max = 10 but freeze the score to s θ ( x , t ) ≡ 0 . Then, we train the time-dependent score s θ ( x , t ) while freezing the reverse physics step. This approach guarantees that any time-independent physics are captured by ˜ P -1 and s θ ( x , t ) can focus on learning small improvements to ˜ P -1 as well as respond to possibly time-dependent data biases. We give additional training details in appendix F.

Evaluation and training variants For evaluation, we consider the MSE and spectral error of the reconstructed initial state ˆ x 0 compared to the reference x 0 . As baselines, during inference, we employ only the learned surrogate model ˜ P -1 without s θ . In addition to that, we evaluate a variant for which we train both the surrogate model and s θ ( x , t ) at the same time. As the two components resemble the drift and score of the reverse-time SDE, this approach is similar to 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 two components.

5 Discussion and Conclusions

We presented a 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 a 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, it requires simulating the system backward in time step by step, which can be costly and alleviated by reduced order methods. Additionally, we assume that ∆ t is sufficiently small and the reverse physics simulator is accurate enough. Determining a 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 SpaTe: 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).

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. P. Brenner. 'Learning data-driven discretizations 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 mechanics'. 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 differential equations'. In: Advances in neural information processing systems 31 (2018).

[Che+18b] T. Q. Chen, Y. Rubanova, J. Bettencourt, and D. Duvenaud. 'Neural Ordinary Differential Equations'. In: Advances in Neural Information Processing Systems . 2018, pp. 6572-6583. URL: https://proceedings.neurips.cc/paper/2018/hash/ 69386f6bb1dfed68692a24c8686939b9-Abstract.html .

[Che+21] N. Chen, Y. Zhang, H. Zen, R. J. Weiss, et al. 'WaveGrad: Estimating Gradients for Waveform Generation'. In: 9th International Conference on Learning Representations . OpenReview.net, 2021. URL: https://openreview.net/forum?id=NsMLjcFaO8O .

[Chu+22] H. Chung, B. Sim, D. Ryu, and J. C. Ye. 'Improving Diffusion Models for Inverse Problems using Manifold Constraints'. In: NeurIPS . 2022. URL: http://papers.nips. cc / paper / 2022 / hash / a48e5877c7bf86a513950ab23b360498 - Abstract Conference.html .

[Chu+23] H. Chung, J. Kim, M. T. McCann, M. L. Klasky, and J. C. Ye. 'Diffusion Posterior Sampling for General Noisy Inverse Problems'. In: The Eleventh International Conference on Learning Representations . OpenReview.net, 2023. URL: https://openreview. net/pdf?id=OnD9zGAGT0k .

[CSY22] H. Chung, B. Sim, and J. C. Ye. 'Come-Closer-Diffuse-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-12412. DOI: 10.1109/CVPR52688. 2022.01209 . URL: https://doi.org/10.1109/CVPR52688.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), P08023.

[DN21] P. Dhariwal and A. Q. Nichol. 'Diffusion Models Beat GANs on Image Synthesis'. In: Advances in Neural Information Processing Systems . 2021, pp. 87808794. URL: https : / / proceedings . neurips . cc / paper / 2021 / hash / 49ad23d1ec9fa4bd8d77d02681df5cfa-Abstract.html .

[Far93] S. J. Farlow. Partial differential equations for scientists and engineers . Courier Corporation, 1993.

[Góm+18] R. Gómez-Bombarelli, J. N. Wei, D. Duvenaud, J. M. Hernández-Lobato, et al. 'Automatic chemical design using a 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] J. Ho, A. Jain, and P. Abbeel. 'Denoising Diffusion Probabilistic Models'. In: Advances in Neural Information Processing Systems . 2020. URL: https://proceedings. neurips . cc / paper / 2020 / hash / 4c5bcfec8584af0d967f1ab10179ca4b 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.

[HLC21] 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 Processing Systems . 2021, pp. 22863-22876. URL: https://proceedings.neurips.cc/ paper/2021/hash/c11abfd29e4d9b4d4b566b01114d8486-Abstract.html .

[Ho+22] J. Ho, T. Salimans, A. A. Gritsenko, W. Chan, et al. 'Video Diffusion Models'. In: NeurIPS . 2022. URL: http://papers.nips.cc/paper%5C_files/paper/2022/ hash/39235c56aef13fb05a6adc95eb9d8d66-Abstract-Conference.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.

[Höp+22] T. Höppe, A. Mehrjou, S. Bauer, D. Nielsen, and A. Dittadi. 'Diffusion Models for Video Prediction and Infilling'. In: CoRR abs/2206.07696 (2022). DOI: 10.48550/ arXiv.2206.07696 . URL: https://doi.org/10.48550/arXiv.2206.07696 .

[HTK20] P. Holl, N. Thuerey, and V. Koltun. 'Learning to Control PDEs with Differentiable Physics'. In: 8th International Conference on Learning Representations . OpenReview.net, 2020. URL: https://openreview.net/forum?id=HyeSin4FPB .

[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://jmlr.org/papers/ v6/hyvarinen05a.html .

[KCT23] G. Kohl, L.-W. Chen, and N. Thuerey. 'Turbulent Flow Simulation using Autoregressive Conditional Diffusion Models'. In: arXiv:2309.01745 (2023).

[Kid+21] P. Kidger, J. Foster, X. Li, and T. J. Lyons. 'Neural SDEs as Infinite-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: http://proceedings.mlr.press/ v139/kidger21b.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/abs/2102.01010 .

[KPB20] I. 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 (2020), pp. 3964-3979.

[KVE21] B. Kawar, G. Vaksman, and M. Elad. 'SNIPS: Solving Noisy Inverse Problems Stochastically'. In: Advances in Neural Information Processing Systems . 2021, pp. 2175721769. URL: https : / / proceedings . neurips . cc / paper / 2021 / hash / b5c01503041b70d41d80e3dbe31bbd8c-Abstract.html .

[LCT22] B. List, L.-W. Chen, and N. Thuerey. 'Learned turbulence modelling with differentiable fluid solvers: physics-based loss functions and optimisation horizons'. In: Journal of Fluid Mechanics 949 (2022), A25.

[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 on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021 . OpenReview.net, 2021. URL: https://openreview.net/forum?id=c8P9NQVtmnO .

[Li+22] S. Li, Z. Huang, T. Du, H. Su, et al. 'Contact Points Discovery for Soft-Body Manipulations 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 / BF01589116 . URL: https://doi.org/10.1007/BF01589116 .

[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, P. 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. DOI: 10.1038/S42256-021-00302-5 . URL: https://doi.org/10.1038/s42256-021-00302-5 .

[Mar+96] G. J. Martyna, M. E. Tuckerman, D. J. Tobias, and M. L. Klein. 'Explicit reversible integrators for extended systems dynamics'. In: Molecular Physics 87.5 (1996), pp. 11171157.

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

[MRO20] D. Maoutsa, S. Reich, and M. Opper. 'Interacting Particle Solutions of Fokker-Planck Equations Through Gradient-Log-Density Estimation'. In: Entropy 22.8 (2020), p. 802. DOI: 10.3390/e22080802 . URL: https://doi.org/10.3390/e22080802 .

[Øks03] B. Øksendal. 'Stochastic differential equations'. In: Stochastic differential equations . Springer, 2003, pp. 65-84.

[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 Machine Learning Research 22.1 (2021), pp. 2617-2680.

[Par21] É. Pardoux. Stochastic Partial Differential Equations: An Introduction . Springer, 2021.

[Pfa+20] T. Pfaff, M. Fortunato, A. Sanchez-Gonzalez, and P. Battaglia. 'Learning mesh-based simulation with graph networks'. In: International Conference on Learning Representations (2020).

[Ram+20] Z. Ramzi, B. Remy, F. Lanusse, J. Starck, and P. Ciuciu. 'Denoising Score-Matching for Uncertainty Quantification in Inverse Problems'. In: CoRR abs/2011.08698 (2020). URL: https://arxiv.org/abs/2011.08698 .

[SC20] S. Schoenholz and E. D. Cubuk. 'Jax md: a framework for differentiable physics'. In: Advances in Neural Information Processing Systems 33 (2020), pp. 11428-11441.

[SE19] Y. Song and S. Ermon. 'Generative Modeling by Estimating Gradients of the Data Distribution'. In: Advances in Neural Information Processing Systems . 2019, pp. 1189511907. URL: https : / / proceedings . neurips . cc / paper / 2019 / hash / 3001ef257407d5a371a96dcd947c7d93-Abstract.html .

[SLG22] C. Salvi, M. Lemercier, and A. Gerasimovics. 'Neural Stochastic PDEs: ResolutionInvariant Learning of Continuous Spatiotemporal Dynamics'. In: NeurIPS . 2022. URL: http : / / papers . nips . cc / paper % 5C _ files / paper / 2022 / hash / 091166620a04a289c555f411d8899049-Abstract-Conference.html .

[Soh+15] J. Sohl-Dickstein, E. A. Weiss, N. Maheswaranathan, and S. Ganguli. 'Deep Unsupervised Learning using Nonequilibrium Thermodynamics'. In: Proceedings of the 32nd International Conference on Machine Learning . Vol. 37. JMLR Workshop and Conference Proceedings. JMLR.org, 2015, pp. 2256-2265. URL: http://proceedings. mlr.press/v37/sohl-dickstein15.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://proceedings.mlr.press/ v115/song20a.html .

[Son+21a] Y. Song, C. Durkan, I. Murray, and S. Ermon. 'Maximum Likelihood Training of ScoreBased Diffusion Models'. In: Advances in Neural Information Processing Systems . 2021, pp. 1415-1428. URL: https://proceedings.neurips.cc/paper/2021/ hash/0a9fdbb17feb6ccb7ec405cfb85222c4-Abstract.html .

[Son+21b] Y. Song, J. Sohl-Dickstein, D. P. Kingma, A. Kumar, et al. 'Score-Based Generative Modeling through Stochastic Differential Equations'. In: 9th International Conference on Learning Representations . OpenReview.net, 2021. URL: https://openreview. net/forum?id=PxTIG12RRHS .

[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://openreview. net/forum?id=vaRCHVj0uGI .

[Sta+21] K. Stachenfeld, D. B. Fielding, D. Kochkov, M. Cranmer, et al. 'Learned Simulators for Turbulence'. In: International Conference on Learning Representations . 2021.

[Thu+21] N. Thuerey, P. Holl, M. Mueller, P. Schnell, et al. Physics-based Deep Learning . WWW, 2021. URL: https://physicsbaseddeeplearning.org .

[Tom+17] J. Tompson, K. Schlachter, P. 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://openreview.net/forum?id=ByH2gxrKl .

[Um+20] K. Um, R. Brand, Y. ( Fei, P. Holl, and N. Thuerey. 'Solver-in-the-Loop: Learning from Differentiable Physics to Interact with Iterative PDE-Solvers'. In: Advances in Neural Information Processing Systems . 2020. URL: https://proceedings.neurips.cc/ paper/2020/hash/43e4e6a6f341e00671e123714de019a8-Abstract.html .

[Vin11] P. Vincent. 'A Connection Between Score Matching and Denoising Autoencoders'. In: Neural Comput. 23.7 (2011), pp. 1661-1674.

[YSM22] R. Yang, P. Srivastava, and S. Mandt. 'Diffusion Probabilistic Modeling for Video Generation'. In: CoRR abs/2203.09481 (2022). DOI: 10.48550/arXiv.2203.09481 . URL: https://doi.org/10.48550/arXiv.2203.09481 .

[ZC21] Q. Zhang and Y. Chen. 'Diffusion normalizing flow'. In: Advances in Neural Information Processing Systems 34 (2021), pp. 16280-16291.

[ZDG96] K. Zhou, J. Doyle, and K. Glover. 'Robust and optimal control'. In: Control Engineering Practice 4.8 (1996), pp. 1189-1190.

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.

Problem setting Let (Ω , F , P ) be a probability space and W ( t ) = ( W 1 ( t ) , ..., W D ( t )) T be a D -dimensional Brownian motion. Moreover, let x 0 be a F 0 -measurable R D -valued random variable that is distributed as p 0 and represents the initial simulation state. We consider the time evolution of the physical system for 0 ≤ t ≤ T modeled by the stochastic differential equation (SDE)

with initial value x 0 and Borel measurable drift P : R D → R D and diffusion g : [0 , T ] → R ≥ 0 . This SDE transforms the marginal distribution p 0 of initial states at time 0 to the marginal distribution p T of end states at time T .

Moreover, we assume that we have sampled N trajectories of length M from the above SDE with a fixed time discretization 0 ≤ t 0 < t 1 < ... < t M ≤ T for the interval [0 , T ] and collected them in a training data set { ( x ( n ) t m ) i M =0 } N n =0 . For simplicity, we assume that all time steps are equally spaced, i.e., t m +1 -t m := ∆ t . Moreover, in the following we use the notation x ℓ : m for 0 ≤ ℓ < m ≤ M to refer to the trajectory ( x t ℓ , x t ℓ +1 , ..., x t m ) .

Assumptions Throughout this paper, we make some additional assumptions to ensure the existence of a unique solution to the SDE (11) and the strong convergence of the Euler-Maruyama method. In particular:

Finite variance of samples from p 0 : E x 0 ∼ p 0 [ || x 0 || 2 2 ] < ∞

Lipschitz continuity of P : ∃ K 1 > 0 ∀ x , y ∈ R D : ||P ( x ) -P ( y ) || 2 ≤ K 1 || x -y || 2

Lipschitz continuity of g : ∃ K 2 > 0 ∀ t, s ∈ [0 , T ] : | g ( t ) -g ( s ) | ≤ K 3 | t -s |

Linear growth condition: ∃ K 3 > 0 ∀ x ∈ R D : ||P ( x ) || 2 ≤ K 3 (1 + || x || 2 )

g is bounded: ∃ K 4 > 0 ∀ t ∈ [0 , T ] : | g ( t ) | ≤ K 4

Euler-Maruyama Method Using Euler-Maruyama steps, we can simulate paths from SDE (11) similar to ordinary differential equations (ODE). Given an initial state X t 0 , we let ˆ X ∆ t t 0 = X t 0 and define recursively

where z t m are i.i.d. with z t m ∼ N (0 , I ) . For t i ≤ t < t i +1 , we define the piecewise constant solution of the Euler-Maruyama Method as ˆ X ∆ t t := ˆ X ∆ t t i . Let X t denote the solution of the SDE (11). Then the Euler-Maruyama solution ˆ X ∆ t t converges strongly to X t .

Lemma A.1. [Strong convergence of Euler-Maruyama method] Consider the piecewise constant solution ˆ X ∆ t t of the Euler-Maruyama method. There is a constant C such that

Proof. See Kloeden et al. [Klo+92, p. 10.2.2]

A.1 1-step Loss and Score Matching Objective

Theorem A.2. Consider a data set { x ( n ) 0: m } N n =1 with trajectories sampled from SDE (11) . Then the 1-step loss

is equivalent to minimizing the score matching objective

where ˜ s θ ( x , t ) = s θ ( x , t ) /g 2 ( t ) as ∆ t → 0 .

Proof.

" ⇐ ": Consider θ ∗ such that ˜ s θ ∗ ( x , t ) ≡ ∇ x log p t ( x ) , which minimizes the score matching objective J ( θ ) . Then fix a time step t and sample x t and x t +∆ t from the data set. The probability flow solution x p t based on equation (14) is

At the same time, we know that the transformation of marginal likelihoods from p t +∆ t to p t follows the reverse-time SDE [And82]

which runs backward in time from T to 0 . Denote by ˆ x ∆ t t the solution of the Euler-Maruyama method at time t initialized with x t +∆ t at time t +∆ t .

Using the triangle inequality for squared norms, we can write

Because of the strong convergence of the Euler-Maruyama method, we have that for the first term of the bound in equation (18)

independent of θ . At the same time, for the Euler-Maruyama discretization, we can write

where z t +∆ t is a standard Gaussian distribution, i.e., z t +∆ t ∼ N (0 , I ) . Therefore, we can simplify the second term of the bound in equation (18)

If θ ∗ minimizes the score matching objective, then ˜ s θ ∗ ( x , t ) ≡ ∇ x log p t ( x ) , and therefore the above is the same as

Combining equations (18), (19) and (25) yields

Additionally, since g is bounded, we even have

For L single ( θ ) , using the above bound (27) and strong convergence of the Euler-Maruyama method, we can therefore derive the following bound

which implies that L single ( θ ) → 0 as ∆ t → 0 .

" ⇒ ": With the definitions from " ⇐ ", let θ ∗ denote a minimizer such that L single ( θ ) → 0 as ∆ t → 0 , i.e., we assume there is a sequence ∆ t 1 , ∆ t 2 , ... with lim n →∞ ∆ t n = 0 and a sequence θ 1 , θ 2 , ... , where θ n is a global minimizer to the objective L ∆ t n single ( θ ) that depends on the step size ∆ t n . If there is θ ∗ such that s θ ∗ ( x , t ) ≡ ∇ x log p t ( x ) , then L ∆ t n single ( θ n ) ≤ L ∆ t n single ( θ ∗ ) . From " ⇐ " we know that lim n →∞ L ∆ t n single ( θ ∗ ) = 0 and therefore lim n →∞ L ∆ t n single ( θ n ) = 0 . This implies that each summand of L single ( θ ) also goes to zero as ∆ t → 0 , i.e., lim ∆ t → 0 E [ || x t -x p t || 2 2 ] = 0 . Again, with the triangle inequality for squared norms, we have that

By the strong convergence of the Euler-Maruyama method and θ = θ ∗ , we obtain

At the same time, for fixed ∆ t > 0 , we can compute

For fixed x t +∆ t , the distribution over z ∼ N (0 , I ) in equation (36) correspond to a noncentral chi-squared distribution [JKB95, Chapter 13.4], whose mean can be calculated as

For each ∆ t > 0 , the above is minimized if and only if ∇ x log p t +∆ t ( x t +∆ t ) = s θ ( x t +∆ t , t +∆ t ) .

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, we consider our proposed probability flow ODE defined by

and for t i < t j define p t j , ODE t i as the distribution obtained by sampling x ∼ p t j and integrating the probability flow with network s θ ( x , t ) equation (7) backward in time until t i . We can choose two arbitrary time points t i and t j with t i < t j because we do not require fixed start and end times of the simulation.

The maximum likelihood training objective of the probability flow ODE (7) can be written as maximizing

Our proposed multi-step loss is 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

where we compute the expectation by sampling x 0: m from the training data set and ˆ x i : i + S -1 is the predicted sub-trajectory that is defined recursively by

In the following, we show that the multi-step loss (4) maximizes a variational lower bound for the maximum likelihood training objective (40).

Theorem A.3. Consider a data set { x ( n ) 0: m } N n =1 with trajectories sampled from SDE (11) . Then the multi-step loss (4) maximizes a variational lower bound for the maximum likelihood training objective of the probability flow ODE (40) as ∆ t → 0 .

Let µ ODE t i ( x t j ) denote the solution of the probability flow ODE (7) integrated backward from time t j to t i with initial value x t j .

For the maximum likelihood objective, we can derive a variational lower bound

where the inequality is due to Jensen's inequality. Since only p t j , ODE t i ( x t i | x t j ) depends on θ , this is the same as maximizing

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, e.g., Kersting et al. [Ker+20, Gaussian ODE filtering]. This allows us to model the conditional distribution p t j , ODE t i | x t j as a Gaussian distribution with mean µ t j , ODE t i ( x t j ) and variance σ 2 = ϵ . Then maximizing (47) reduces to matching the mean of the distribution, i.e., minimizing

independent of ϵ > 0 . Since this is true for any time step t j > t i and corresponding simulation state x t j given x t i , we can pick the pairs ( x t i , x t i +1 ) , ( x t i , x t i +2 ) , ( x t i , x t i +3 ) and so on. Then, we can optimize them jointly by considering the sum of the individual objectives up to a maximum sliding window size

As ∆ t → 0 , we compute the terms µ ODE t k ( x t j ) on a single trajectory starting at x t i with sliding window S covering the trajectory until x t j via the Euler method, i.e., we can define recursively

By varying the starting points of the sliding window x t i , this yields our proposed multi-step loss L multi ( θ ) .

A.3 Denoising Score Matching for Deterministic Simulations

So far, we have considered physical systems that can be modeled by an SDE, i.e., equation (11). While this problem setup is suitable for many scenarios, we would also like to apply a similar

methodology when the system is deterministic, i.e., when we can write the problem as an ordinary stochastic equation

In the case of chaotic dynamical systems, this still represents a hard inverse problem, especially when information is lost due to 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 7b and 7c illustrate two additional data setup and inference variants for deterministic physical systems modeled by the ODE (51). 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 -P ( x ) and s θ ( x , t ) is separated into two steps. The temporal evolution from t i +1 to t i is then defined entirely by physics. We apply an additive noise to the system and the update step by s θ ( x , t ) , which can be interpreted as denoising for a now slightly perturbed state ˜ x t i . In this case, we show that the network s θ ( x , t ) still learns the correct score ∇ x log p t ( x ) during training using denoising score 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

where z t i ∼ N (0 , I ) . If the physics system is deterministic and ∆ t is small enough, then we can approximate x t i ≈ ˆ x t i and for the moment, we assume that we can write

In this case, we can rewrite the 1-step loss L single ( θ ) from (14) to obtain the denoising score matching loss

which is the same as minimizing

Now, the idea presented in Vincent [Vin11] is that for score matching, we can consider a joint distribution p t i ( x t i , ˜ x t i ) of sample x t i and corrupted sample ˜ x t i . Using Bayes' rule, we can write p t i ( x t i , ˜ x t i ) = p σ (˜ x t i | x t i ) p t i ( x t i ) . The conditional distribution p σ ( ·| x t i ) for the corrupted sample is then modeled by a Gaussian with standard deviation σ = √ ∆ t g ( t i ) , i.e., we can write ˜ x = x + √ ∆ t g ( t i ) z for z ∼ N (0 , I ) similar to equation (55). Moreover, we can define the distribution of corrupted data q σ as

If σ is small, then q σ ≈ p t i and KL( q σ || p t i ) → 0 as σ → 0 . Importantly, in this case, we can directly compute the score for p σ ( ·| x ) as

Moreover, the theorem proven by Vincent [Vin11] means that we can use the score of the conditional distribution p σ ( ·| x ) to train s θ ( x , t ) to learn the score of q σ ( x ) , i.e.

By combining (61) and (59), this exactly equals the denoising loss L DSM ( θ ) in (57). As q σ ≈ p t i , we also obtain that ∇ x log q σ ( x ) ≈ ∇ x log p t i ( x ) , so the network s θ ( x , t i ) approximately learns the correct score for p t i .

We have assumed (55) that the only corruption for ˆ x noise t i is the Gaussian noise. This is not true, as we have

so there is an additional source of corruption, which comes from the numerical errors due to the term x t i +1 -∆ t P ( x t i +1 ) -x t i . The conditional distribution p σ ( ·| x ) is only approximately Gaussian. Ideally, the effects of numerical errors are dominated by the Gaussian random noise. However, even small errors may accumulate for longer sequences of inference steps. To account for this, we argue that the multi-step loss L multi ( θ ) should be used. During training, with the separation of physics update and denoising, the simulation state is first progressed from time t i +1 to time t i using the reverse physics solver. This only yields a perturbed version of the simulation at time t i due to numerical inaccuracies. Then a small Gaussian noise is added and, via the denoising network s θ ( x , t ) , the simulation state is projected back to the distribution p t i , 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.

B Architectures

ResNet We employ a simple ResNet-like architecture, which is used for the score function s θ ( x , t ) and the convolutional neural network baseline (ResNet) for the stochastic heat equation in section 3.2 as well as in section 3.4 again for the score s θ ( x , t ) .

For experiments with periodic boundary conditions, we apply periodic padding with length 16, i.e., if the underlying 2-dimensional data dimensions are N × N , the dimensions after the periodic padding are ( N +16) × ( 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 t is concatenated as an additional constant channel to the 2-dimensional input data when this architecture represents the score s θ ( x , t ) .

The encoder part of our network begins with a single 2D-convolution 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-convolution, and Leaky ReLU. All 2D convolutions have 32 filters with kernel size four and stride 1. The encoder part ends with a single 2D convolution with one filter, kernel size 1, and no activation. Then, in the decoder part, we begin with a transposed 2D convolution, 32 filters, and kernel size 4. Afterward, there are four consecutive residual blocks, analogous to the residual encoder blocks, but with the 2D convolution replaced by a transposed 2D convolution. Finally, there is a 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 [Mue+22], Appendix A.1, 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.

FNO The FNO-2D architecture introduced in [Li+21] with k max ,j = 12 Fourier modes per channel is used as a baseline for the stochastic heat equation experiment in section 3.2 and the learned physics surrogate model in section 3.4.

Dil-ResNet The Dil-ResNet architecture is described in [Sta+21], Appendix A. This architecture represents the score s θ ( x , t ) in the buoyancy-driven flow with obstacles experiment in section 3.3. We concatenate the constant time channel analogously to the ResNet architecture. Additionally, positional information is added to the network input by encoding the x -position and y -position inside the domain in two separate channels.

Architecture

Parameters

Table 2: Summary of architectures.

C 1D Toy SDE

For the 1D toy SDE discussed in section 3.1, we consider the SDE given by

with λ 1 = 7 and λ 2 = 0 . 03 . The corresponding reverse-time SDE is

Throughout this experiment, p 0 is a categorical distribution, where we draw either 1 or -1 with the same probability. In figure 8, we show trajectories from this SDE simulated with the Euler-Maruyama method. Trajectories start at 1 or -1 and approach 0 as t increases.

Neural network architecture We employ a neural network s θ ( x, t ) parameterized by θ to approximate 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-VR]. In all cases, the neural network is a 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 hyperparameters as described in the original paper [KB15]. The learning rate, batch size, and the number of epochs depend on the data set size (100% with 2 500 trajectories, 10%, or 1%) and are chosen to ensure convergence of the training loss.

(a)

λ

2

= 0

(b)

λ

2

= 0

.

03

Training - 1-step loss For the 1-step loss and all data set sizes, we train for 250 epochs with a learning rate of 10e-3 and batch size of 256 . In the first phase, we only keep every 5th point of a trajectory and discard the rest. Then, we again train for 250 epochs with the same batch size and a learning rate of 10e-4 but keep all points. Finally, we finetune the network with 750 training epochs and a 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-step loss, which resembles a sliding window size of 2. We initially train for 1 000 epochs with a batch size of 512 and a learning rate of 10e-3, where we keep only every 5th point on a trajectory and discard the rest. Then, with a decreased learning rate of 10e-4, we begin training with a sliding window size of S = 2 and increment it every 1 000 epochs by one until S max = 10 . In this phase, we train on all points without any removals.

Training - ISM For ISM, we compute the partial derivative ∂s θ ( x ) i /∂ x i using reverse-mode automatic differentiation in JAX (jax.jacrev). For 100% and 10% of the data set, we train for 2 000 epochs with a learning rate of 10e-3 and batch size of 10 000. Then we train for an additional 2 000 epochs with a learning rate 10e-4. For 1%, increase the number of epochs to 20 000.

Training - SSM-VR For sliced score matching with variance reduction [Son+19, SSM-VR], we use the same training setup as for ISM.

(a) ISM learned score.

(b) Multi-step learned score.

(c) ISM reverse-time SDE trajectories.

(d) Multi-step reverse-time SDE trajectories.

(e) ISM probability flow trajectories.

(f) Multi-step probability flow trajectories.

75

75

1.0

1.0

50

50

0.5

25

0.5

25

0.0

0

0.0

0

-0.5

-25

-0.5

-25

-50

-50

-1.0

1.0

75

-75

0

2

4

6

8

10

0

2

4

6

8

10

t

t

1.0

1.0

0.5

0.5

0.0

0.0

-0.5

-0.5

-1.0

-1.0

10

8

6

t

4

2

0

10

8

6

t

4

2

0

1.0

1.0

0.5

0.5

0.0

0.0

-0.5

-0.5

-1.0

-1.0

10

8

6

4

2

0

10

8

6

4

2

0

t

t

Table 3: Posterior metric

Q

for 1 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 in figure 9 trained on the full data set. The learned score of ISM and the multi-step loss in figure 9a 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 are more accurate compared to ISM. For example, in figure 9e, a trajectory explodes to negative infinity. Also, trajectories from the multi-step loss end in either -1 or 1 , while ISM trajectories are attenuated and do not fully reach -1 or 1 exactly, particularly for the probability flow ODE.

Results of Table 1 in Main Paper We include the standard deviations of table 1 from the main paper in table 3 above. The posterior metric Q is very sensitive to the learned score s θ ( x , t ) . Overall, our proposed multi-step loss gives the most consistent and reliable results.

Empirical verification of Theorem 3.1 For the quadratic SDE equation 63, the analytic score is non-trivial, therefore a direct comparison of the learned network s θ ( x , t ) and the true score ∇ x log p t ( x ) is difficult. However, we can consider the SDE with affine drift given by

with λ = 0 . 5 and g ≡ 0 . 04 . Because this SDE is affine and there are only two starting points -1 and 1 , we can write the distribution of states starting in x 0 as a Gaussian with mean µ ( t ; x 0 ) = x 0 e -λt and variance σ 2 ( t ; x 0 ) = g 2 2 λ (1 -e -2 λt ) , see [Øks03]. Then, the score at time t and position x conditioned on the starting point x 0 is ( x -µ ( t ; x 0 )) /σ 2 ( t ; x 0 ) . See figure 10 for a visualization of the analytic score and a 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.

0

2

4

6

8

10

1.0

0.5

0.0

0.5

1.0

x

(a) Paths and density

0

2

4

6

8

10

1.0

0.5

0.0

0.5

1.0

x

50

0

50

(b) Analytic score field

0

2

4

6

8

10

t

1.0

0.5

0.0

0.5

1.0

x

50

0

50

(c) 1-step training

0

2

4

6

8

10

t

1.0

0.5

0.0

0.5

1.0

x

50

0

50

(d) Multi-step training

D Stochastic Heat Equation

Spectral solver The physics of the 2-dimensional heat equation for x ∈ R d × d can be computed analytically. The solver P ∆ t h ( x ) simulates the systems x forward in time by a fixed ∆ t using the (shifted) Fourier transformation F . In particular, we can implement the solver with

where ◦ denotes element-wise multiplication and A (∆ t ) ij ∈ R d × d is a matrix with entries A (∆ t ) ij := exp ( -∆ t · min( i, j, d -i, j -i )) . The power spectrum of x is scaled down by A (∆ t ) , and higher frequencies are suppressed more than lower frequencies (for ∆ t > 0 ). If noise is added to x , then this means that especially higher frequencies are affected. Therefore the inverse transformation (when ∆ t > 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 1 and s 2 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

with a weighting vector w ∈ R d and w k = 1 for k ≤ 10 and w k = 0 otherwise.

Training For inference, we consider the linear time discretization t n = n ∆ t with ∆ t = 0 . 2 / 32 and t 32 = 0 . 2 . During training, we sample a random time discretization 0 ≤ t 0 < t ′ 2 < .... < t ′ 31 < t 32 for each batch based on t n via t ′ n ∼ U ( t n -∆ t/ 2 , t n + ∆ t/ 2) for n = 1 , ..., 31 and adjust the reverse physics step based on the time difference t i -t i -1 . In the first training phase, we consider the multi-step loss with a sliding window size of S = 6 , 8 , ..., 32 steps, where we increase S every two epochs. We use Adam to update the weights θ with learning rate 10 -4 . We finetune the network weights for 80 epochs with an initial learning rate of 10 -4 , which we reduce by a factor of 0 . 5 every 20 epochs.

s θ 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 s θ ( 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 s θ ( 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 x 0 and end state x T .

Test-time distribution shifts We have tested the effects of test-time distribution shifts for the heat equation experiment. We train the score network s θ for a specific combination of diffusivity α 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 11. 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.

Additional results We show visualizations for the predictions of different methods in figure 12 and figure 13.

10 -4

10 -3

MSE

10 -5

10 -4

10 -3

10 -2

10 -1

10 0

-30% -20% -10%

1.0

+10%+20%+30%

Diffusivity

10 -1

Spectral error

-30% -20% -10%

0.1

+10%+20%+30%

Noise

10 -2

10 -1

10 0

(a) Reverse-time SDE

10 -4

MSE

10 -5

10 -4

10 -3

10 -2

10 -1

10 0

-30% -20% -10%

1.0

+10%+20%+30%

Diffusivity

10 -2

Spectral error

-30% -20% -10%

0.1

+10%+20%+30%

Noise

10 -2

10 -1

10 0

(b) Probability flow ODE

E Buoyancy-driven Flow with Obstacles

We use semi-Lagrangian advection for the velocity and MacCormack advection for the hot marker density within a fixed domain Ω ⊂ [0 , 1] × [0 , 1] . The temperature dynamics of the marker field are modeled with a 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 2 until S max = 20 .

Separate vs. joint updates We compare a joint update of the reverse physics step and corrector function s θ , see figure 7b, and a separate update of reverse physics step and corrector function, see figure 7c. An evaluation regarding the reconstruction MSE and perceptual distance is shown in

(a) Ground truth

(b) Input

(c) SMDP - ODE

(d) ResNet

(e) FNO

(f)

s

θ

only - ODE

(g)

s

θ

(

x

, t

)

only - SDE

(h) BNN

(i) SMDP - SDE

(a) Ground truth

(b) Input

(c) SMDP - ODE

(d) ResNet

(e) FNO

(f)

s

θ

only - ODE

(g)

s

θ

only - SDE

(h) BNN

(i) SMDP - SDE

approaches. Overall, there are no apparent differences for the ODE inference performance but slight benefits for the SDE inference when separating physics and corrector update.

Limited-memory BFGS We use numerical optimization of the marker and velocity fields at t = 0 . 35 to match the target smoke and velocity fields at t = 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 [Pas+19]. As arguments, we use history _ size = 10 , max _ iter = 4 and line _ search _ fn = strong _ wolfe . Otherwise, we leave all other arguments to the default values. We optimize for 10 steps which takes ca. 240 seconds per sample on a single NVIDIA RTX 2070 gpu.

Diffusion posterior sampling DPS An additional baseline for this problem is diffusion posterior sampling for general noisy inverse problems [Chu+23, DPS]. As a first step, we pretrain a diffusion model on the data set of marker and velocity fields at t = 0 . 35 . We use the mask for obstacle positions as an additional conditioning input to the network to which no noise is applied. 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 -4 using cosine annealing with warm restarts ( T 0 = 50000 , η min = 5 × 10 -6 ). The measurement operator A is implemented using our differentiable forward simulation. We consider the Gaussian version of DPS, i.e., Algorithm 1 in [Chu+23] with N = 1000 . We fix the step size ζ i at each iteration i to 1 / || y -A (ˆ x 0 ( x i )) || . For each inference step, we are required to backpropagate gradients through the diffusion model and the forward simulation. Inference for a single sample requires ca. 5000 seconds on a single NVIDIA RTX 2070 gpu.

Additional results We provide more detailed visualizations for the buoyancy-driven flow case in figure 16 and figure 17. 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.

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 in this scenario leaves only little room for substantial changes in the states.

˜

P -

1

only

s

θ

only

joint update

separate update

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

MSE (

t

=

0

.

35)

SMDP SDE

SMDP ODE

(a) Reconstruction MSE (

t

= 0

.

35

).

˜

P -

1

only

s

θ

only

joint update

separate update

10

-

2

10

-

1

LPIPS (

t

=

0

.

35)

SMDP SDE

SMDP ODE

(b) Perceptual distance LPIPS (

t

= 0

.

35

).

G

Ground truth

SMDP ODE

SMDP SDE

t

=

0

.

65

˜

P -

1

only

t

=

0

.

55

t

=

0

.

45

t

=

0

.

35

t

=

0

.

25

t

=

0

.

15

Ground truth

SMDP ODE

SMDP SDE

t

=

0

.

65

˜

P -

1

only

t

=

0

.

55

t

=

0

.

45

t

=

0

.

35

t

=

0

.

25

t

=

0

.

15

x

(b) Velocity field (

-direction)

Ground truth

SMDP ODE

SMDP SDE

t

=

0

.

65

˜

P -

1

only

t

=

0

.

55

t

=

0

.

45

t

=

0

.

35

t

=

0

.

25

t

=

0

.

15

Ground truth

SMDP ODE

SMDP SDE

t

=

0

.

65

˜

P -

1

only

t

=

0

.

55

t

=

0

.

45

t

=

0

.

35

t

=

0

.

25

t

=

0

.

15

(b) Velocity field (

x

-direction)

Ground truth

100 steps

200 steps

300 steps

400 steps

500 steps

750 steps

1000 steps

2000 steps

Ground truth

100 steps

200 steps

300 steps

400 steps

500 steps

750 steps

1000 steps

2000 steps

p-1

only

only

F Isotropic turbulence

Training For the physics surrogate model ˜ P -1 , we employ an FNO neural network, see appendix B, with batch size 20 . We train the FNO for 500 epochs using Adam optimizer with learning rate 10 -3 , which we decrease every 100 epochs by a factor of 0 . 5 . We train s θ ( x , t ) with the ResNet architecture, see appendix B, for 250 epochs with learning rate 10 -4 , decreased every 50 epochs by a factor of 0 . 5 and batch size 6 .

Refinement with Langevin Dynamics Since the trained network s θ ( x , t ) approximates the score ∇ x log p t ( x ) , it can be used for post-processing strategies [WT11; SE19]. We do a fixed point iteration at a single point in time based on Langevin Dynamics via:

for a number of steps T and ϵ = 2 × 10 -5 , cf. figure 18 and figure 19. For a prior distribution π t , x 0 t ∼ π t and by iterating (68), the distribution of x T t equals p t for ϵ → 0 and T → ∞ . There are some theoretical caveats, i.e., a Metropolis-Hastings 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 trajectories in figure 20 and figure 21.

Ground truth

1  only

SMDP-ODE

t

=10

SMDP-SDE

t

=9

t

=8

t

=6

t

=4

Ground truth

1  only

SMDP-ODE

t

=10

SMDP-SDE

t

=9

t

=8

t

=6

t

=4

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 Sampling for General Noisy Inverse Problems'. In: The Eleventh International Conference on Learning Representations . OpenReview.net, 2023. URL: https://openreview. net/pdf?id=OnD9zGAGT0k .

[HJA20] J. Ho, A. Jain, and P. Abbeel. 'Denoising Diffusion Probabilistic Models'. In: Advances in Neural Information Processing Systems . 2020. URL: https://proceedings. neurips . cc / paper / 2020 / hash / 4c5bcfec8584af0d967f1ab10179ca4b 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: NeurIPS 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://jmlr.org/papers/ v6/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 on Learning Representations . 2015. URL: http://arxiv. 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, ICML 2020, 13-18 July 2020, Virtual Event . Vol. 119. Proceedings of Machine Learning Research. PMLR, 2020, pp. 51985208. URL: http://proceedings.mlr.press/v119/kersting20a.html .

[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 on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021 . OpenReview.net, 2021. URL: https://openreview.net/forum?id=c8P9NQVtmnO .

[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 / BF01589116 . URL: https://doi.org/10.1007/BF01589116 .

[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, F. Massa, A. Lerer, et al. 'PyTorch: An Imperative Style, HighPerformance Deep Learning Library'. In: Advances in Neural Information Processing Systems . 2019, pp. 8024-8035. URL: https://proceedings.neurips.cc/paper/ 2019/hash/bdbca288fee7f92f2bfa9f7012727740-Abstract.html .

[RHS22] S. Rissanen, M. Heinonen, and A. Solin. 'Generative Modelling With Inverse Heat Dissipation'. In: arXiv:2206.13397 (2022).

[SE19] Y. Song and S. Ermon. 'Generative Modeling by Estimating Gradients of the Data Distribution'. In: Advances in Neural Information Processing Systems . 2019, pp. 1189511907. URL: https : / / proceedings . neurips . cc / paper / 2019 / hash / 3001ef257407d5a371a96dcd947c7d93-Abstract.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://proceedings.mlr.press/ v115/song20a.html .

[Sta+21] K. Stachenfeld, D. B. Fielding, D. Kochkov, M. Cranmer, et al. 'Learned Simulators for Turbulence'. In: International Conference on Learning Representations . 2021.

[Vin11] P. Vincent. 'A Connection Between Score Matching and Denoising Autoencoders'. In: Neural Comput. 23.7 (2011), pp. 1661-1674.

[WT11] 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 2, 2011 . Omnipress, 2011, pp. 681-688.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n### 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.\n\nWEAKNESSES:\n- 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\nQUESTIONS:\n- 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.\n\nLIMITATIONS:\nThe authors mention limitations in their conclusion.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n*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.\n\nWEAKNESSES:\n* 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.\n\nQUESTIONS:\n* 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..\n\nLIMITATIONS:\nNA", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-05"}

comments: 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.

STRENGTHS:
The authors propose multi-step loss to capture long-range dependence in physics system, which has potential implication for general diffusion models.
The 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. 

Minor/errata

Hats are missing in equation (4)

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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2023-07-07

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\nThe 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.\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-13"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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\nSTRENGTHS:\nThe 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\nWEAKNESSES:\nThere 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\nQUESTIONS:\n– 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\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-25"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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).\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\n- Releasing the code could be crucial as there are many delicate details one might need to get right for the method to work.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-28"}