SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Benjamin Holzschuh  $^{1,2*}$  Simona Vegetti  $^{2}$  Nils Thuerey  $^{1}$
$^{1}$ Technical University of Munich, 85748 Garching, Germany  $^{2}$ Max Planck Institute for Astrophysics, 85748 Garching, Germany
We propose to solve inverse problems involving the temporal evolution of physics systems by leveraging recent advances from diffusion models. Our method moves the system's current state backward in time step by step by combining an approximate inverse physics simulator and a learned correction function. A central insight of our work is that training the learned correction with a single-step loss is equivalent to a score matching objective, while recursively predicting longer parts of the trajectory during training relates to maximum likelihood training of a corresponding probability flow. We highlight the advantages of our algorithm compared to standard denoising score matching and implicit score matching, as well as fully learned baselines for a wide range of inverse physics problems. The resulting inverse solver has excellent accuracy and temporal stability and, in contrast to other learned inverse solvers, allows for sampling the posterior of the solutions. Code and experiments are available at https://github.com/tum-pbs/SMDP.
Many physical systems are time-reversible on a microscopic scale. For example, a continuous material can be represented by a collection of interacting particles [Gur82; BLL02] based on which we can predict future states. We can also compute earlier states, meaning we can evolve the simulation backward in time [Mar+96]. When taking a macroscopic perspective, we only know average quantities within specific regions [Far93], which constitutes a loss of information, and as a consequence, time is no longer reversible. In the following, we target inverse problems to reconstruct the distribution of initial macroscopic states for a given end state. This problem is genuinely tough [ZDG96; Góm+18; Del+18; LP22], and existing methods lack tractable approaches to represent and sample the distribution of states.
Our method builds on recent advances from the field of diffusion-based approaches [Soh+15; HJA20; Son+21b]: Data samples  $\mathbf{x} \in \mathbb{R}^D$  are gradually corrupted into Gaussian white noise via a stochastic differential equation (SDE)  $d\mathbf{x} = f(\mathbf{x}, t)dt + g(t)dW$ , where the deterministic component of the SDE  $f: \mathbb{R}^D \times \mathbb{R}_{\geq 0} \to \mathbb{R}^D$  is called drift and the coefficient of the  $D$ -dimensional Brownian motion  $W$  denoted by  $g: \mathbb{R}_{\geq 0} \to \mathbb{R}_{\geq 0}$  is called diffusion. If the score  $\nabla_{\mathbf{x}} \log p_t(\mathbf{x})$  of the data distribution  $p_t(\mathbf{x})$  of corrupted samples at time  $t$  is known, then the dynamics of the SDE can be reversed in time, allowing for the sampling of data from noise. Diffusion models are trained to approximate the score with a neural network  $s_\theta$ , which can then be used as a plug-in estimate for the reverse-time SDE.
However, in our physics-based approach, we consider an SDE that describes the physics system as  $d\mathbf{x} = \mathcal{P}(\mathbf{x})dt + g(t)dW$ , where  $\mathcal{P}:\mathbb{R}^D\to \mathbb{R}^D$  is a physics simulator that replaces the drift term of diffusion models. Instead of transforming the data distribution to noise, we transform a simulation state at  $t = 0$  to a simulation state at  $t = T$  with Gaussian noise as a perturbation. Based on a given
*Correspondence to: benjamin HOLzschuh@tum.de
end state of the system at  $t = T$ , we predict a previous state by taking a small time step backward in time and repeating this multiple times. Similar to the reverse-time SDE of diffusion models, the prediction of the previous state depends on an approximate inverse of the physics simulator, a learned update  $s_{\theta}$ , and a small Gaussian perturbation.
The training of  $s_{\theta}$  is similar to learned correction approaches for numerical simulations [Um+20; Koc+21; LCT22]: The network  $s_{\theta}$  learns corrections to simulation states that evolve over time according to a physics simulator so that the corrected trajectory matches a target trajectory. In our method, we target learning corrections for the "reverse" simulation. Training can either be based on single simulation steps, which only predict a single previous state, or be extended to rollouts for multiple steps. The latter requires the differentiability of the inverse physics step [Thu+21].
Importantly, we show that under mild conditions, learning  $s_{\theta}$  is equivalent to matching the score  $\nabla_{\mathbf{x}}\log p_t(\mathbf{x})$  of the training data set distribution at time  $t$ . Therefore, sampling from the reverse-time SDE of our physical system SDE constitutes a theoretically justified method to sample from the correct posterior.
While the training with single steps directly minimizes a score matching objective, we show that the extension to multiple steps corresponds to maximum likelihood training of a related neural ordinary differential equation (ODE). Considering multiple steps is important for the stability of the produced trajectories. Feedback from physics and neural network interactions at training time leads to more robust results.
In contrast to previous diffusion methods, we include domain knowledge about the physical process in the form of an approximate inverse simulator that replaces the drift term of diffusion models [Son+21b; ZC21]. In practice, the learned component  $s_{\theta}$  corrects any errors that occur due to numerical issues, e.g., the time discretization, and breaks ambiguities due to a loss of information in the simulation over time.
Figure 1 gives an overview of our method. Our central aim is to show that the combination of diffusion-based techniques and differentiable simulations has merit for solving inverse problems and to provide a theoretical foundation for combining PDEs and diffusion modeling. In the following, we refer to methods using this combination as score matching via differentiable physics (SMDP). The main contributions of our work are: (1) We introduce a reverse physics simulation step into diffusion models to develop a probabilistic framework for solving inverse problems. (2) We provide the theoretical foundation that this combination yields learned corrections representing the score of the underlying data distribution. (3) We highlight the effectiveness of SMDP with a set of challenging inverse problems and show the superior performance of SMDP compared to a range of stochastic and deterministic baselines.
Diffusion models and generative modeling with SDEs Diffusion models [Soh+15; HJA20] have been considered for a wide range of generative applications, most notably for image [DN21], video [Ho+22; Hop+22; YSM22], audio synthesis [Che+21], uncertainty quantification [CSY22; Chu+22; Son+22; KVE21; Ram+20], and as autoregressive PDE-solvers [KCT23]. However, most approaches either focus on the denoising objective common for tasks involving natural images or the synthesis process of solutions does not directly consider the underlying physics. Models based on Langevin dynamics [Vin11; SE19] or discrete Markov chains [Soh+15; HJA20] can be unified in a time-continuous framework using SDEs [Son+21b]. Synthesizing data by sampling from neural SDEs has been considered by, e.g., [Kid+21; Son+21b]. Contrary to existing approaches, the drift in our method is an actual physics step, and the underlying SDE does not transform a data distribution to noise but models the temporal evolution of a physics system with stochastic perturbations.
Methods for solving inverse problems for (stochastic) PDEs Differentiable solvers for physics dynamics can be used to optimize solutions of inverse problems with gradient-based methods by backpropagating gradients through the solver steps [Thu+21]. Learning-based approaches directly learn solution operators for PDEs and stochastic PDEs, i.e., mappings between spaces of functions, such as Fourier neural operators [Li+21], DeepONets [Lu+21], or generalizations thereof that include stochastic forcing for stochastic PDEs, e.g., neural stochastic PDEs [SLG22]. Recently, there have been several approaches that leverage the learned scores from diffusion models as data-driven
(a) Training and inference overview.
(b) Sliding window and optimizing  $s_\theta$ .
Figure 1: Overview of our method. For training, we fit a neural ODE, the probability flow, to the set of perturbed training trajectories (a, top). The probability flow is comprised of a reverse physics simulator  $\tilde{\mathcal{P}}^{-1}$  that is an approximate inverse of the forward solver  $\mathcal{P}$  as well as a correction function  $s_\theta$ . In many cases, we can obtain  $\tilde{\mathcal{P}}^{-1}$  from  $\mathcal{P}$  by using a negative step size  $\Delta t$  or by learning a surrogate model from data. For inference, we simulate the system backward in time from  $\mathbf{x}_T$  to  $\mathbf{x}_0$  by combining  $\tilde{\mathcal{P}}^{-1}$ , the trained  $s_\theta$  and Gaussian noise in each step (a, bottom). For optimizing  $s_\theta$ , our approach moves a sliding window of size  $S$  along the training trajectories and reconstructs the current window (b). Gradients for  $\theta$  are accumulated and backpropagated through all prediction steps.
regularizations for linear inverse problems [Ram+20; Son+22; KVE21; CSY22; Chu+22] and general noisy inverse problems [Chu+23]. Our method can be applied to general non-linear inverse physics problems with temporal evolution, and we do not require to backpropagate gradients through all solver steps during inference. This makes inference significantly faster and more stable.
Learned corrections for numerical errors Numerical simulations benefit greatly from machine learning models [Tom+17; Mor+18; Pfa+20; Li+21]. By integrating a neural network into differential equation solvers, it is possible to learn to reduce numerical errors [Um+20; Koc+21; BWW22] or guide the simulation towards a desired target state [HTK20; Li+22]. The optimization of  $s_{\theta}$  with the 1-step and multi-step loss we propose in section 3.1 is conceptually similar to learned correction approaches. However, this method has, to our knowledge, not been applied to correcting the "reverse" simulation and solving inverse problems.
Maximum likelihood training and continuous normalizing flows Continuous normalizing flows (CNFs) are invertible generative models based on neural ODEs [Che+18a; KPB20; Pap+21], which are similar to our proposed physics-based neural ODE. The evolution of the marginal probability density of the SDE underlying the physics system is described by Kolmogorov's forward equation [Øks03], and there is a corresponding probability flow ODE [MRO20; Son+21b]. When the score is represented by  $s_{\theta}$ , this constitutes a CNF and can typically be trained with standard methods [Che+18b] and maximum likelihood training [Son+21a]. Huang et al. [HLC21] show that minimizing the score-matching loss is equivalent to maximizing a lower bound of the likelihood obtained by sampling from the reverse-time SDE. A recent variant combines score matching with CNFs [ZC21] and employs joint learning of drift and corresponding score for generative modeling. To the best of our knowledge, training with rollouts of multiple steps and its relation to maximum likelihood training have not been considered so far.
Problem formulation Let  $(\Omega, \mathcal{F}, P)$  be a probability space and  $W(t) = (W_1(t), \dots, W_D(t))^T$  be a  $D$ -dimensional Brownian motion. Moreover, let  $\mathbf{x}_0$  be a  $\mathcal{F}_0$ -measurable  $\mathbb{R}^D$ -valued random variable that is distributed as  $p_0$  and represents the initial simulation state. We consider the time evolution of
the physical system for  $0 \leq t \leq T$  modeled by the stochastic differential equation (SDE)
$$
d \mathbf {x} = \mathcal {P} (\mathbf {x}) d t + g (t) d W \tag {1}
$$
with initial value  $\mathbf{x}_0$  and Borel measurable drift  $\mathcal{P}:\mathbb{R}^D\to \mathbb{R}^D$  and diffusion  $g:[0,T]\rightarrow \mathbb{R}_{\geq 0}$ . This SDE transforms the marginal distribution  $p_0$  of initial states at time 0 to the marginal distribution  $p_T$  of end states at time  $T$ . We include additional assumptions in appendix A.
Moreover, we assume that we have sampled  $N$  trajectories of length  $M$  from the above SDE with a fixed time discretization  $0 \leq t_0 < t_1 < \ldots < t_M \leq T$  for the interval  $[0, T]$  and collected them in a training data set  $\{(\mathbf{x}_{t_m}^{(n)})_{i=0}^M\}_{n=0}^N$ . For simplicity, we assume that all time steps are equally spaced, i.e.,  $t_{m+1} - t_m \coloneqq \Delta t$ . Moreover, in the following we use the notation  $\mathbf{x}_{\ell:m}$  for  $0 \leq \ell < m \leq M$  to refer to the trajectory  $(\mathbf{x}_{t_\ell}, \mathbf{x}_{t_{\ell+1}}, \ldots, \mathbf{x}_{t_m})$ . We include additional assumptions in appendix A.
Our goal is to infer an initial state  $\mathbf{x}_0$  given a simulation end state  $\mathbf{x}_M$ , i.e., we want to sample from the distribution  $p_0(\cdot|\mathbf{x}_M)$ , or obtain a maximum likelihood solution.
3.1 Learned Corrections for Reverse Simulation
In the following, we furthermore assume that we have access to a reverse physics simulator  $\tilde{\mathcal{P}}^{-1}:\mathbb{R}^D\to \mathbb{R}^D$  , which moves the simulation state backward in time and is an approximate inverse of the forward simulator  $\mathcal{P}$  [HKT22]. In our experiments, we either obtain the reverse physics simulator from the forward simulator by using a negative step size  $\Delta t$  or by learning a surrogate model from the training data. We train a neural network  $s_{\theta}(\mathbf{x},t)$  parameterized by  $\theta$  such that
$$
\mathbf {x} _ {m} \approx \mathbf {x} _ {m + 1} + \Delta t \left[ \tilde {\mathcal {P}} ^ {- 1} \left(\mathbf {x} _ {m + 1}\right) + s _ {\theta} \left(\mathbf {x} _ {m + 1}, t _ {m + 1}\right) \right]. \tag {2}
$$
In this equation, the term  $s_{\theta}(\mathbf{x}_{m + 1},t_{m + 1})$  corrects approximation errors and resolves uncertainties from the Gaussian perturbation  $g(t)dW$ . Below, we explain our proposed 1-step training loss and its multi-step extension before connecting this formulation to diffusion models in the next section.
1-step loss For a pair of adjacent samples  $(\mathbf{x}_m, \mathbf{x}_{m+1})$  on a data trajectory, the 1-step loss for optimizing  $s_\theta$  is the  $L_2$  distance between  $\mathbf{x}_m$  and the prediction via (2). For the entire training data set, the loss becomes
$$
\mathcal {L} _ {\text {s i n g l e}} (\theta) := \frac {1}{M} \mathbb {E} _ {\mathbf {x} _ {0: M}} \left[ \sum_ {m = 0} ^ {M - 1} \left[ \left| \left| \mathbf {x} _ {m} - \mathbf {x} _ {m + 1} - \Delta t \left[ \tilde {\mathcal {P}} ^ {- 1} \left(\mathbf {x} _ {m + 1}\right) + s _ {\theta} \left(\mathbf {x} _ {m + 1}, t _ {m + 1}\right) \right] \right| \right| _ {2} ^ {2} \right] \right]. \tag {3}
$$
Computing the expectation can be thought of as moving a window of size two from the beginning of each trajectory until the end and averaging the losses for individual pairs of adjacent points.
Multi-step loss As each simulation state depends only on its previous state, the 1-step loss should be sufficient for training  $s_\theta$ . However, in practice, approaches that consider a loss based on predicting longer parts of the trajectories are more successful for training learned corrections [Bar+19; Um+20; Koc+21]. For that purpose, we define a hyperparameter  $S$ , called sliding window size, and write  $\mathbf{x}_{i:i+S} \in \mathbb{R}^{S \times D}$  to denote the trajectory starting at  $\mathbf{x}_i$  that is comprised of  $\mathbf{x}_i$  and the following  $S-1$  states. Then, we define the multi-step loss as
$$
\mathcal {L} _ {\text {m u l t i}} (\theta) := \frac {1}{M} \mathbb {E} _ {\mathbf {x} _ {0: M}} \left[ \sum_ {m = 0} ^ {M - S + 1} \left[ \left| \left| \mathbf {x} _ {m: m + S - 1} - \hat {\mathbf {x}} _ {m: m + S - 1} \right| \right| _ {2} ^ {2} \right] \right], \tag {4}
$$
where  $\hat{\mathbf{x}}_{i:i + S - 1}$  is the predicted trajectory that is defined recursively by
$$
\hat {\mathbf {x}} _ {i + S} = \mathbf {x} _ {i + S} \quad \text {a n d} \quad \hat {\mathbf {x}} _ {i + S - 1 - j} = \hat {\mathbf {x}} _ {i + S - j} + \Delta t \left[ \tilde {\mathcal {P}} ^ {- 1} \left(\hat {\mathbf {x}} _ {i + S - j}\right) + s _ {\theta} \left(\hat {\mathbf {x}} _ {i + S - j}, t _ {i + S - j}\right) \right]. \tag {5}
$$
3.2 Learning the Score
Denoising score matching Given a distribution of states  $p_t$  for  $0 < t < T$ , we follow [Son+21b] and consider the score matching objective
$$
\mathcal {J} _ {\mathrm {S M}} (\theta) := \frac {1}{2} \int_ {0} ^ {T} \mathbb {E} _ {\mathbf {x} \sim p _ {t}} \left[ | | s _ {\theta} (\mathbf {x}, t) - \nabla_ {\mathbf {x}} \log p _ {t} (\mathbf {x}) | | _ {2} ^ {2} \right] d t, \tag {6}
$$
i.e., the network  $s_{\theta}$  is trained to approximate the score  $\nabla_{\mathbf{x}}\log p_t(\mathbf{x})$ . In denoising score matching [Vin11; Soh+15; HJA20], the distributions  $p_t$  are implicitly defined by a noising process that is given by the forward SDE  $d\mathbf{x} = f(\mathbf{x},t)dt + g(t)dW$ , where  $W$  is the standard Brownian motion. The function  $f:\mathbb{R}^D\times \mathbb{R}_{\geq 0}\to \mathbb{R}^D$  is called drift, and  $g:\mathbb{R}_{\geq 0}\to \mathbb{R}_{\geq 0}$  is called diffusion. The process transforms the training data distribution  $p_0$  to a noise distribution that is approximately Gaussian  $p_T$ . For affine functions  $f$  and  $g$ , the transition probabilities are available analytically, which allows for efficient training of  $s_{\theta}$ . It can be shown that under mild conditions, for the forward SDE, there is a corresponding reverse-time SDE  $d\mathbf{x} = [f(\mathbf{x},t) - g(t)^2\nabla_{\mathbf{x}}\log p_t(\mathbf{x})]dt + g(t)d\tilde{W}$  [And82]. In particular, this means that given a marginal distribution of states  $p_T$ , which is approximately Gaussian, we can sample from  $p_T$  and simulate paths of the reverse-time SDE to obtain samples from the data distribution  $p_0$ .
Score matching, probability flow ODE and 1-step training There is a deterministic ODE [Son+21b], called probability flow ODE, which yields the same transformation of marginal probabilities from  $p_T$  to  $p_0$  as the reverse-time SDE. For the physics-based SDE (1), it is given by
$$
d \mathbf {x} = \left[ \mathcal {P} (\mathbf {x}) - \frac {1}{2} g (t) ^ {2} \nabla_ {\mathbf {x}} \log p _ {t} (\mathbf {x}) \right] d t. \tag {7}
$$
For  $\Delta t\to 0$  we can rewrite the update rule (2) of the training as
$$
d \mathbf {x} = \left[ - \tilde {\mathcal {P}} ^ {- 1} (\mathbf {x}) - s _ {\theta} (\mathbf {x}, t) \right] d t. \tag {8}
$$
Therefore, we can identify  $\hat{\mathcal{P}}^{-1}(\mathbf{x})$  with  $-\mathcal{P}(\mathbf{x})$  and  $s_{\theta}(\mathbf{x},t)$  with  $\frac{1}{2} g(t)\nabla_{\mathbf{x}}\log p_t(\mathbf{x})$ . We show that for the 1-step training and sufficiently small  $\Delta t$ , we minimize the score matching objective (6).
Theorem 3.1. Consider a data set with trajectories sampled from SDE (1) and let  $\tilde{\mathcal{P}}^{-1}(\mathbf{x}) = -\mathcal{P}(\mathbf{x})$ . Then the 1-step loss (3) is equivalent to minimizing the score matching objective (6) as  $\Delta t \to 0$ .
Proof. See appendix A.1
Maximum likelihood and multi-step training Extending the single step training to multiple steps does not directly minimize the score matching objective, but we can still interpret the learned correction in a probabilistic sense. For denoising score matching, it is possible to train  $s_{\theta}$  via maximum likelihood training [Son+21a], which minimizes the KL-divergence between  $p_0$  and the distribution obtained by sampling  $\mathbf{x}_T$  from  $p_T$  and simulating the probability flow ODE (8) from  $t = T$  to  $t = 0$ . We derive a similar result for the multi-step loss.
Theorem 3.2. Consider a data set with trajectories sampled from SDE (1) and let  $\tilde{\mathcal{P}}^{-1}(\mathbf{x}) = -\mathcal{P}(\mathbf{x})$ . Then the multi-step loss (4) maximizes a variational lower bound for maximum likelihood training of the probability flow ODE (7) as  $\Delta t \to 0$ .
Proof. See appendix A.2
To conclude, we have formulated a probabilistic multi-step training to solve inverse physics problems and provided a theoretical basis to solve these problems with score matching. Next, we outline additional details for implementing SMDP.
We start training  $s_{\theta}$  with the multi-step loss and window size  $S = 2$ , which is equivalent to the 1-step loss. Then, we gradually increase the window size  $S$  until a maximum  $S_{\mathrm{max}}$ . For  $S > 2$ , the unrolling of the predicted trajectory includes interactions between  $s_{\theta}$  and the reverse physics simulator  $\tilde{\mathcal{P}}^{-1}$ . For inference, we consider the neural SDE
$$
d \mathbf {x} = \left[ - \tilde {\mathcal {P}} ^ {- 1} (\mathbf {x}) - C s _ {\theta} (\mathbf {x}, t) \right] d t + g (t) d W, \tag {9}
$$
which we solve via the Euler-Maruyama method. For  $C = 2$ , we obtain the system's reverse-time SDE and sampling from this SDE yields the posterior distribution. Setting  $C = 1$  and excluding the noise gives the probability flow ODE (8). We denote the ODE variant by SMDP ODE and the SDE variant by SMDP SDE. While the SDE can be used to obtain many samples and to explore the posterior distribution, the ODE variant constitutes a unique and deterministic solution based on maximum likelihood training.
(b) Inference
(c) 1-step vs. multi-step
Figure 2: Overview of our 1D toy SDE. (a) Training with a data set of trajectories and known temporal dynamics given by  $\mathcal{P}(\mathbf{x})\coloneqq -\mathrm{sign}(\mathbf{x})\mathbf{x}^2$  and  $g\equiv 0.1$ . We estimate the score  $\nabla_{\mathbf{x}}\log p_t(\mathbf{x})$  with our proposed method using an MLP network for  $s_\theta (\mathbf{x},t)$ . Negative values (blue) push down the trajectories, and positive ones (red) push them up. Together with the dynamics, this can be used to reverse the system as shown in (b) either with the reverse-time SDE or the probability flow ODE. A successful inversion of the dynamics requires the network  $s_\theta$  to be robust and extrapolate well (c). Inference using GRID trained with the 1-step loss causes trajectories to explode, as the network does not extrapolate well. Training GRID with the multi-step loss solves this issue.
We show the capabilities of the proposed algorithm with a range of experiments. The first experiment in section 4.1 uses a simple 1D process to compare our method to existing score matching baselines. The underlying model has a known posterior distribution which allows for an accurate evaluation of the performance, and we use it to analyze the role of the multi-step loss formulation. Secondly, in section 4.2, we experiment with the stochastic heat equation. This is a particularly interesting test case as the diffusive nature of the equation effectively destroys information over time. In section 4.3, we apply our method to a scenario without stochastic perturbations in the form of a buoyancy-driven Navier-Stokes flow with obstacles. This case highlights the usefulness of the ODE variant. Finally, in section 4.4, we consider the situation where the reverse physics simulator  $\tilde{\mathcal{P}}^{-1}$  is not known. Here, we train a surrogate model  $\tilde{\mathcal{P}}^{-1}$  for isotropic turbulence flows and evaluate how well SMDP works with a learned reverse physics simulator.
4.1 1D Toy SDE
As a first experiment with a known posterior distribution we consider a simple quadratic SDE of the form:  $dx = -\left[\lambda_1 \cdot \mathrm{sign}(x)x^2\right]dt + \lambda_2 dW$ , with  $\lambda_1 = 7$  and  $\lambda_2 = 0.03$ . Throughout this experiment,  $p_0$  is a categorical distribution, where we draw either 1 or -1 with the same probability. The reverse-time SDE that transforms the distribution  $p_T$  of values at  $T = 10$  to  $p_0$  is given by