Score Matching via Differentiable Physics
Authors: Benjamin Holzschuh
Simona Vegetti
Nils Thuerey
Venue: Submitted to ICLR 2023 (2023)
Keywords: diffusion models
physics simulations
stochastic partial differential equations

Abstract: Diffusion models based on stochastic differential equations (SDEs) gradually perturb a data distribution $p(\mathbf{x})$ over time by adding noise to it. A neural network is trained to approximate the score $\nabla_\mathbf{x} \log p_t(\mathbf{x})$ at time $t$, which can be used to reverse the corruption process. In this paper, we focus on learning the score field that is associated with the time evolution according to a physics operator in the presence of natural non-deterministic physical processes like diffusion. A decisive difference to previous methods is that the SDE underlying our approach transforms the state of a physical system to another state at a later time. For that purpose, we replace the drift of the underlying SDE formulation with a differentiable simulator or a neural network approximation of the physics. At the core of our method, we optimize the so-called probability flow ODE to fit a training set of simulation trajectories inside an ODE solver and solve the reverse-time SDE for inference to sample plausible trajectories that evolve towards a given end state. We demonstrate the competitiveness of our approach for different challenging inverse problems.

Under review as a conference paper at ICLR 2023
SCORE MATCHING VIA DIFFERENTIABLE PHYSICS
Anonymous authors Paper under double-blind review
ABSTRACT
Diffusion models based on stochastic differential equations (SDEs) gradually per- turb a data distribution p(x) over time by adding noise to it. A neural network is trained to approximate the score Vx log (x) at time t, which can be used reverse the corruption process. In this paper, we focus on learning the score field that is associated with the time evolution according to a physics operator in the presence of natural non-deterministic physical processes like diffusion. A decisive difference to previous methods is that the SDE underlying our approach transforms the state of a physical system to another state at later time. For that purpose, we replace the drift of the underlying SDE formulation with a differen- tiable simulator or neural network approximation of the physies. our method, we optimize the so-called probability flow ODE to fit training set of simulation trajectories inside an ODE solver and solve the reverse-time time SDE for inference to sample plausible trajectories that evolve towards given end state. We demonstrate the competitiveness of our approach for different challenging inverse problems.
the core of
INTRODUCTION
Many physical systems are time-reversible on microscopic scale. For example, a continuous mate- rial can be represented by collection of interacting particles (Gurtin 1982 Blanc al. 2002) based on which we can predict future states of the material. We can also compute earlier states, meaning we can evolve the simulation backwards in time (Martyna al.. (996). When taking a macroscopic perspective, we only know the average quantities within specific regions (Farlow 1993), which con- stitutes a loss of information. It only then that time-reveriibility is no longer possible, since many macroscopic and microscopic initial states exist that evolve to yield the same macroscopic state.
In the following, we target inverse problems to reconstruct the distribution of initial macroscopic states for given end state. This genuinely tough problem has applications in many areas of scien- tific machine learning (Zhou et al., 1996; Goonez-Bommbrelll et al., 2018 Delaquis et al., 2018 Lim & Psaltis 2022). and existing methods lack tractable approaches represent and sample the distribution of states. We address this issue by leveraging continuous approaches for diffu- sion models in the context of physical simulations. In particular, our work builds on the reverse- diffusion theorem (Anderson, 1982). Given the functions f(.,t) Ra + R called drift, and g(-) : IR + E. called diffusion, it can be shown that under mild conditions, for the forward stochas- tic differential equation (SDE) dx = f(x, t)dt + g(t) dw there is a corresponding reverse-time SDE dx = [f(x, t) V log P: (x) dt + (t. In particular, this means that given a marginal distribution of states Pol (x) at time t = O and pr (x) at a T such that the forward SDE transforms px(x) to pT (x), then the reverse-time SDE runs backward in time and transforms pr(x) into po (x).
The term Vx log (x) is called the score.
This theorem is a central building block for SDE--based diffusion models and denoising score match- ing (Song et al. 202 Moolicocur-Martineae et al. 2021). which parameterize the drift and diffusion in such a way that the forward SDE corrupts the data and transforms it into random noise. By train- ing neural network represent the score, the eeverse-time SDE can be deployed model, which transforms samples from random noise (x) to the data distribution (x()).
generative
In this paper, show that similar methodology processes. We replace the drift fxx t) by a physics model P(x) Rd - RR, which is implemented by differentiable solver neural network that represent the dynamics of physical system, thus
likewise be employed to model physical

---

Under review as a conference paper at ICLR 2023
Training: Probubility Dow dx [P(x) (1)) 4)
physics model P(x): diffusion ((t) 3
trajoctories ((x))))
Inforenee: time
(P(x) (1) ()) a(l)dw
marginal
Score field Vx log approximated by so(x, t)
likelihood
po(x)
marginal
likelihood pr(x)
Data generation dx = P(x))d g(t)
Time
Figure I: Overview: we employ a physics simulator P to learn the score field V. log (x() with neural network in the presence of noise uncertainties. The trained model allows for sampling the posterior of po, i.c. different states that explain observation pr. via probability flow solving the reverse-time SDE,
deeply integrating physical knowledge into our method. The end state at = T on which the forward SDE acts is not fully destroyed by the diffusion g(t), but instead, the noise acts perturbation of the system state over time. An overview of our method is shown in figure
To the best of our knowledge, our work is the first to leverage the reverse- difffusii theorem method for solving inverse problems of physical systems. As such, our primary aim is to demon- strate how existing algorithms from this field can be used in the context physics simulations. We showcase the efficacy of the score matching viewpoint on physics problems with range of challenging inverse problems. Specifically, our contributions are:
1. We develop framework in which we incorporate the everse-diffusion theorem and score matching into a method for solving inverse problems that involve the time evolution of physical systems. We demonstrate its competitiveness against common baseline ap- proaches using the heat equation as an example.
2. We highlight the effectiveness of our method with more challenging inverse problem where we simulate luiddbased transport process backwards in time in the presence of randomized obstacles. Here, we compare our method to different strategies for learned solvers.
3. Finally, we show that this approach can even be used when the underlying SDE is unknown. Our approach can be combined with operator learning methods and demonstrate its effectiveness for learning the Navier-Stokes equation in the turbulent regime.
2 BACKGROUND AND RELATED WORK
Learned solvers: Numerical simulations benefit greatly from machine learning models (Tompson et al. 2017 Morton et al.) 2018 Pfaff et al. 2020 Li et al. 2020). By integrating neural network inside differential equation solvers, it is possible learn to reduce numerical errors (Tompson al.
2017 Kochkov et al.) 2021 Brandstetter et al. 2022)
guide the simulation towards a desired target state (Holl et al. 2020b) et al. 2022). As errors may accumulate quickly over time, trained networks benefit from gradients that are backpropagated over multiple time steps (Um et al. 2020).
Diffuusion models: Diffusson models (Ho et al. 2020; Song et al. 2021c) have been considered for wide range of applications. Most notably, diffusion models have been proposed for image (Dhariwal & Nichol 2021). video (Ho et al. 20022] Höppe et al. 2022 Yang et al. (002) and audio synthesis (Chen al. 2021). Recently, [ansal al. (2022 have proposed to train generalized diffusion models for arbitrary transformations and suggest that fully deterministic models without any noise sufficient for generative behaviour.
2

---

Under review as a conference paper at ICLR 2023
Specifically for uncertainty quantification, solving inverse problems and conditional sampling many
methods have been proposed (Chung et al., 2021 2022 Song et al. 2021b Kawar et al. 2021 Ramzi et al. 2020). However, most approaches either focus on the denoising objective that com- mon for tasks involving natural images, or the synthesis process of solutions does not take the underlying physics directly into account.
Generative modeling via SDEs: Classical denoising score matching approaches based on Langevin dynamics (Vincent 2011 Sonn & Ermon 2019 SMLD) and based on discrete Markov chains, e-g- et al., 2001 Ho et al. 2020 DDPM), can
Denoising Diffusion Probabilistic Models
be unified in a continuuss framework using SDEs (Song et al. (021c). Given distribution of states Po (x) at time t = 0, an SDE transforms (x to tractable distribution (x
dx f(x + g(t)dw,
with w the standard Brownian motion, a drift f(.,)) RR + R and diffiusion g()) : R R, which for X (x() yields diffusion process (x, (x=) As outlined above, the reverse-time SDE of the reverse-diffusion theorem (Anderson) 1982) is given by
dx [f t) g(t) Vx log]( + g())dū,
with Vx log Px being the score. By sampling from the tractable distribution pr(x) and simu- lating the reverse-time SDE equation 2 we can generate samples from po Since the score Vx (x) is not known analytically, it is approximated by neural network.
Continuous normalizing flows: Continuous normalizing flows (CNFs) are invertible generative models based on neural ODEs. Given an initial distribution 2g Pag (2o) and a function f (=((t) t; O) represented by a neural network, CNF trained to solve the ODE (z = f(=(t), 1;0) with boundary conditions =(4)) = 2g and =(4 =(41) x. where x E Rd is sample from the training data set. A useful property of CNFs is that is cheap to compute the log-likelihood of data samples, due
to the instantancous change of variables formula Ô log p(=(t))
- -Tr (t)(t)) (Grathwohl
jet al. (2019). The evolution of the marginal prohability density (x( for the SDE in equation (T)
deseribed by Kolmogorov's forward equation (Øksendal] 2003). Maoutsa et al. (2020) and Song iet al.] (2021c) show that there exists an ODE with the same Kolmogorov forward equation. This ODE is called probability flow ODE and is given by
dx
f(x,
1 gg(()
2 V log (x() dt.
The probability flow ODE equation
3 represents CNF and, if f(x, t) is known, a network sa parameterized by 9 representing log Pt (x) can be trained via maximum likelihood using standard methods (Chen et al. 2018). While the evolution of pt (x) is the same between the probability flow
ODE from equation (3) and the reverse-time SDE from equation (2), there are caveats due to the
approximation by X, t) (Song et al. 2021b; Lu et al. 2022). Huang et al. (2021) show that minimizing the score-matching loss is equivalent to maximizing a lower bound of the likelihood obtained by sampling from the reverse-ting time SDE.
A recent variant combines score matching with CNFs (Zhang & Chen 2021). and employs a joint
training of the drift and score with an integration backwards in time.
3 METHOD
Modeling assumptions We consider known physics model P(x) Rd - Rd that is differen- tiable and approximates the time evolution, i.e. Xim-1 FO Xt,.. (tm+1 - tm)) -P(X,... One of key modelling choices is to describe the time evolution of the physical system by a stochastic differential equation
dx P(x)dt + g(t)dw,
with a diffusion process g(.) : R + IR that perturbs the simulation states. We can simulate paths from this SDE using Euler-Maruyame steps, i.e. for a time discretization to tt < tM and initial state Xt. we obtain the iteration rule
Xtat1 Xtm. + (tm+1 tm) .P(xe.)) Vtnt1 t. g(tm)
(5)
3

---

Under review as a conference paper at ICLR 2023
Training
Inference
SMDP-ODE
mag//aao 2o0t .TK
bithrectinn Indirecionnll rogniaa
dx = P(x) - (00 logp:(x) dt
Batt-aat Mudt Whableaur
+ Dowioie
SMDP-SDE
(P(x) (V log (xx]dd glt)db
SMDP-ODE
SMDP-SDE
Figure 2: During the training phase, we optimize (x t) that approximates the score Vx log (x) inside the probability flow ODE to fit the data trajectories. For the inference part, we simulate trajectories of the reverse-time SDE.
where 2t... are i.i.d. with 2t... No. I). The random noise ((mm)2 that is added to the data at each time step can be regarded as either noise inherent to the physical problem, or as noise from a measurement process. As training data, we consider a set of N trajectories (X...))= n sampled with equation 5 and which describe the evolution of physical system.
Inverse Problem Given an end state XT we are interested in recovering a likely tra- jectory (xx ) that evolves towards Xy. More formally, the set of trajectories ((xr..)im0..... M. i N implicitly defines marginal likelihoods px(x) at every time step & which are linked through the SDE equation ( of the physical system by the Kolmogorov forward equa- tion. The solution trajectory may be unique, so we want sample from the full posterior instead of obtaining only a maximum likelihood solution, i.c we want to sample from Po (xl xr).
Method In line with previous work in score-based generative modeling (Song & Ermon 2019 Song et al. 2021c). we approximate the score Vx Pt (x) of the marginal likelihoods by a neural network sp (x, t). We optimize 5x t) via maximum likelihood training of the probability flow ODE, as discussed in section 2 For this, we maximize a variational lower bound for the maximum
likelihood objective, which we estimate by minimizing the following loss
M
C ((Xx) M 0) =
Σ
||x,
ODE XT,
(6)
ODE XT f DDX) 1 X
(t)ss (XD t)dt,
where we sample SDE trajectory (xx)) M from the training set. We give theoretical justification for this objective in appendix A Intuitively, our method fits bijective and deterministic trajectories of the probability flow ODE to the non non-determinisic SDE trajectories. In contrast to previous work, our method deeply integrates a prior about the physical system in the form of the simulation operator T is not fully destroyed by the
P(x) into the training process. In this context, the end state noise, but instead the noise acts as a perturbation of the system state over time. An overview of our method is shown in figure
Given an end state XT, we can solve the probability flow ODE backwards in time using the trained score function (x,, t) to obtain trajectory (xpreed ) However, this will only give sin- gle, deterministic solution and not allow for sampling from the posterior p(x [x)). We simulate trajectories from the reverse-time SDE (see section I via
dx [P(x) - g2 (1))))) (x, O) dt + g(t)dw.
The evolution of marginal probabilities pe (x) for this SDE is the same for the probability flow
ODE equation 7 (Song et al.) 2021c). Moreover, by the reverse-diffusion theorem (Anderson 1982]. SDE equation (8) the time-reverse of the physical system SDE from equation fore, we can appxocimate sampling from p(x by simulating trajectories from SDE equation { using any traditional SDE solver. In the following, we refer to the integration of the physics model P(x) into the score-based modelling approach as score matching via differentiable physics, or
There-
SMDP in short, We denote trajectories from the probability flow ODE by SMDP-ODE, and those obtained by simulating the reverse-time SDE by SMDP-SDE.
4

---

Under review as a conference paper at ICLR 2023
Training and Inference Algorithm 1 gives
Algorithm SMDP-ODE, SMDP-SDE
an overview of SMDP inference for the ODE as
well as the SDE variant when using the explicit
Require: Xlar, {tmm) JMM (9tm P for m t M to do
Euler method as ODE solver. For simplicity, we also employ the explicit Euler method for train-
1:
2:
p P(xim)
ing and backpropagate gradients through multi-
+- Si. s9 (m) 2m)
ple solver steps when computing the ODE tra- jectory in equation 7 to obtain updates for A. We also refer this procedure as unrolling the dynamics, Our training setup is similar to Um al. (020)) which was originally developed for training correction functions in the con- text of controlling numerical errors for physical simulations. In particular, in our implementa- tion, we consider a sliding window for unrolling the dynamics, which makes our training very flex- consider single-step updates well unrolling the entire simulation. We give
SMDP-ODE then
5:
Xem-1 + Xem (fm Em 1m
6:
if SMDP-SDE then
7:
Xt Xem (fi 1m- ~~ (0,1)
9:
X1m-1 t X1m + Vlm - 1m--12 12
10: return Xto
ible, i.e. more detailed descriptions about our training setup in appendix We consider an additional variant of SMDP, for which we apply bidirectional training, i.e. instead of only training the probability flow ODE for the time-backwar direction T + O we alternate with optimizing the time-forward direction M, i.c. equation becomes
ODE Xy [ P(XDDE) 2
1 (t)s (XDEE ,t)dt,
4 EXPERIMENTS
We conduct several experiments demonstrate the advantages of SMDP compared number of baseline methods, The source code for all experiments will be made available upon acceptance. We first consider the 2D heat equation in section 4.1 where objective is to find possible initial states at time / given noisy end state at time
In our second experiment in section 4.2 we transfer the established practices more challenging problem, where we are interested in reconstructing the trajectory of buoyancy-driven flow within closed simulation domain given an end state at time T.. What makes this problem challenging is that for each simulation, we place different obstacles at different positions within the simulation domain. Then, in section 4.3 we consider the situation where the physics of the system is unknown. For this purpose, we consider training network that approximates the solutions to the Navier- Stokes equations and a network that approximates the score field. We demonstrate that by doing we obtain improved performance for inverse problems and the learned score be used to refine predictions in post-processing. We provide additional results in the appendix. In appendix compare our method with implicit score matching (Hyvärinen) 2005) in a toy experiment, which demonstrates the importance of include physics dynamics in the training. We analyze how many steps are required when unrolling the dynamics for obtaining stable trajectories in appendix Finally, in appendix we give an additional evaluation of quality and diversity of the posterior distribution we obtain when sampling from the reverse-time SDE in the heat equation experiment.
C
0.2
diffusion
SMDP-ODE
SMDP-SDE
Figure 3: Heat diffusion case. We simulate a Gaussian random field at t 0 forwards in time using equa- tion Given either solve the probability flow ODE simulate trajectories of the reverse- time SDE obtain solutions for the state at 0.
5

---

Under review as a conference paper at ICLR 2023
Method
MSE [10- $
Spectral error
Full posterior
SMDP-ODE
0.74
3.62
X
SMDP-SDE
5.56
0.56
V
Data
SMDP-ODE
Spestral denaity
SMDP SDE
ResNet-S
2.17
1.67
X
ResNee-S
ResNet-P
2.30
1.09
X
ResNer-P
BNN-S
3.47 X 10
1.25
10
BNN-P
BNN-P
3.81 x 102
0.99
V
FNO-S
FNO-S
2.54 x 104
1.60
FNO-P
10-5
FNO-P
2.50 x 104
1.47
X
Difusion-noiee
HeatGen 1.39
4.45
X
10
HeatGen+noise
4.45
3.24
V
Figure 4: Spectral density different scales, the red line indicating ground truth. The closer method the ground truth, the better it produces structures of a similar scale.
Table 1: Evaluation of reconstruction MSE and spectral cr- ror for SMDP and baselines. The columm full posterior indi- cates whether models yield p estimates or allow to sample from the posterior.
4.1 HRAT EQUATION
ă5 o&u which plays fundamental role in many physical systems. Here, we set the diffusivity constant to = and initial conditions at f = O are generated from Gaussian random fields with n 4 at resolution 32 x 32. We simulate the heat diffusion using a spectral method until f = 0.2 with fixed step size At = 6.25 x 10-3 using the iteration rule from equation S with g 0.1. Our training data set consists of 2.500 initial conditions with their corresponding trajectories sampled with varying step size At and end states at t 0.2. The test set comprised of 500 initial conditions and corresponding end states generated directly without any noise.
We consider the heat equation du
Training We consider a small ResNer-like architecture based on an encoder and decoder part (see appendix Eb as representation for the score function sp (x, t). The physics model P is implemented differentiable programming JAX (Schoenholz Cubuk 2020). For better comparison with the baseline methods, these are trained with Gaussian random noise of 0.1 added to the inputs. This noise applied to all network inputs during testing.
Baseline methods As baseline methods, we consider the ResNet-like architecture from above, in addition Bayesian neural network (BNN) based on a U-Net architecture with spatial dropout (Mueller et al.) [022). as well as Fourier neural operator (FNO) network (Li et al., 2020). For each of these three methods, we consider two variants: the first variant is trained with a supervised loss, i.e. the training data consists of pairs (x,,Xy) with initial state X, and end state xy- The supervised loss corresponds to the squared L2 distance between the network prediction x, xfred and the ground truth, i.e. (xn pred - 2
2 For the second variant, the reconstruction loss, we rely on the differentiable solver and only make use of the end state X such that the loss becomes (P(xpred, pred, x²) -x)) i.e we which we compare to the desired end state Xy. We denote the supervised variant by S and the physics-based one by P. Additionally, we consider an adopted generative model from Rissanen al. (2022). denoted by HeatGen. We train this network similarly to SMDP-ODE, but without the solver P. such that the network has to learn the score and the physics at the same time.
simulate the network output forward in time using to obtain state
Reconstruction accuracy vs. fitting the data manifold We give an evaluation of our method and the baselines by considering the reconstruction MSE on the test set: how well a predicted initial state &n that is simulated forward in time yields states that correspond to the reference end state in terms of MSE. This metric has the disadvantage that it does not measure how well the prediction matches the training data manifold, iee. for this case whether the prediction resembles the statistics of the Gaussian random field. For that reason, we additionally compare the power spectral density of the states as the spectral loss. The corresponding measurements are given in table 1 which

---

Under review as a conference paper at ICLR 2023
Ground trunh
Markes donstry
Valudly (dd(x)
SADDE
0.55
0.35
0.25
Figure Buoyancy flow case. Ground truth shows marker density and velocity field the r-direction different points of the simulation trajectories. The simulation end state at = 0.65 is the input to SMDP-ODE and SMDP-SDE.
show that method SMDP-ODE performs best in terms of the reconstruction MSE. However, solutions obtained by SMDP-ODE are very smooth and do not contain the small-scale structures of the references, which is reflected in high spectral error that is also visually prominent, as shown in figure 4 SMDP-SDE on the other hand performs very well in terms of spectral error and yields visually convincing solutions with only slight increase in the reconstruction loss. We note that there is a natural tradeoff between both metrics, and SMDP-ODE and SMDP-SDE perform best for both cases respectively while using the
of weights.
Ablation study We performed ablation study to highlight several design choices of the proposed method. In particular note that despite fundamental differences between SMDP-ODE and SMDP-SDE, as explained section 3] the main difference constant factor for and the noise term g(t)d for SMDP-SDE, cf. equations (7) and (8). We investigated how the change in noise integration affects the perfor- mance of SMDP-ODE, and considered a variant SMDP. ODE+noise that includes the addition of the noise term during inference, but is otherwise identical ODE. As shown with crosses in figure 5] this method has a slightly higher reconstruction loss but in contrast to SMDP-SDE does not improve upon the spectral error. This indicates that SMDP-SDE can recover small-scale distributions of the references. while the ODE variant by construction tends to favour smooth solutions when faced with uncertainties. We additionally investigated the effect of the proposed bidirectional training scheme. The error measurements of Figure show that ODE variant is not affected, but SMDP-SDE significantly benefits. The re- sulting model robustly handles wide range of different temporal discretizations for inference, which deviate from the training discretization, with high ac- curacy. In conclusion, our SMDP-SDE model with bidirectional training yields the best performance for wide range of hyperparameter settings. In appendix F we additionally evaluate the effects of logarithmic time discretization and a physics-conditioned score field.
inference time is the
10'1
MSR
10
10
SMDP.
10
x0.25
x0.5
x01
x8
Infos
slops
Figure 5: Reconstruction MSE and spectral eтrors for the bidirectional (top) and regular variant (bottom). The x-axis shows the rela- tive increase of the numher of time steps dur- ing inference compared to training. Large errors are truncated at the top of each graph.
4.2 BUOYANCY-DRIVEN FLOW WITH OBSTACLES
Next, we test our methodology on a more challenging problem. For this purpose we consider sim- ulations of buoyancy-driven flow within fixed domain Ω c [0 [0, and randomly placed

---

Under review as a conference paper at ICLR 2023
obstacles. We make use of semi--agrangii advection for the velocity and MacCormack advec- tion for the hot marker density. The temperature dynamics of the marker field are modeled with a Boussinesq approximation. Each simulation runs from time t = 0.0 t 0.65 with step size of At = 0.01. The inflow at (0.2, 0.5) is active until f = 0.2. Our objective is to employ SMDP-ODE and SMDP-SDE obtain trajectories that reconstruct plausible flow given an end state of the marker density and velocity fields at time & = 0.65.
Our training data set consists of 250 simulations with corresponding trajectories. For the data gen- eration, we make use of the differentiable phiflow solver (Holl et al. (0200)). 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. The testing set comprises 5 simulations. In contrast to the previous task, we generate the training data set without any noise, but add a Gaussian random noise with standard deviation o state of the training trajectories.
V to each simulation
Training and Results We em- ploy a neural network architec- ture based on dilated convolu- tions (Stachenfeld et al.] 2021 for (x, i), see appendix E The physics operator Pxx) imple mented by using a negative step size AAt for time integration. The training of SMDP-ODE is analogous the previous exper- iments on heat equation in sec- tion 4.1 However, we do not un- roll the entire trajectory and ap- ply sliding window with size
Configuration
MSEL
LPIPS,,
[10-3]
ODE
SDE
ODE
SDE
SMDP
0.05
0.78
5,82
147.17
Rollout noise Physics cond. score
0.12
1.03
9.10
145.72
0.15
1.25
5.94
360.54
Bidir. training Decoupled score&phys.
0.4
0.61
24.19
134.21
0.07
0.72
7.78
67.23
Fully learned Physics only
0.28
12.54
0.5
346.31
Table 2:
Evaluation of variants for the buoyancy obstacle case in terms reconstruction MSE and LPIPS of the marker field.
20. For this case, we compare our method to fully learned baseline method. A quantitative eval- uation in terms of reconstruction MSE and LPIPS metrics (Zhagg et al., (018)) is given in table 2 It becomes apparent that our ODE method clearly outperform the learned baseline. Interestingly, the SDE variant performs less well for this test case. This behavior can be explained by the highly nonlinear system dynamics, and the comparatively approximate reverse simulator which yields the substantial errors for the Physics only version in table 2 These errors causes the score network to inadvertently learn significant corrections of the physics operator, which deteriorates the quality of the score field. Nonetheless, as qualitatively shown in figure 66 both variants are able to accurately recover the initial states despite the complex motion of the fluid around the obstacles.
Algorithmic variants We evaluate several altered configurations of our SMDP method to further illustrate its behavior. We consider adding noise during rollouts, i.c. adding a noise term o z for ~ (O, I) to the state after applying the score and physics updates during training. Additionally, we experiment with a physics conditioned score, i.e. we extend the input dimension of the score function accept concatenated inputs of the form sn [x, P(x)], l). We also evaluate the effects of the bidirectional training. Finally, we evaluate a version that decouples score and physics, i.e. we do not evaluate the score and physics update on the same input x, but instead we first apply the physics update P P(x). and evaluate the score function afterwards. Overall, the eттor measurements in table justify our choices for the baseline SMDP algorithm.
4.3 ISOTROPIC TURBULENCE
As third example, consider problem where the physics operator is unknown, i.e. we approxi- mate both P and the score Vx log)x(() by neural networks. We consider the problem of learning the time evolution of isotropic, forced turbulence as determined by the 2D Navier-Stokes equations with a viscosity of P = 10 (Li et al. 2020). The training data set consists of vorticity fields from 1000 simulation trajectories from 64 x 64. Our objective is to predict trajectory Su:10 that reconstructs the true trajectory Xu:10 given an end state x of the solution, whereas in the original paper, the objective was learn an operator mapping predicting the vorticity at later point in time.
until T 10 with At 1 and a spatial resolution of
8

---

Under review as a conference paper at ICLR 2023
Ground truth
Learned physics
SMDP-ODE
SMDP-SDE
SMDP-SDE+LD
Figure Turbulence case. Comparison of reconstructed trajectories (
Methods We represent the physics operator (xx) by FNO neural network and the score by the baseet-based architecture from section 4.1 As first step of our method, we optimize the physics model P using our SMDP training setup from the previous two experiments, but fix the score to sx(x,t) t) = O. We denote this method by Learned physics. As a second step, we train SMDP-ODE with the decoupled score and physics variant by optimizing the time-dependent score sy sg(x, t) while freeziig the physics model P. This approach guarantees that any time- line-nddependen physics are captured by P and 6x 1) can focus on learning small improvements to as well as respond to possibly time-dependent data biases. As this test employs two trained networks for both components of the SDE, we compare to DiffFlow Zhang & Chen 2021 as an additional baseline.
Evaluation We give evaluation of the improvements of SMDP over the learned variants in table Learned physics variant, our methods im- prove the mean squared error (MSE) be- tween the ground truth trajectories and the reconstructed trajectories slightly, while there is an substantial decrease in the spec- error. This can be seen qualitatively figure 7 In this scenario DiffFlow has severe difficulties to learn state updates and score field, sulting in large differences terms of MSE. As before, the SMDP-SDE method performs best in terms of spectral error at the expense of a slightly increased MSE.
Method
MSEJ [10-2]
Spectral Error
Compared the
SMDP-ODE SMDP-SDE
14.3
0.26
15.3
0.17
DiffFlow
221.2
0.58
Learned physics
16.3
0.30
Table 3: Evaluation the turbulence case.
Outlook: Refinement with Langevin dynamics Since the score Vx log Pr (x) represents data gradient, we can use gradient-based optimization methods to find local optima of the data distribu-
tion (x) that are close to Inspired by stochastic gradient Langevin dynamics (Welling & Teh
2011 we consider the iteration rule x ++1 x + x log Pt (xf) + /2Ez,, for 2 10-5 where ~ (N. (details in appendix H.I Denoted by SMDP-SDE+LLE in figure 7 this method
manages to extract even finer details from the reverse- SDE solution. As such provides an
interesting starting point for a further refinement of the SMDP results.
5 CONCLUSION
We presented SMDP, derivative of score matching in the context of physical simulations and dif- ferentiable physics. We demonstrated its competitiveness against different baseline methods and in challenging inverse physics problems. We demonstrated the versatility SMDP with two variants: while the neural ODE variant focuses on high MSE accuracies, the neural SDE variant allows for sampling the posterior and yields an improved coverage of the target data manifold. Despite the promising initial results, our work gives rise to many interesting questions. In particular, the time discretizations is crucial issue, as for training data generation and differentiable solvers, we would favor larger step size and less evaluations due to computational constraints, while for conventional diffusion models, a large number of smaller time steps often yields substantial improvements in terms of quality. Determining good balance between accurate solutions with few time steps and diverse solutions with many time steps will remain an important direction for future research this area.

---

Under review as a conference paper at ICLR 2023
REFERENCES
Brian DO Anderson. Reverse-time diffusion equation models. Stochastic Processes and their Ap- plications, 12(3):31-332, 1982.
Arpit Bansal, Eitan Borgnia, Hong-Min Chu, Jie S. Li, Hamid Kazemi, Furong Huang, Micah Goldblum, Jonas Geiping, and Tom Goldstein. Cold diffusion: Inverting arbitrary image trans- forms without noise. CoRR, abs/2208.09392, 2022. doi: 10.48550/a2Xxv.2208.0933 URL https: //doi .rg.orggggoo /t0. B08885500aiiii 2208. 09392
Xavier Blane, Claude Le Bris, and P-L Lions. From molecular models to continuum mechanics. Archive for Rational Mechanics and Analysis, 164(4):341-381, 2002.
Valentin De Bortoli, James Thorton, Jeremy Heng, and Arnaud Doucet. Diffusion schrödinger bridge with applications to score-based generative modeling. In Mare' Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan (eds.), Advances Neural Information Processing Systems 34: Annual Conference
Neu-
ral Information Processing Systems 2021, NeurlPS 2021, December 6-14, 2021, virtual, pp. 17699-17709, 2021. URL htttss:/ www.iinnnnnnneu.rneeeeeeeeee..ee www.nnnurip... wweurips.ccppppeee..ee ccpppppppp/p2////// www.e:a17777tteeee1-aaattaaaaaae htmll
Johannes Brandstetter, Daniel Worrall, and Max Welling. Message passing neural pde solvers. International Conference on Learning Representations, 2022.
Nanxin Chen, Yu Zhang, Heiga Zen, Ron Weiss, Mohammad Norouzi, and William Chan. Wave- grad: Estimating gradients for waveform generation. In 9th International Conference on Learning Representattons, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net 2021. URL
net www.oooroorrsaooooooo jcF a080
Tian Qi Chen, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud. Neural ordinary dif- ferential equations. In Advances in Neural Information Processing Systems 31: Annual Con- ference on Neural Information Processing Systems 2018, NeurIPS 2018, December 3-8, 2018, Montréal, Canada, pp. 6572 6583, 2018. URL http:ttttttttteeeedeeeese eign wwwrinuur..... cc/ Abstract .tml
69 633836622 http:///.aaaaa....oooooooea22222
Tianrong Chen, Guan-Horng Liu, and Evangelos A. Theodorou. Likelihood training of schrödinger bridge using kawardd-bakwaad sdes theory. In The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022. OpenReview.net, 2022. URL
httpps:/cpeeneeview net/ forum? www.oooooooooooonooooeoooo id=nioAdkCEaR.
Hyungjin Chung, Byeongsu Sim, and Jong Chul Ye. Come--looer-difftus--fftee Accelerat- ing conditional diffusion models for inverse problems through stochastic contraction. CoRR, abs/2112.051 2021. URL httthhttss//eec /awpiv.aaaivvv/aa//ea org hbsgg//b///111 05146
Hyungjin Chung, Byeongsu Sim, Dohoon Ryu, and Jong Chul Ye. Improving diffusion models for inverse problems using manifold constraints. CoRR, abs/2200.000411, 2022. doi: 10.48550/arXiv. 2206.00941. URL httpt:ttt////t /doii wwt. wtw.:////.....//a.../////..//// 2206.....222...... 00941
Delaquis, Jewell, I Ostrovskiy, Weber, T Ziegler, Dalmasson, LJ Kaufman, Richards, JB Albert, G Anton, et al. Deep neural networks for energy and position reconstruction in exo- 200. Journal of Instrumentation, P3(00) 2018.
Prafulla Dhariwal and Alexander Quinn Nichol. Diffusion models beat gans on image synthe- sis. In Advances in Neural Information Processing Systems 34: Annual Conference Neu- ral Information Processing Systems 2021, NeurlPS 2021, December 6-14, 2021, virtual, pp.
8780 8794, 2021. URL https: http://///cocccecpsee//em www.ieeriii. wttpppppe//2////aaaaeeaa/a 11d000222266ffffffaaaaa Abstract html
Stanley Farlow. Partial differential equations for scientists and engineers. Courier Corporation, 1993.
10

---

Under review as a conference paper at ICLR 2023
Rafael Coamez-Bombarelli Jennifer N Wei, David Duvenaud, José Miguel Hernández-Lobato, Timothy D Hirzel,
Benjamín Baanchez-Lengeling Dennis Sheberla, Jorge Ryan P Adams, and Alán Aspuru-Guzik. Automatic chemical design using data-driven contin- uous representation of molecules. ACS central science, 4(2):268-276, 2018.
Will Grathwohl, Ricky T. Q. Chen, Jesse Bettencourt, Ilya Sutskever, and David Duvenaud. FFJORD: free-form continuous dynamics for scalable reversible generative models. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net, 2019. URL https::/ net forum?id-
EJxgknCeK7
Morton E Gurtin. An introduction to continuum mechanics. Academic press, 1982.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic mod- els. In Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12. 2020, virtual, 2020. URL 4c5ccfec33.aaa0001111110.... htt.:://1..cc..aattttttttttttt. html
httppssscccpppppp cc/papprr2222/haatt
Jonathan Ho, Tim Salimans, Alexey A. Gritsenko, William Chan, Mohammad Norouzi, and David J. Fleet. Video diffusion models. CoRR, abw/2004.03458, 2022. doi: 1044855WanXxx22244.0488
URLhtttps /top::://dd.... .org/ 10. 48540 wt.55//arxi... 201.:lvv2204404.55.5 03458
Philipp Holl, Vladlen Koltun, Kiwon Um, and Nils Thuerey. phiflow: A differentiable pde solving framework for deep learning via physical simulations. In NeurlPS Workshop, volume 2, 2020a.
Philipp Holl, Nils Thuerey, and Vladlen Koltun. Learning to control pdes with differentiable physics. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net, 2020b. URL httts::/ wtt.:/rrueeeeeeoommmm forum? id=HyeSin4FPH
net net
Tobias Höppe, Arash Mehrjou, Stefan Bauer, Didrik Nielsen, and Andrea Dittadi. Diffusion models for video prediction and infilling. CoRR, abw//200.006666 2022, doi: 10.48550/rrxiv.2206.7766
URL h/tp::://doo.ooo
org/ 4ww.//////...22222
Chin-Wei Huang, Jae Hyun Lim, and Aaron C. Courville. A variational perspective on diffusion-based generative models and score matching. In Marc' Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan (eds.), Ad- vances in Neural Information Processing Systems 34: Annual Conference
Neural Infor- mation Processing Systems 2021, NewrIPS 2021, December 6-14, 2021, virtual, pp. 22863- 22876, 2021. URL 1tt.00011111100000aa--oaaataaaaa
http:/pproceed.......i...ss..ss nww.ips..pp...... htc/pappr///2///// haash/ c1 21labfa22eecc44dd556666 6b0
Abst html
Aapo Hyvärinen. Estimation of non-normalized statistical models by score matching. J. Mach. Learn. Res., 6:695-709, 2005. http:
html
Alexia Dolicoeur-Martinen Ke Li, Rémi Piché-Taillefer, Tal Kachman, and Ioannis Mitliagkas. Gotta go fast when generating data with score-based models. arXiv preprint 1xXx::2005.14080 2021.
Bahjat Kawar, Gregory Vaksman, and Michael Elad. SNIPS: solving noisy inverse problems stochastically. In Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurlPS 2021, December 6-14, 2021, virtual,
pp. 21757-21769, 2021. URL https: http://owceedii.....eeeee.i....e
wwt.:sccc-ssttttt.tcc html
Dmitrii Kochkov, Jamie A. Smith, Ayya Alieva, Qing Wang, Michael P. Brenner, and Stephan Hoyer. Machine learning accelerated computational fluid dynamics. CoRR, abs/2102.01010, 2021. URL
hutpps://aae
org/ abs 2102. 01010
Sizhe Li, Zhiao Huang, Tao Du, Hao Su, Joshua B Tenenbaum, and Chuang Gan. Contact points discovery for bodt-body manipulations with differentiable physics. International Conference on Learning Representations, 2022.
11

---

Under review as a conference paper at ICLR 2023
Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, An- drew Stuart, and Anima Anandkumar. Fourier neural operator for parametric partial differential equations. arXiv:2010.08895, 2020.
Joowon Lim and Demetri Psaltis. Maxwellnet: Physics-driven deep neural network training based on maxwell's equations. Apl Photonics, 7(1):011301, 2022.
Cheng Lu, Kaiwen Zheng, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Maximum likelihood training for score-based diffusion odes by high order denoising score matching. In International Conference on Machine Learning. ICML 2022, 17-23 July 2022, Baltimore, Maryland, USA, volume 162 of Proceedings of Machine Learning Research, pp. 14429-14460. PMLR, 2022. wwu2f..htmm
URL
pwwss/eevvvvv... vi htt.s/ss....prreseeeeeeeeeeeeeee
mlr
Dimitra Maoutsa, Sehastian Reich, and Manfred Opper. Interacting particle solutions of fokker- planck equations through gradient-log-denstty estimation. Entrpyy 22(8):802, 2020. doi: 10. 3390/e22080802. URL www.https:.. /tiii wrw.10........0000......0. 0ttp://20..3300022208000000
Glenn J Martyna, Mark E Tuckerman, Douglas Tobias, and Michael I. Klein. Explicit reversible integrators for extended systems dynamies. Molecular Physics, 87(5):11771117,, 1996.
Jeremy Morton, Antony Jameson, Mykel J Kochenderfer, and Freddie Witherden. Deep dynamical modeling and control of unsteady fluid flows. In Advances in Neural Information Processing Systems, 2018
Maximilian Mueller, Robin Greif, Frank Jenko, and Nils Thuerey, Leveraging stochastic predictions of bayesian neural networks for fluid simulations. arXiv preprint arXiv:2205.0122 2022.
Bernt Øksendal. Stochastic differential equations. In Stochastic differential equations, pp. 65-84. Springer, 2003.
Tobias Pfaff, Meire Fortunato, Alvaro Sanchez-Gonzale,, and Peter Battaglia. Learning mesh-based simulation with graph networks. International Conference on Learning Representations, 2020.
Zaccharie Ramzi, Benjamin Remy, François Lanusse, Jean-Luc Starck, and Philippe Ciuciu. Denois- ing score-matching for uncertainty quantification in inverse problems. CoRR, abs/201 I. 100699, 2020. utLpttps: / arxiv orq/ abb///////// htw.b////1100005.
Severi Rissanen, Markus Heinonen, and Amn Solin. Generative modelling with inverse heat dissi- pation. arXiv:2206. 13397, 2022.
Samuel Schoenholz and Ekin Dogus Cubuk. Jax md: a framework for differentiable physics. Ad- vances in Neural Information Processing Systems, 33:11428-11441, 2020.
Jascha Sohl-Dickstein, Eric A. Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsu- pervised learning using nonequilibrium thermodynamics. In Proceedings of the 32nd Inter- national Conference on Machine Learning, ICML 2015, Lille, France, 6-11 July 2015, vol- ume 37 of JMLR Workshop and Conference Proceedings, pp. htt.t::/prrcee
JMLR.org, 2015. URL 5tmml
ngs mm wrw.s/esss/ee... preess/vi vv7 0
soh
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribu- tion. Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurlPS 2019, December 8-14, 2019, Vancouver, BC, Canada, pp. 11895-11907, 2019. URI, 2019 20099hashh3000 hhwshh/00eee25/11 2tt.:0/eee2227/1/aa111a111113.3.
wersseeurisss hwtp:pppcccppprr//pppp wwt.eesttatttttttt.a.aeacaa..e.. htmi
Yang Song, Conor Durkan, Iain Murray, and Stefano Ermon. Maximum likelihood training of score-based diffusion models. In Advances in Neural Information Processing Systems 34: An- nual Conference
Neural Information Processing Systems 2021, NeurIPS 2021, December 6- 14, 2021, virtual, pp. 1415-1423 2021a. paper 2001
htt.ttpppe/ppooccennnnn.eeeee 2a.00/hhhhhhham /twhhaaahhhaaaaaaaa 0a 9 ddb.dbbl1 hol.ao/oooocccec0000000222222223
cc/
Yang Song, Liyue Shen, Lei Xing, and Stefano Ermon. Solving inverse problems in medical imaging with basee generative models. CoRR, aabs/2111.080at 2021b. URL https: ///axxv..
org/abs/ 2111 08005
12

---

Under review as a conference paper at ICLR 2023
Yang Song, Jascha Sohl-Dickstein, Diederik P. Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In 9th International Conference Learning Representations. ICLR 2021, Virtall Event. Austria, May 3-7, 2021. OpenReview.net, 2021c. URL cettps::/openreviie ew. net. forum? d=
PxTIG1 2RRHS)
Kim Stachenfeld, Drummond Buschman Fielding, Dmitrii Kochkov, Miles Cranmer, Tobias Pfaff, Jonathan Godwin, Can Cui, Shirley Ho, Peter Battaglia, and Alvaro Sanchez Gonzalezz Learned simulators for turbulence. In International Conference on Learning Representations, 2021.
Jonathan Tompson, Kristofer Schlachter, Pablo Sprechmann, and Ken Perlin. Accelerating eule- rian fluid simulation with convolutional networks. In 5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Workshop Track Proceedings. net. forum? ianruum?di=BB dd
OpenReview.net, 2017. URL htttttt://pppeeertrt//
Kiwon Um, Robert Brand, Yun (Raymond) Fei, Philipp Holl, and Nils Thuerey. Solver- in-the-loop: Learning from differentiable physics to interact with iterative pde-solvers.
In Advances in Neural Information Processing Systems 33: Annual Conference on
Neural Information Processing Systems 2020, NeurlPS 2020, December 6-12, 2020,
virtual, 2020. URL https:://proceeddiii newreeiddiiinnnnnniie.eee.....e neurips /t/ppaper/222//ash/ html
Pascal Vincent. A connection between matching and denoising autoencoders. Neural Com- put.,
2011. doi: 10.1162/NECO, a), 00142. URL httth::////oo org/
10. 11611 NECO 00142]
Max Welling and Yee Whye Teh. Bayesian learning via stochastic gradient langevin dynamics.
In Lise Getoor and Tobias Scheffer (eds.), Proceedings of the 28th International Conference on Machine Learning, ICML 2011, Bellevue, Washington, USA, June 28- July 2011. pp. 681-688. Omnipress, 2011. URL |///
1 wammmpappr... .dff
Ruihan Yang. Prakhar Srivastava, and Stephan Mandt. Diffusion probabilistic modeling for video doi: 10.485550kxxi.2223.044443
generation. CoRR, abs/2203.0948.. 2022.
URL https:
Qinsheng Zhang and Yongxin Chen. Diffusion normalizing flow. Advances in Neural Information Processing Systems, 14:16280- 2021.
Richard Zhang, Phillip Isola, Alexei A Efros, Eli Shechtman, and Oliver Wang. The unreasonable effectiveness of deep features as a perceptual metric. In Proceedings of the IEEE conference on computer vision and patlern recognition, pp. 586-595, 2018.
Kemin Zhou, JC Doyle, and Keither Glover. Robust and optimal control. Control Engineering Practice, 4(8): (100:1119-119999 1190, 1996.
13

---

Under review as a conference paper at ICLR 2023
APPENDIX
A ADDITIONAL DETAILS OF TRAINING METHODOLOGY
Below we summarize the problem formulation from the main paper, and provide details about the training procedure together with further information how the solution of the probability flow ODE solution relates to solutions of the SDE.
Problem setting We consider the time evolution of the physical system modeled by the stochastic differential equation
dx P(x)dt + g(t)dw,
(10)
with drift P Rd - Rd and diffusion g T RT which transforms the marginal distribu- tion Po of initial states at time 0 to the marginal distribution pr of end states at time T. Morcover, assume that we have sampled N trajectories of length .M from the above SDE with a fixed time discretization S to < f1 < < tM I T for the interval [0, 7] and collected them in a training data set { ((xt,n))= (i...))== /n=D.....N
We are interested in training neural network sy (x(x, t) parameterized by 0 to approximate the score vx log (x),, i.e. minimize the score matching objective
JSM (0; (0))(-)))-0))))-(())) 2 L Expp. ||s((x - V log dt,
(11)
where A [0, + R+ is weighting function. In the case of densoing score matching, where the underlying SDE dx = f(,t g(t for affine functions f(-) and g(-), the score can be learned by minimizing the denoising score matching objective using transition kernels, which enables an efficient training of diffusion models, see [Song et al. (2021c) for a reference.
In our case P is an arbitrary function describing the dynamics of the physical system, and hence we can not rely on an analytical expression for the transition kernel.
Training via Continuous Normalizing Flows Score-based diffusion models can be transformed into continuous normalizing flows (Chen et al. 2018 CNFs), which allows for tractable computa- tion of the likelihood. We can train the corresponding CNF given by
1 (1))((()) f) dt
dx P(x)
(12)
using maximum likelihood training, i.c. maximizing
Ex ~p
[lg Po
DDE (xo)]
(13)
s.t. XT = Xu + / P(xi) 2 gg² (t)s t)dt,
(14)
where pe ODE is the distribution obtained by sampling XT pr and simulating x using ODE equation T2p The logelikel can be computed using the instantaneous change of variables formula (Chen al.] (2088) and by using the fact that Pr approximately Gaussian. For denoising score matching, maximum likelihood training of the corresponding CNF is usually not done, because requires running an ODE solver for every optimization step and the training with the denoising
score matching objective is more efficient (Huang al. 2021 Song et al. (021a). It was shown by Song et al. 12021a) that there is connection between the Kullback- divergence and the score matching objective. In particular
(L 1 SDF Jsax((: + KL [|=))
(15)
for a prior distribution т and p3 using the score approximation so sx(x t).
SDE defined by sampling Xy ~ a and solving the reverse- time SDE
14

---

Under review as a conference paper at ICLR 2023
Chaining CNFs For SDE equation 10. the distribution Pr describes the perturbed end states and an exact likelihood computation is no longer possible. Equivalent to maximizing the likelihood of the CNF, we minimize the Kullback-Leible divergence (L || ()DR).
Since in our problem setting, we do not require Po or Pr to be a simple distribution for which we evaluate the likelihood, can think of the CNF from time 0 time T multiple smaller CNFs chained together, e.g. we consider a chain of CNFs for the time discretizations t to to- l2 & and on. Each of those CNFs transforms the marginal probabilities pt, Pt.c via the probability flow ODE 12 by minimizing (L,Ptt Pt, ODE44+1) during training, where Pt, defined by sampling Xt.41 Pt.+1 and simulating the probability flow ODE from time 4++1 until fz-
ODE,44 is
For example, we can sample XT ~ pr and generate data from po by recursively using the CNFs to generate xt, given x until we reach Xp. Since CNFs are bijective, this also works in the reverse direction, i.e. we can sample Xg ~ Po and simulate X from Po with the same method.
In the following, we will derive method using the regular, increasing flow of time, i.e. 0 + T. The direction 7 - follows analogously by replacing the SDE (1O) with the corresponding reverse- time SDE (Anderson 1982). Then, the new objective obtained by chaining the CNFs becomes the minimization of
M Σ
(L l|p ODEA))
(16)
where ODE, Pt++1 is now defined for the time increasing direction, i.c. sampling Xt, ~ Pt. and simulat- ing t44 via the probability flow ODE. If sx t) = Vx Px(),, then by the theory of probability flow ODEs (Song al., 2021c), the above objective will become 0, since the marginal probabilities Pr of the SDE will coincide with the marginal likelihoods of the probability flow CNF at every time
In this case, then we also obtain (L p)DE) 0, i.e. when sampling X ~ Po and solving the probability flow ODE, we obtain exactly the distribution pr.
Additional assumptions In the following, we make the same additionally assumptions Song et al. a Appendix A. Specifically, we require that
(i) aC, > 0 E Rd t [0, T): V po s C₁(1 + ||x||2)
(ii) 3C, > 0 Vx E gd [0, []: [x, & C1 ||x|12
(iii) 3Cp 0 _d [x) |l2 S Cp(1+ [xx
(iv) Vt E [0. 7| 3 (x) E O(e Oce-1xx115 as ||x||2 + 0c
Minimizing the Kullback-Leine divergence fied to
The i-th summand of equation 16 can be simpli-
Ptet (Xteet) ODE,t: Ptis
KL(P.
|l P
PP441 ODE,1: = Ex 10411
log
(17)
(xtr+s)
Pessa (Xtrar)
Exp Exi, NPtct 1814
log
(18)
ODEE (X1111)
-Exx, mpp,, Ex Pieta |x4;
log
ODE, (Xten)) ] (Pt
+C,
(19)
where C' is a constant independent of 0. Thus, we can maximize the expectation:
Ex,, rp,, Exigy opet [xs, log
(Pi,1
ODE, (xx,+))
(20)
Locality of CNFs and estimating Pt.+s ODE,
ODE (Xt11)
With the law of iterated expectations, we can
write the probability Pt.s (xx..) ) in equation
20
ODE Ph44
(x(.)) Ex,, upc,
Npe,
ODE P
(x1,.1 (8c..]
(21)
15

---

Under review as a conference paper at ICLR 2023
We show in the following that can approximate the expectation in equation 21 with only neg- ligible error & > 0 on a different distribution Pt. (Xx,,,) that depends on x and K. Importantly, the support of Pt. (xx__.) is a bounded set with diam(supp(Pe, (Xt.4 ))) - 0 as 2 all X1 Ed Intuitively, the score Vx log Pr (x) is determined by the (local) probabilities Pa in
+ 0 for
small environment around x instead of the global properties of p:. This insight is an important part of our methodology and loss estimation.
To prove this, given X we first define set of points x for which the density Pt++1 greater than chosen i.e. we define
ODE, (Xe,43 (x) is
(*(x,..)) = xxE Rdd
Ph1+1
ODE
(Xr,41 (x) 2 x3)
(22)
Then, for this set, we consider the indicator function 1x x (-), which is for elements in (*(()(4))) and 0 otherwise. Then, we can rewrite equation
as
ODE, P (xt, 18x))
Ex,, Mpy
(23)
Ex., ups, - (x&,, P
ODE (Xt.41 (xx,.) + (1 1 x (x)))
ODE, (xL.+) [x).]
(24)
By the definition of the set St (xx441) ) we can derive the following bounds ODE.
Ex,, ~ps
(1 20,41 (Rx,, ))
(xt, [2x.)] S n
(25)
and
E Ex,; mp:
1 1 (St,, Pt:11
C.EEE (X1e41 (8e.) (1 I 1x (kr,))
(.P ODE, (Xec41 (x(.)
(26)
Ex, -pt,
1 1 (Rer Pemn CDDE, (Xt;11 (x11111211)
(27)
ODE,N Phil
(X1111)= - -K
(28)
For the next part, we need to make an approximation for Pt,41
ODEE (. (x,))
Approximating Pt,41 (-Ixe,) as Gaussian Since the CNF bijective, choosing Pt,4 P (
ODE
(-Ix..)
as Dirac delta function with spike located at (XDE(X(,) is the correct choice. Here, (ODE(X(.) is defined as the solution of the probability flow ODE equation (12) for x,, integrated from time to 4++1. However, solvers to compute MODE(X) (x anyway, Therefore, we make the assumption that given Xx,, solving the probability flow ODE equation 122 until t441 will give a solution x,DE that is approximately Xt, +(1+1 (P - 1 s0 (x (:)) and standard deviation GODE e(li+1 - 4.), for an arbitrary, but small >> 0. This approximation makes use of the explicit
are assuming that we are limited by machine precision and inexact ODE
Gaussian with mean (ODE(x;;)
Euler method and thus also relies on the time step 1, - tę being sufficiently small to ensure stability of the time integration.
Given the above approximation with Euler steps, we can derive the following bound on the distance between x and HODE(X) (x) using assumptions (i), (ii) and (iii)
| - MODE M2
||x, xr, (t141 - (P(t,) + 29, so(x,, f))
(29)
$ (t++1 - Pt),) + 2 (S,,, +))1l2)
(30)
5 (li41 4.) [11 + 18r,12 + 19 (1 + C.)
(31)
$ (ti+1 - Ct)]Coell |12
(32)
Morover, the approximation using Gaussians gives us the following equivalence for the term
Prope,, 3a (Xt,11 1& which we have used to define the set S* (x()); in equation (22)
P:ODE , (Xt (xx.) Z K.
(33)
(2m) exp (Xt.. - (XDDE(X,))) (Xi.. - (ODE((())) )) 2 K,
(34)
16

---

Under review as a conference paper at ICLR 2023
where Σ Tone.. The above is then equivalent to
| - ( )112 -log kiv (2T)
TODE
(35)
Note that JODE depends and (li+1 1i). We can now finally define the distribution Pr. (X4441)- ). For this, we define the set
8th (Xt=r))
{x R I - x112 S (t141 to)
(36)
XES XES* (xt1441
Now, by combining equation For the indicator function
(35 Hh
and equation 32
obtain diam (5* (x4,41 )) + O 1441
(.) on Stx (x )), we get the bound
1x (x) I 1 x (x)
Vx IR
(37)
and therefore also
Ex,, upe.
(1-
I Xe4t
(Zt, 3)
ODE,
(xe,+s (x()
(38)
Given S* (xt,)))) we define the distribution pr, (X4(1) based on Pt; but with support restricted to St (xx,,)) Then, we get the following equality
IN (
(
O A (Xr..: [8..) Ex,, (x,4)))) Cxe+: Pt+s
ODE, (x [Be.)) (39)
where Ox is normalizing constant depending on 6, ti+1 t and Xx+1-
Maximizing the variational lower bound tion 38), we obtain a lower bound equation 20
From Jensen's inequality, equation and equa-
Ex, -p, Ex pp (8c, [log
Pt.m
ODE,
(Xt: ]
(40)
Ex,, rp,, Ex opep (xc, log (Ex, (x
ODE, Cxe+i Ptot1
(xe,+s [str.)
(41)
2 Exi,ppi Ex Ma (xc, Taxe, mey, (Xe,
log
Cx
ODE, (Xiet1 (kte))
(42)
Exi, Ex,, op (xc, Bx,, me, (Xects)
log Pta11 (Xtesa (kt:)) + log(CCxeil
ODE
(43)
which is the same as maximizing
[log
ODE, (Xtaus (x
Ex,, Pe, Ex Peppe (x,, Ex,, pp,, (x1))))
Pte41
(44)
Thus, instead of the original objective equation 20 equation
we instead maximize the lower bound from
Deriving the 1.2 loss equivalent minimizing
Since Pt--1
ODE, ( (8x,) ) is the density of a Gaussian, the objective equation 44h
Ex,, <p: Exeep peet (x,, Pe Pe, (Xee44))
[[|xtey - MODE (xx))1//oooo
(45)
Ex,, mp: Exica pris bx,, Exs, ruBe, (Keeua))
[[xxt41 MODE MODE (1)
(46)
where
(ODE(Xt.) = &c, + (4441 - 4)) (((x))) 50 20000(x0 (x((,))4)
(47)
17

---

Under review as a conference paper at ICLR 2023
Training trajectory
X
X
XtM
Initialization
L2 loss
L2 loss
(ODE(Xe))
MODE (x1))
M MODE (xto)
( {GC
Sp X 40)
(t to)
P(x,))
Data flow
Cradient flow
Figure 8: Training overview for the trajectory (Xt,y Xt, X Gradients are backpropagated over multiple time steps via automatic differentiation. This requires that the physics operator P is differentiable. Incoming gradients at (x,, 42) are used to obtain gradients for 0, which are summed over all steps i. The network weights 0 are then updated based on the optimizer, e.g. stochastic gradient descent Adam.
Estimating the loss The training data set {(xt, ((=t,.) O I N is sampled directly from the emp induced by the train- SDE of the physical system equation 110 so the empirical distribution Pt, ing data set at time discretization
close the marginal distribution Pt, for 0 S sufficiently large N, We make use of this fact to approximate the sampling of Xr, ~ Pr, and
M and
X6+1 Pt,+1 Xt, in equation (46), by drawing data sample Xt, at time t; and its next successor X1,+1 on the trajectory at time f from the training data set.
We approximate sampling from the distribution Pt, (xt_.)) by reusing the sampled x,, Intuitively, for the CNF from & to f4+1. our optimization fits a deterministic (bijective) process defined by the probability flow to the non-deterministic SDE trajectories, which start the same point.
Extension to multiple time steps tory (Xt,,..,,)) with 0 < < S
We jointly train multiple time steps. First, we sample trajec- M from the training data set. Then, the loss becomes
|
(48)
k=1+1
where (xDE(x;)) is the discretized trajectory from the probability flow ODE, i.e.
VODE(xe)) Xt;
(49)
1 (oDe((x.) (xDE(X))) + (Ie- (x-1) ((orone(x..)) T.. so (HoDD (x.,) (xx,), (.))))))))))))))
(50)
For multiple steps, instead of reusing Xt, as a sample from pi, (Xt,11) ) we use the previous ODE solution p (xt.). ). An overview of the training for multiple steps is shown in Figure
Bidirectional training Since CNFs are bidirectional and our loss formulation does not make any specific assumptions about Pt, and Pt,,. we can consider training the reverse time direction. this case, the SDE describing the physical system becomes the reverse-time SDE, Analogous to the in- creasing time direction, we sample the trajectory (x.....,, Xxy ). but the time direction is reversed, ie.
> t' and tn = t's and f = fg- We train the combined objective by sampling trajectory from the training data set and alternating between optimizing the time-forward direction equation 448). and the corresponding time-backwardh direction loss. Note that training both the directions (i.e. po to pr and vice versa) has been explored in other works as well, for example for training Schrödinger Bridges (Bortoli et al. 2021 Chen et al. 2022). where the discrepancy between the final marginal distribution (i.e. either p? or ph SDI. for the time-forwart and respectively time- backward (or SDI
reverse-time) SDR is minimized.
18

---

Under review as a conference paper at ICLR 2023
Rollout length The training of SMDP requires unrolling algorithm shown in figure Addi- tionally, we adopt a training method based on sliding windows. For a window size S, we consider the points (Xt - .... ................................ xt ) from training trajectory and unroll the SMDP algorithm I for S steps. We compute the loss equation (48) and backpropagate gradients through all steps to obtain updates for 0. Then, move the window by 1, f.e. consider the points (X XtM -27 (tx---1) and compute the updates for O. We repeat the above M times until we have covered the entire trajec- tory. If the training data trajectories are short, use M, otherwise, pick lower window size S.
g
S 1, then we do not require differentiability of the physies operator P, as SMDP reduces to predicting the next point on the trajectory given the previous point.
Starting with an untrained score network, long rollouts may yield divergent trajectories because of the physics dynamics. Therefore losses may be very high and the training becomes unstable. We therefore typically start training with short sliding window, e.g. 2. We train for few epochs and then increase the sliding window size by constant. We repeat this until we reach a sufficiently high rollout length, which yields stable trajectories for the entire simultion.
We give details of this for each specific experiment either directly in the main text or the accompa- nying appendix.
Comparison to DiffFlow Zhang & Chen 2021) train DiffFlow by considering forward and back- ward processes in the context of generative modelling. In their setting, the drift f(x, () is also learnable neural network and they consider po be the data distribution and pr to resemble simple noise distribution. Specifically, they implement the forward and backward processes as
X x; + [i(x:)A + 9 VAL
(51)
x; = x - [f++(( - 1
$141 (1 + 9 V At
(52)
for two samples of 5 (0, 7) and time discretization (4 and A/; (li+1 4:)-
Zhang & Chen (002)) directly minimize the KL. divergence between the trajectory distribution for the forward and backward process, i.e. they minimize
(l(pF(7)llen(T)) Erapp (r (x))]] + REr-p [logpn(x))]
(53)
N-1
log p=(xx(x;-1) PB 178 xx))
Σ
ETapF
(54)
Using the forward and backward process discretizations well as the fact that pa (xx) is Gaussian distribution, they are able to derive a loss based on the squared difference between the forward and backward process as well as an additional likelihood term for log (B(xx). see equation (15) in their paper.
Our method on the other hand directly minimizes the KL divergence between marginal distribution Pt and the ones produced by deterministic probability flow ODE PODE, Thus, our loss likewise minimizes the difference between a forward and backward trajectory. Moreover, in our case, Pr not constrained to simple noise distribution.
19

---

Under review as a conference paper at ICLR 2023
B ADDITIONAL EXPERIMENT: ID PROCESS
We discuss an additional experiment, where we compare SMDP with Implicit Score Matching (ISM) HHväärinen 2005). For this task, we consider the SDE given by
dæ =- -2x1 sign(z)a* dt + A2dw.
(55)
The corresponding reverse-time SDE is given by dr [A1 sign(r)** X V logPe(x)) di A2dw.
(56)
Throughout this experiment, Po is categorical distribution, where we draw either 1 or -1 with the same probability. In figure / we show trajectories from this SDE simulated with the Euler- Maruyama method. Trajectories either start at 1 -1 and approach trajectory value at t = 10, it is no longer possible to infer the origin of the trajectory at =tt=0,
increases. Given the
We employ neural network sp(2,,t) t) parameterized by 0 to approximate the score via ISM and compare it to SMDP. In both cases, the neural network is a simple multilayer perceptron with tanh activations and hidden layers with 20 neurons for the first four hidden layers and 10 neurons for the last hidden layer.
Our training data set consists of 250 simulated trajectories from 0 until & 10 and At 0.02. Therefore each training trajectory has length M = 500.
Implicit Score Matching Implicit Score Matching (Hyvärinen] 2005) is a score matching method that leverages the fact that for random vector E [ga with probability density function p, mini- mizing the score matching objective
J(0) 2 Exup [x]s)(x) Vx logp(x)]
(57)
is equivalent to minimizing the following objective
(x);;
J'(0) = Exxu
sx())
(58)
Ox.
Note that for ISM, there is no explicit time dimension, so we are absorbing the time dimension into x, i.e. for the trajectory (2'1 F2,
time with x(9) = (I,, .2,,.... sampled at time (41,
R2. We collect x; for 42, ..., 1 we concatenate value and 1.... M from all trajectories in a new training data set. When training sx()), we therefore lose the information, to which trajectory value-time pair originally belonged. We obtain the time-dependent score Vx, logPx, Pxg (xi) from the first coordinate of the output of se (xx, ie. sa(x)1.
We compute the partial derivative (x) /0x; using reverse- automatic differentiation in JAX (jxxjacrev)) We train s8 (x) with the Adam optimizer for 15.000 epochs with learning rate 10 which we decrease by a factor 0.1 every 5.000 epochs.
SMDP Training The training SMDP follows the experiments in section 4 with slight problem- specific modifications. Since the length of the trajectories is very long (M 500), we subsample the trajectories, and keep every 5th point in time. We train with a sliding window of size 4, which we increase every 500 epochs by 2 until we reach an window size of 40 with the Adam optimizer and learning rate 10-5, Then, we finetune the network and train on the complete trajectories for an additional 500 epochs and sliding window size 70 with learning rate 10
Comparison We show a direct comparison of the learned score, the reverse-time SDE trajectories and the probability flow trajectories between ISM and SMDP in figure 10 The learned score of ISM and SMDP in figure 10a and figure 10b) is similar until & 2. Then, the score for both networks becomes positive for points marginally above O, i.e the score pushes points up, leading them to at t = O. Analogously, points marginally below 0 are pushed down. For SMDP, the absolute value of the score in this region is considerably higher than ISM, thus having a more significant shearing effect.
20

---

Under review as a conference paper at ICLR 2023
This affects the reverse-time SDE trajectories. Because of the stronger shearing, SMDP trajectories result in states either 1 or -1, see figure Tod On the other hand, for ISM in figure T0c where the shearing is less pronounced, many trajectories end up in-between and -1. These states not valid samples from the posterior distribution of the solution from equation (55)
Interestingly, the probability flow solutions for both SMDP and ISM have comparable quality, cf. figure 10e and figure Tor In both cases, trajectories that start above at t 10 end close to at t = 0, and vice versa.
Overall, this case illustrates that the learned score of SMDP is more accurate than the one of ISM, because of its stronger shearing. This results in better quality of the reverse-time SDE trajectories. SMDP performs better here because it is directly trained with an integration of the physics dynam- ics, whereas ISM purely data-driven and does incorporate the physics model with its temporal evolutions at training time.
21

---

Under review as a conference paper at ICLR 2023
1.0
1.0
0.5
0.5
0.0
0.0
10
10
(a) A2
0
(b) )2 0.03
Figure 9: Trajectories from SDE equation
55
with A2 0 (a) and A2 0.03 (b).
75
75
0.5
0.5
4
0.0
0.0
-0.5
25
-0.5
25
10
-75
10
(a) ISM learned score.
(b) SMDP learned
1.0
-05
-05
-1.0
-1.0
(c) ISM reverse-time SDE trajectories.
(d) SMDP reverse-time SDE trajectories.
1.0
1.0
0.5
0.5
8 0.0
4 0.0
10
10
(e) ISM probability flow trajectories.
(f) SMDP probability flow trajectories.
Figure 10: Comparison of Implicit Score Matching (ISM, left) and Score Matching via Differentiable Physics
(SMDP, right). Colormap in (a) and (b) truncated to [-75, 75].
22

---

Under review as a conference paper at ICLR 2023
10-1
[0-l
10"1
Reconstuctoon MSE
Recandtrccton MSE
Recomsrrustenn Resonstrutioon MSE
10-2
Reeoeet cudion
10-3
10"
10'7
1on5
20-1
10"
1071
10"
longth
badipropapation stops
Dackpropagation steps
(a) SMDP-ODE
(b) SMDP-SDE
Figure 11: Heat equation reconstruction MSE for different sliding window sizes S (Rollout length) and number gradient backpropagation steps for SMDP-ODE (a) and SMDP-SDE (b).
C ABLATION FOR GRADIENT BACKPROPAGATION AND ROLLOUT LENGTH
Gradient backpropagation steps is possible to unroll for steps, but stop the gradient back- propagation after G < S steps. For example, when unrolling the algorithm for S = 32 steps, can unroll the first G 8 steps, save the intermediate results, then backpropagate the gradient to optimize 0. Then, continue unrolling with the intermediate results new initialization, again stopping and updating 0 after G steps. For this example, we can repeat this S/G = 4 times in total until we obtain combined rollout of S steps.
An advantage of stopping the gradient after G steps is that it reduces memory requirements, we do need store all intermediate results and therefore this allows for training with large window sizes.
For all other experiments in the paper, we have
S always.
Effect reconstruction MSE We evaluate the effect of changing the window size $ and the number of steps until we stop the gradient backpropagation using the setup of the heat equation experiment in section 4.1 In figure [1] show the results the reconstruction MSE for SMDP- ODE and SMDP-SDE. We keep the time discretization fixed, i.e. full simulation trajectory consists of 32 steps from always backpropagate gradients through all unrolling steps. In the second case, where we change the number of gradient backpropagation steps, we keep the window size S fixed at 32, which is the entire simulation trajectory. For both SMDP-ODE and SMDP-SDE. our evaluation shows that both longer rollouts and more gradient backpropagation steps improve the reconstruction MSE. The improvements are more significant for SMDP-DDE.
0 until L 0.2, however we vary the window size S. In this case,
23

---

Under review as a conference paper at ICLR 2023
D ID HEAT EQUATION: EVALUATION OF POSTERIOR
To demonstrate the diversity and high quality of the reverse-time SDE trajectories we provide an additional evaluation of the posterior for one-dimensional heat equation case. We compare the SDE solutions with samples obtained from filtering a large data set. To simplify the comparison, we consider 1D processes based on the 2D Gaussian random fields from section 4.1 Each process has a length of 100 and corresponds to Gaussian random field of size x 100. Analogously to section
we use the heat equation simulate the states forward in time from 0.0 until 0.2.
Figure 12 shows some examples of the ID processes we consider here.
The network (e(x, t) representing the score is trained as described for the heat equation experiment, section 4.1] using the same ResNet architecture with padding in y-direction removed. We consider a time discretization with A/ 0.01. We begin training with an initial rollout of 6 steps and increase the rollout length every 5 epochs by 2 until reach 14 rollout steps. For every epoch, we train on 20 randomly generated 1D processes. We use the Adam optimizer with learning rate 10-4
Training of SMDP
In the following, we randomly generate a ID process P. We describe how we generate samples for t 0.0 conditioned the simulation end state of P at 0.2.
Reverse-SDE posterior We initialize the state based on P at f. = 0.2. Then, we simulate the reverse-time SDE with the learned score se (x, t) via Suler-Maruyama steps. A visualization 100 samples is shown in figure 14
Empirical distribution We sample 106 processes and form pairs of initial state = 0 and end states / 0.2. We filter the 100 process end states that closest the end state of in terms of the L2 distance. As solutions from the empirical sampling, we consider the 100 corresponding initial states as shown in figure 13 This empirical distribution makes qualitative comparison possible, They indicate that the reverse-time SDE solutions are diverse while matching very well. Simulating the obtained solutions forward in time gives end state that are in excellent agreement with
as shown in figures 14 and 13
(a) 1 = 0.0
(b) = 0.2
Figure 12: Examples of training data for 1D heat equation with initial states 0.0 (a) and end states
0.2 (b).
24

---

Under review as a conference paper at ICLR 2023
(a) 0.2
0.0
Figure 13: Empirical distribution: we generate 10 Gaussian processes and simulate them forward in time from 0.0 until f = 0.2. We pick one specific process P and sort all other processes in ascending order based on the distance to time 0.2. Then we pick 100 and visualize them time 0.0 (a) and 0.2 (b). The top row both plots shows the process
(a) 0.0
(b) 0.2
Then, we simulate 100 trajectorics from the reverse-time SDE with learned score and P at = 0.2 as initialization. We sort the states at = 0.0 based their distance to P and visualize them (a). We then simulate all states from forward in time again until i 0.2, see (b). The forward simulated trajectories almost exactly match P 0.2.
Figure 14: Reverse-time SDE: We pick the process from figure 13
25

---

Under review as a conference paper at ICLR 2023
E ARCHITECTURES
ResNet We employ a simple ResNet-like architecture, which is used in Section 4.1 for the score function (x, and the convolutional neural network baseline (ResNet-S and ResNet-P) well as in Section 4.3 again for the score sø (x(, t).
Since in both experiments, there are periodic boundary conditions, we apply periodic padding with length 16, i.c. the underlying 2-dimensional data dimensions are N N, the dimenions after (N + 16). We implement the periodic padding by first tiling the
the periodic padding are (N++1 + input times in and y-direction and then cropping to the correct sizes. The time is concatenated as an additional constant channel to the 2-dimensional input data when this architecture is used to represent the score Sø sx, t).
The encoder-part our network begins with a single 2D-convolution encoding layer with 32 filters, kernel size 4 and activation function. This is followed by 4 consecutive residual blocks, each consisting of 2D-convolution, LeakyReLU, 2D-convolution and Leaky ReLU, All 2D convolutions have 32 filters with kernel size 4 and stride 1. The encoder part ends with single 2D convolution with filter, kernel size and no activation. Then, in the decoder part, we begin with a transposed 2D) convolution, 32 filters, kernel size 4. Afterwards, there 4 consecutive residual blocks, analogous to the encoder residual blocks, but with the 2D convolution replaced by a transposed 2D convolution. Finally there is final 2D convolution with 1 filter and kernel size 5, Parameters statisties of this model. as well as the others are given in table
UNet We use the UNet architecture with spatial dropout as given in (Mueller et al., 2022), Ap- pendix A.I. 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 For all experiments, we consider the FNO-2D architecture introduced in (Li al. 2020 with kmax.j = 12 Fourier modes per channel.
Dil-ResNet The Dil-ResNet architecture described in (Stachenfeld et al. 2021). Appendix A. Since this architecture represents the score in Section 4.2 we concatenate the constant time channel analogously to the ResNet architecture. Additionally, positional information is added the network input by encoding the x-position and y-position inside the domain in two separate channels.
Architecture Parameters
ResNet
330.754
UNet
706.035
DilatedConv
336.915
FNO
465.377
Table 4: Summary of architectures.
26

---

Under review as a conference paper at ICLR 2023
F HEAT EQUATION
F.1 ADDITIONAL TRAINING DETAILS
Spectral loss
We consider a spectral error based on the two-dimensional power spectral density. For two 2d-images, we compute their 2d Fourier transform and compute the radially averaged power spectrum $1 and 82. Then, we define the spectral error as the difference between the log of the spectral densities
(88,, 82) = |log(>1) log(s2)]-
(59)
SMDP-ODE For inference, we consider the linear time discretization in n with AI 0.2/32 and la2 0.2. During training, we sample random time discretization i, for each batch based on in by fn ~ Utn, 2t/2, in + At/2) for n 1, .....31 to not overfit on the step size. In the warmup phase of training, we unroll A for N 6, 8.... 32 steps, where we increase N every 2 epochs. We employ Adam to update the weights with learning rate 10-4 After the warmup is finished, we finetune the network weights for 80 epochs with an initial learning rate of 10-4 which we reduce by factor of 0.5 every 20 epochs.
Diffusion model The diffusion model is inspired by (Rissanen et al. 2022 We make use of the same ResNet architecture as SMDP-ODE, however keep the time discretization in fixed. The diffusion model does not include the physics operator P in the rollout. Therefore, it has to learn the score and physics the same time. Except for those changes, the training setup is identical to SMDP-ODE.
Baseline methods All other baseline methods are trained for 80 epochs using the Adam optimizer algorithm with initial learning rate of 10 which is decreased by factor of 0.5 every 20 epochs. For the training data, we consider solutions to the heat equation consisting of initial state Xy and end state Xy that are noise-free and add Gaussian noise with standard deviation 0.1 to the network input.
F.2 LOGARITIIMIC TIME DISCRETIZATION
For this experiment, we consider a logarithmic time discretization during training and inference. This discretization is finer around t = O, where most of the small-scale structures are smoothened out. We find that this method also works well, but the linear time discretization performs better for most cases, especially when AL is changed for inference, cf. figure and figure 15
10-1
10-1
SMDP-SDE
SMDP-SDE
SMDP-ODE
SMDP-ODE SMDP-ODE
10
SMDP-ODE
10
erre pundds
10 sauu eapsod
10-3
10-3
MSB
10e
MSE
10
10-4-
10
10-5
10-6
x0.5
x0.25
x0.5
Inference time steps
Inference time steps
(a) Regular training
(b) Bidirectional training
Figure 15: Logarithmic time discretization. Reconstruction MSE and spectral errors for varying time steps of regular variant (a) and bidirectional variant (b). Large errors are truncated the top each graph.
27

---

Under review as a conference paper at ICLR 2023
F.3 PHYSICS-CONDITIONE SCORE
We extend the definition of the score function, to also include information about the physics at time i.e. Pxx.). Thus, replace sx(x, in algorithm by s0([x, P(x)]: l). where concatenate both inputs. An evaluation is shown in figure Although the stability of SMDP-ODE and SMDP- ODE noise is greatly increased, we do not obtain results with low spectral error for SMDP-SDE. Interestingly, the bidirectional training also seems harmful in this case.
10
10
SMDP-ODE
SMDP-ODE
10-2
SMDP-ODE noise
10"2
SMDP-ODE + nOise
10 3
MSE
100
MSE
-10
10
10
X
10-1
10-1
10-6
10-
x0.25
x0.5
x1
x4
x8
x16
x0.25
x0.5
3.4
x16
Inference
steps
Inference steps
(a) Regular training
(b) Bidirectional training
Figure 16: Physics-comee conditioned score. Reconstruction MSE and spectral errors for varying time steps of regular variant (a) and bidirectional variant (h). Large errors are truncated at the top of each graph.
F,4 HIGHER TRAINING NOISE
We increase the diffusion coefficient from 0.1 to 1.0 during training, but for inference, we set it back to the lower value. An evaluation is shown figure 17 Overall, SMDP-SDE produces results with low spectral eиror and SMDP-ODE has much lower MSE, similar to the standard case. However, overall, the MSE is much higher, indicating shortcomings of this approach.
10
10
SMDP-SDE SMDR-ODE SMDP-ODE
SMDP-SDE
X
SMDP-ODE
10-2
10-2
SMDP-ODE
X
X
X
run papods
100 Speciil
10
10
MSE
10
MSE
10-
10"
10-5
10-5
10
Inference time steps
Inference time steps
(a) Regular training
(b) Bidirectional training
Figure 17: Higher training noise. Reconstruction MSE and spectral errors for varying time steps of regular variant (a) and bidirectional variant (b). Large errors are truncated top each graph.
F.5 ADDITIONAL RESULTS
28

---

Under review as a conference paper at ICLR 2023
(a) Ground truth
(b) Input
(c) SMDP-ODE
(d) ResNet-S
(c) ResNet-P
(0) FNO-P
(g) FNO-S
(h) Diffusion
(i) Diffusion noise
(j) SMDP-SDE
(k) BNN-S
Figure 18: Predictions of different methods for the heat equation problem (example 1 of 2).
(a) Ground truth
(b) Input
(c) SMDP-ODE
(d) ResNet-S
(c) ResNet-P
(0) FNO-P
(g) FNO-S
(h) Diffusion
(i) Diffusion noise
(i) SMDP-SDE
(k) BNN-S
Figure 19; Predictions of different methods for the heat equation problem (example 2 of 2). Neither the BNN nor the Dife model are able to produce small-sce scale structures.
29

---

Under review as a conference paper at ICLR 2023
Ground truth
SMDP-SDE
Figure 20: Comparison of SMDP-SDE predictions and ground truth for buoyancy-driven flow at t 0.36.
G BUOYANCY-DRIVEN FLOW WITH OBSTACLES
G.1 TRAINING
We train all networks with Adam and learning rate X 10 with batch size 16. We begin with unrolling N 2 steps, which increase every 30 epochs by 2 until we reach N 20.
G.2 ADDITIONAL RESULTS
We give more detailed time evolutions of results for the buoyancy flow case in figure 21 and figure 22 These again highlight the difficulties of the physics simulator to recover the initial states by itself. The SMDP variants significantly improve upon this behavior.
In figure 20 we also show an example of the posterior sampling for this case. becomes apparent that the inferred small-scale structures of the different samples change. However, in contrast cases like the heat diffusion example, the physics simulation this scenario leaves only little room for substantial changes of the states.
30

---

Under review as a conference paper at ICLR 2023
Ground truth
SMDP .DE
SMDP SDE
Phy only
0.65
1= 0.55
0.45
0.35
0.25
= 0.15
(a) Marker field
Ground Iruth
SMDP ODE
SMDP SDE
Physics only
0.65
0.55
0.45
0.35
0.25
(b) Velocity field (æ-direction)
Figure 21: Predictions for buoyancy- -rivvvn flow with obstacles (example I of 2).
31

---

Under review as a conference paper at ICLR 2023
Ground truth
SMDP -DD
SMDP SDE
Physics only
0.65
== 0.55
0.45
0.35
t= 0.25
= 0.15
(a) Marker field
Ground truth
SMDP ODE
SMDP SDE
Physies only
0.65
0.55
0.45
0.35
(b) Velocity field (æ-direction)
Figure 22: Predictions for buoyancy- -rivvvn flow with obstacles (example 2 of 2).
32

---

Under review as a conference paper at ICLR 2023
Ground truth
Learned physics
100 steps
200 steps
300 steps
steps
steps
steps
1000 steps
steps
Figure 23: Steps of Langevin dynamics for e x 10
Ground truth
Learned physics
steps
200 steps
300 steps
400 steps
500 steps
750 steps
1000 sleps
2000 steps
Figure 24: Steps with Langevin dynamics for x 10-5 and an additional step with At sa(xx () which smoothes the images.
H ISOTROPIC TURBULENCE
For the learned physics network, we employ FNO neural network 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 factor of 0.5. We train SMDP-ODE with the ResNet architecture for 250 epochs with learning rate 10-4, decreased every 50 epochs by factor of 0.5 and batch size 6.
H.1 REFINEMENT WITH LANGEVIN DYNAMICS
We do a fixed point iteration at a single point in time via:
x x + Vx lo 22ez,
(60)
for number of steps T and 2 10-5 post-processing and refinement strategy, ef. fig-
ure 23 and figure 24 This is motivated by established methods in score-based generative modelling Welling & Teh 2011 Song & Ermon 2019). For prior distribution Ti, x म and by iterating
equation (60), the distribution of x equals P: for e - and T - oc. There are some theoretical caveats, i.e. an additional Metropolis- Hastings update needs be added in equation (60) and there are regularity conditions (Song & Ermon 2019).
33

---

Under review as a conference paper at ICLR 2023
Ground truth
Leamed physics
SMDP ODE
SMDP33D S0S
t= 10
t=
t=
t=6
t=4
Figure 25: Predictions for isotropic turbulence (example 1 of 2).
34

---

Under review as a conference paper at ICLR 2023
Ground truth
Leamed physics
SMDP ODE
SMDP SDD SDE
t= 10
t=
t= 8
t=6
t=4
Figure 26: Predictions for isotropic turbulence (example 2 of 2).
35


Reviews:
[{"review_id": "vDh12ceHIH", "reviewer": "Reviewer_N4rm", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper uses score matching SDEs for learning inverse problems with physical simulators. Different from common diffusion models, the forward SDE is changed: the drift is given by a physical model and the diffusion term does not grow to infinity as time increases. Training and inference algorithms based on SDEs and probabilistic ODEs are discussed. Empirical results are evaluated on several inverse problems and compared against baselines, demonstrating the empirical advantage of the proposed SMDP method.", "summary_of_the_review": "While this seems to be an interesting application of diffusion normalizing flows to physical inverse problems, I am a bit concerned by the clarity in the method section. There are some notation issues, and even the loss function is not precisely defined (either in the main paper or in the appendix). But I do think the method itself is valid and experiments generally support the validity of the method.", "strength_and_weaknesses": "Strength:\n- Empirical performance seems to be better compared to baselines.\n- The idea is an interesting take at inverse problem solving using the idea of diffusion models.\n\nWeakness:\n\nThe method section is not very clear. \n- For example, in $\\{x_t^n\\} \\subset R^{d \\times T}$ is confusing. What is $d$? With notation $d \\times T$ does this imply that the trajectories are sampled at discrete times? The $t_{m}$ notation should come after you discuss Euler-Maruyama discretization steps?\n- What is the loss function, and are there any principled justification to this? My understanding is that this is the loss function in (Zhang & Chen, 2021) but with the drift being non-trainable. If this is true, then it would help to make this explicit, as it helps understanding the paper by a great deal. While the technical novelty is lowered, this is still a non-trivial application of diffusion normalizing flows (since Zhang & Chen mostly considered generative modeling problems).\n\nIt is unclear why both ODE and SDE can be used. \n- The goal in inverse problems is to sample from the full posterior. In that sense, ODE will only give one solution, which as the paper suggests, does not cover many candidate solutions. This is also because ODE does not recover the measure of the entire path, just the marginals. Wouldn't this suggest that, assuming the forward model is correct, the SDE is the **only** correct way of sampling from the posterior, since it should match the path measures.\n\nEvaluations are not precisely targeted towards posterior sampling.\n- If my understanding is correct, then in Table 1, only one path is compared against. However, if we have a non-zero g, then there are naturally many paths with similarly high probability of occurring, and comparing against one path is going to place posterior sampling at a disadvantage. I wonder if it is possible to obtain a \"ground truth\" posterior and compare path distributions with that? For example, using metrics like maximum mean discrepancy or divergences based on KDEs.\n- The paper address the above by having g=0 during evaluation. However, would this generate the same score function as in g=0.1 during training?\n\n\nNot weakness, but curious questions:\n- Why does the FNO perform so badly in Table 1?", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "W-5kgphD5dO", "reviewer": "Reviewer_37Sw", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "In this paper, the authors adapt the formalism of denoising diffusion models to learn physics trajectories. In particular, they consider a forward diffusion given by a differentiable physics operator. The noise of the forward diffusion is assumed to be inherent noise or measurement error. Once this is done, the authors learn the backward dynamics by backpropagating through the dynamics using the square distance between the observed forward point and generated point. They consider both Ordinary Differential Equation (ODE) and Stochastic Differential Equation (SDE) settings. They apply their method to 3 problems: a toy heat equation problem, a buoyancy-driven flow with obstacles and an isotropic turbulence problem. The method outperforms other neural network approaches to learn dynamics.", "summary_of_the_review": "I think this paper represents a nice application of diffusion models to physics.\nApart from the training procedure which seems strange I think the methodology and experimental settings are valid.\nI think this paper opens the door to more fruitful interactions between the diffusion models and physical systems.", "strength_and_weaknesses": "STRENGTHS:\n* I think the connection of differentiable physics and denoising diffusion models is interesting and is a nice application of this generative model framework to physics applications. I believe that this paper is a step towards the use of generative model for simulating trajectories in physics.\n* The experimental section is quite strong. I especially enjoyed some of the findings regarding the differences between SDE and ODE which I think are both intriguing and interesting.\n\nWEAKNESSES:\n* For me the main weakness of the paper is how the training is conducted. By using the loss function the authors use in the paper they lose the flexibility of diffusion models which do not require unrolling the dynamics. If I'm right I think the authors did not apply classical diffusion models here because they did not have access to the transition kernel of the forward dynamics (since it depends on the differentiable physics operator). However there are ways to circumvent this issue by looking at the Implicit Score Matching loss, see for instance [1] (and [2] for an older reference in the score matching literature). If one does not want to use that loss then one can leverage the fact that the Euler-Maruyama discretization of the forward dynamics yield Gaussian steps (locally), see [3]. Both these approaches allow for efficient training of diffusion models even in the case where the forward dynamics is given by a differentiable physics operator. These losses are more computationally friendly and yield good results in practice\n\n* Regarding the bidirectional optimization, recent papers on Schrodinger Bridges using diffusion models [3,4] rely on some bidirectional models. I think the authors should discuss the link between this approach and these papers. Finally, I think that the authors missed the opportunity to discuss [5] which advocates for dropping the noise in diffusion models. I think that this is not always beneficial, dropping the noise makes sense here, as the authors justify the presence of the noise by saying that it \"is added to the data at each time step [and] can be regarded as either noise to the physical problem or as noise from a measurement process\". Gaussian noise with such covariance seems rather arbitrary and seems to be present only for the convenience of having well-defined probability measures.\n\nOTHER COMMENTS:\n* In the reconstruction loss in the first experiment (Section 4.1). I don't really understand how this can work. If $T$ is large then the operator $P$ destroys all information and converges to a stationary state. Then how can we check that we recover $x_T$ (or are all $x_T$ the same in that case)?\n* I would have appreciated more details on the bidirectional method.\n* I also would have appreciated more explanations as of why SDE does not perform well in the buoyancy driven flow example.\n* I am not competent enough to assess if the baselines chosen by the authors are relevant and will not comment on them.\n\n[1] A Variational Perspective on Diffusion-Based Generative Models and Score Matching - Huang, Courville\n[2] Estimation of Non-Normalized Statistical Models by Score Matching - Hyvarinen\n[3] Diffusion Schrödinger Bridge with Applications to Score-Based Generative Modeling - De Bortoli, Thornton, Heng, Doucet\n[4] Likelihood Training of Schrödinger Bridge using Forward-Backward SDEs Theory - Chen, Liu, Theodorou\n[5] Cold Diffusion: Inverting Arbitrary Image Transforms Without Noise - Bansal, Borgnia, Chu, Li, Kazemi, Huang, Goldblum, Goldstein", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "OO0Oc4dSso", "reviewer": "Reviewer_wxSJ", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper describes an ODE-based model for inverting the time-evolution of a physical system, aiming to predict a distribution of initial states which could have given rise to this system. Their approach assumes access to a differentiable simulator of the physical system, and represents the reverse-time evolution of the system as either a PDE or an ODE that combines the simulator with a neural network force term. The paper proposes training this force term such that the reversed trajectories under the ODE are close (under L2 norm) to true (forward) trajectories generated by the physical system running in the normal direction. The authors present results for a few physical processes, including heat diffusion, buoyancy-driven flow, and isotropic turbulence, and find that their method yields good approximations of the initial state.\n\n~~Although this approach is motivated by score matching and denoising diffusion generative models, and the authors spend considerable time explaining the background of these models, the actual proposed approach deviates in a number of ways from proper SDE-based diffusion approaches. It seems to me that this is not really a diffusion model in the sense used by the generative model community, and the probability-flow ODE appears to be used \"incorrectly\" from the perspective of score matching and distribution fitting. I think the proposed methodology would be better understood as a (empirically successful) parameterization and denoising-based training strategy for an inverse model, without reference to generative diffusion models or probability flow.~~\n\n*Updated summary:* After discussion with the authors, it does seem like there is a connection to proper SDE-based diffusion, but this connection is not very well explained. In particular, while their training objective isn't exactly a maximum-likelihood objective for the SDE itself or for the probability-flow ODE, it can be seen as a maximum-likelihood objective for a different, perturbed ODE flow model. And if the \"sliding window\" used during training is sufficiently small, the optimum of their objective will approximate the score function of the original SDE/ODE. I still feel that these connections aren't explained very well in the current revision of the paper.", "summary_of_the_review": "This paper has some interesting ideas for using learned ODEs to find the initial states of a physical system conditioned on a final state, and the method seems to work well empirically. However, I think in its current form the paper makes many unjustified and incorrect claims about the relationship between this model and SDE-based score-matching diffusion generative models. The paper also omits many important details about the training procedure, which seems substantially different than prior work with diffusion models.\n\n~~As such I do not think the paper should be accepted in its current state. I would be willing to raise my score if the authors clearly explain what their actual training algorithm is and remove the misleading claims about the relationship to score matching and probabilistic inference (or, alternatively, if they can justify the sense in which their objective is equivalent to score matching and explain why my concerns are unfounded).~~\n\n*Updated review:* The authors have added the missing details about the training procedure, which is now much better explained. They have also added a derivation which claims their objective is a probabilistic lower bound that is (indirectly) connected to score matching. I have some issues with the way this is presented in the paper, but after discussion with the authors (see below) I've been convinced that there is indeed a sense in which the actual objective used (for a sufficiently small sliding window) could be viewed as approximately doing score matching. However, this connection and the dependence on the window size are not very well explained in the paper, and there are still some misleading claims about maximum-likelihood training. I have raised my score from 3 to 5, since I no longer think there are critical problems, but still I'm hesitant to recommend acceptance due to these remaining clarity issues.", "strength_and_weaknesses": "Strengths:\n- Using a learned model to guide a differentiable simulator is a good idea, and using noise to train it seems reasonable as well.\n- Their proposed model does well on a variety of experiments, and in particular does better than either directly learning everything with a neural network, or simply running the simulator backwards.\n\nWeaknesses:\n- ~~The paper gives very few details about how the model is actually trained. Many details are missing regarding the actual loss used, the distribution of samples in the training set, what is fixed and what is learned, etc. The authors briefly mention \"bidirectional training\" but it is never described what this actually means.~~\n- ~~Based on my understanding of the training procedure, it is not actually trained via score matching, the typical way generative diffusion models are trained. This is why it is necessary to differentiate through the physics; I do not believe such differentiation would be necessary using an actual score-matching objective.~~\n  - *After discussion with the authors, it does seem like their objective can be viewed as an approximation of score matching. However, this approximation seems to depend on the \"sliding window\" used during training being small, and this requirement isn't discussed very clearly in the paper (although there's some discussion in the appendix).*\n- ~~The authors appear to be mis-using the relationship between the diffusion SDE and probability flow ODE. In particular, running the probability flow ODE should *not* be interpreted as producing a trajectory sample. The probability flow ODE is a continuous flow such that the marginal distribution at each time is correct, but any two points along a given flow path are not necessarily paths from the same trajectory. However, the authors discuss using the probability flow ODE to generate a \"solution\" trajectory, and appear to train their model so that this trajectory is close (in L2) distance to entire sampled trajectories. I similarly do not think it is appropriate to treat samples from the SDE as being drawn from a posterior, given that the training process does not seem to be based on KL divergence (or any other mechanism for fitting a sensible posterior over trajectories).~~\n  - *After discussion with the authors, it seems like the L2 distance objective can be seen as a maximum-likelihood for a different probabilistic model (not clearly described in the paper right now, but only mentioned in the appendix). And for a small enough sliding window, training this second model would require learning the correct score function, which could then be used to draw samples from the reverse SDE despite their objective not being a KL divergence with the reverse SDE..*\n- ~~Overall there are multiple choices that seem very questionable from the perspective of probabilistic generative modeling.~~\n- After discussion with the authors, the main remaining weakness I see is that the connections to score matching are not very well explained, and that some details which are presented as minor computational efficiency improvements are actually necessary modeling choices that enable their model to approximate the true score function.\n  - The training process actually uses a small window around parts of the trajectory. This is briefly mentioned as an implementation detail in the new \"Training and Inference\" section, but after reading the appendix and discussing this with the authors, it seems that the small window is actually necessary for their model to recover the true score function. If this window is not used, the learned \"score\" function might differ substantially from the true score, and the connection to the real SDE/ODE might break down.\n  - The authors refer to their objective as a maximum-likelihood bound for the probability flow ODE, but after discussion it appears that the objective is actually a maximum-likelihood bound for a *different* probabilistic model, involving additional Gaussian noise perturbations (which the authors motivate as modeling the numerical inaccuracy of the ODE solver). This perturbed model is not explicitly described in the current submission, except for a partial description in the appendix.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject